;;; SPDX-License-Identifier: LGPL-3.0-or-later
;;; SPDX-FileCopyrightText: Copyright © 2015-2024 Tony Garnock-Jones <tonyg@leastfixedpoint.com>

#lang racket/base

(provide (for-syntax preserves-pattern-registry
                     register-preserves-pattern!
                     analyse-pattern
                     analyse-pattern-bindings
                     instantiate-pattern
                     analyse-match-pattern)
         define-preserves-pattern
         :pattern
         :parse

         pattern->shape
         pattern->constant-values
         pattern->constant-paths
         pattern->check-paths
         pattern->capture-paths

         !dump-registered-preserves-patterns!

         transform-pattern-binding
         current-pattern-binding-let
         pattern-binding-let

         (all-from-out "schemas/dataspacePatterns.rkt"))

(require (for-syntax racket/base))
(require (for-syntax racket/match))
(require (for-syntax racket/list))
(require (for-syntax racket/struct-info))
(require (for-syntax racket/syntax))
(require (for-syntax syntax/id-table))
(require (for-syntax syntax/stx))

(require preserves)
(require preserves-schema)
(require racket/match)
(require racket/list)
(require racket/stxparam)
(require "pattern-expander.rkt")
(require "smart-pattern.rkt")
(require "schemas/dataspacePatterns.rkt")

;;---------------------------------------------------------------------------

(define-for-syntax preserves-pattern-registry (make-free-id-table))

(define-for-syntax (register-preserves-pattern! id-stx transformer)
  (free-id-table-set! preserves-pattern-registry id-stx transformer))

(define-syntax (define-preserves-pattern stx)
  (syntax-case stx ()
    [(_ top-type-name parser-name (ctor-stx field-stxs ...) pattern-stx bindings-stx)
     #`(begin (begin-for-syntax
                (register-preserves-pattern!
                 #'ctor-stx
                 (lambda (mode s)
                   (match mode
                     ['pattern
                      (syntax-case s ()
                        [(_ field-stxs ...) pattern-stx]
                        [_ (raise-syntax-error 'ctor-stx "Invalid pattern" s)])]
                     ['match-pattern
                      (syntax-case s ()
                        [(_ field-stxs ...) #`(app parser-name (ctor-stx field-stxs ...))]
                        [_ (raise-syntax-error 'ctor-stx "Invalid match-pattern" s)])]
                     ['bindings
                      (syntax-case s ()
                        [(_ field-stxs ...) bindings-stx]
                        [_ (raise-syntax-error 'ctor-stx "Invalid binding-pattern" s)])]))))
              (void))]))

;;---------------------------------------------------------------------------

(define-for-syntax orig-insp
  (variable-reference->module-declaration-inspector (#%variable-reference)))

(define-syntax-parameter transform-pattern-binding (lambda (id-stx expanded-stx) expanded-stx))

(define-syntax-parameter current-pattern-binding-let
  (lambda (stx)
    (syntax-case stx ()
      [(_ [(binding ...) bindings] body ...)
       (syntax/loc stx
         (match-let ([(list binding ...) bindings]) body ...))])))

(define-syntax pattern-binding-let
  (lambda (stx)
    (syntax-case stx ()
      [(_ [pat bindings] body ...)
       (with-syntax (((binding ...) (analyse-pattern-bindings #'pat)))
         (syntax/loc stx
           (current-pattern-binding-let [(binding ...) bindings] body ...)))])))

(begin-for-syntax
  (define (dollar-id? stx)
    (and (identifier? stx)
         (char=? (string-ref (symbol->string (syntax-e stx)) 0) #\$)))

  (define (undollar stx)
    (and (dollar-id? stx)
         (datum->syntax stx (string->symbol (substring (symbol->string (syntax-e stx)) 1)) stx)))

  (define (id=? expected stx)
    (and (identifier? stx)
         (free-identifier=? expected stx)))

  (define (id-value stx)
    (and (identifier? stx)
         (syntax-local-value stx (lambda () #f))))

  (define (discard-id? stx) (id=? #'_ stx))
  (define (list-id? stx) (id=? #'list stx))
  (define (hash-or-hasheqv-id? stx) (or (id=? #'hash stx) (id=? #'hasheqv stx)))

  (define (constructor-registered? stx)
    (free-id-table-ref preserves-pattern-registry stx #f))

  (define (struct-label/arity-check id-stx actual-count)
    (match-define (list _type-stx ctor-stx pred-stx accessor-stxs _mutator-stxs _super)
      (extract-struct-info (id-value id-stx)))
    (define expected-count (length accessor-stxs))
    (when (not (= expected-count actual-count))
      (raise-syntax-error ':pattern
                          (format "Expected ~a arguments to ~v, but got ~a"
                                  expected-count
                                  ctor-stx
                                  actual-count)))
    (syntax-e ctor-stx))

  (define (append-map-pairs f xs)
    (match xs
      ['() '()]
      [(list _) (raise-syntax-error #f "Odd number of elements in hash-like pattern")]
      [(list* k v more) (append (f k v) (append-map-pairs f more))]))

  (define (transform-binding id-stx expanded-stx)
    ((syntax-parameter-value #'transform-pattern-binding) id-stx expanded-stx))

  (define (walk-hash walk pieces-stx)
    (append-map-pairs (lambda (key-stx pat-stx) (list key-stx (walk pat-stx)))
                      (syntax->list pieces-stx)))

  (define (analyse-pattern stx
                           #:check-destructuring [check-destructuring (lambda (stx) stx)]
                           #:wrap-literal [wrap-literal (lambda (stx) stx)])
    (define (walk stx)
      (define disarmed-stx (syntax-disarm stx orig-insp))
      (syntax-case disarmed-stx ($ quasiquote unquote quote)
        [expander
         (pattern-expander-form? #'expander)
         (pattern-expander-transform disarmed-stx
                                     (lambda (result) (walk (syntax-rearm result stx))))]

        [(ctor args ...)
         (constructor-registered? #'ctor)
         ((free-id-table-ref preserves-pattern-registry #'ctor) 'pattern disarmed-stx)]

        ;; Extremely limited support for quasiquoting and quoting
        [(quasiquote (unquote p)) (walk #'p)]
        [(quasiquote (p ...)) (walk #'(list (quasiquote p) ...))]
        [(quasiquote p) #`(literal->literal-pattern #,(wrap-literal #''p))]
        [(quote p) #`(literal->literal-pattern #,(wrap-literal #''p))]

        [(unquote p) #'p]

        [(ctor piece ...)
         (struct-info? (id-value #'ctor))
         (let ((label (struct-label/arity-check #'ctor (length (syntax->list #'(piece ...))))))
           (check-destructuring #`(rec '#,label (list #,@(map walk (syntax->list #'(piece ...)))))))]

        [(list-stx piece ...)
         (list-id? #'list-stx)
         (check-destructuring #`(arr (list #,@(map walk (syntax->list #'(piece ...))))))]

        [(hash-stx piece ...)
         (hash-or-hasheqv-id? #'hash-stx)
         (check-destructuring #`(dict (hash #,@(walk-hash walk #'(piece ...)))))]

        [id
         (dollar-id? #'id)
         (transform-binding (undollar #'id) #`(Pattern-bind (Pattern-discard)))]

        [($ id p)
         (transform-binding #'id #`(Pattern-bind #,(walk #'p)))]

        [id
         (discard-id? #'id)
         #`(Pattern-discard)]

        [(c l (list-stx piece ...))
         (and (id=? #'rec #'c)
              (list-id? #'list-stx))
         #`(rec* l (list #,@(map walk (syntax->list #'(piece ...)))))]

        [(c (list-stx piece ...))
         (and (id=? #'arr #'c)
              (list-id? #'list-stx))
         #`(arr* (list #,@(map walk (syntax->list #'(piece ...)))))]

        [(c (hash-stx piece ...))
         (and (id=? #'dict #'c)
              (hash-or-hasheqv-id? #'hash-stx))
         #`(dict* (hash-stx #,@(walk-hash walk #'(piece ...))))]

        [other #`(literal->literal-pattern #,(wrap-literal #'other))]))
    (walk stx))

  (define (analyse-pattern-bindings stx)
    (let walk ((stx stx))
      (define disarmed-stx (syntax-disarm stx orig-insp))
      (syntax-case disarmed-stx ($ quasiquote unquote quote)
        [expander
         (pattern-expander-form? #'expander)
         (pattern-expander-transform disarmed-stx
                                     (lambda (result) (walk (syntax-rearm result stx))))]

        [(ctor args ...)
         (constructor-registered? #'ctor)
         ((free-id-table-ref preserves-pattern-registry #'ctor) 'bindings disarmed-stx)]

        ;; Extremely limited support for quasiquoting and quoting
        [(quasiquote (unquote p)) (walk #'p)]
        [(quasiquote (p ...)) (walk #'(list (quasiquote p) ...))]
        [(quasiquote _p) '()]
        [(quote _p) '()]

        [(ctor piece ...)
         (struct-info? (id-value #'ctor))
         (append-map walk (syntax->list #'(piece ...)))]

        [(list-stx piece ...)
         (list-id? #'list-stx)
         (append-map walk (syntax->list #'(piece ...)))]

        [(hash-stx piece ...)
         (hash-or-hasheqv-id? #'hash-stx)
         (append-map-pairs (lambda (_k v) (walk v)) (syntax->list #'(piece ...)))]

        [id
         (dollar-id? #'id)
         (list (undollar #'id))]

        [($ id p)
         (cons #'id (walk #'p))]

        [id
         (discard-id? #'id)
         '()]

        [(c l (list-stx piece ...))
         (and (id=? #'rec #'c)
              (list-id? #'list-stx))
         (append-map walk (syntax->list #'(piece ...)))]

        [(c (list-stx piece ...))
         (and (id=? #'arr #'c)
              (list-id? #'list-stx))
         (append-map walk (syntax->list #'(piece ...)))]

        [(c (hash-stx piece ...))
         (and (id=? #'dict #'c)
              (hash-or-hasheqv-id? #'hash-stx))
         (append-map-pairs (lambda (_k v) (walk v)) (syntax->list #'(piece ...)))]

        [other '()])))

  (define (instantiate-pattern stx lookup)
    (let walk ((stx stx))
      (define disarmed-stx (syntax-disarm stx orig-insp))
      (syntax-case disarmed-stx ($ quasiquote unquote quote)
        [expander
         (pattern-expander-form? #'expander)
         (pattern-expander-transform disarmed-stx
                                     (lambda (result) (walk (syntax-rearm result stx))))]

        ;; Extremely limited support for quasiquoting and quoting
        [(quasiquote (unquote p)) (walk #'p)]
        [(quasiquote (p ...)) (walk #'(list (quasiquote p) ...))]
        [(quasiquote _p) '()]
        [(quote _p) '()]

        [(hash-stx piece ...)
         (hash-or-hasheqv-id? #'hash-stx)
         (quasisyntax/loc stx (hash-stx #,@(walk-hash walk #'(piece ...))))]

        [id
         (dollar-id? #'id)
         (lookup (undollar #'id) #f)]

        [($ id p)
         (lookup #'id (lambda () (walk #'p)))]

        [id
         (discard-id? #'id)
         (raise-syntax-error 'instantiate-pattern "Cannot instantiate discard patterns" stx)]

        [(f args ...)
         (quasisyntax/loc stx (f #,@(map walk (syntax->list #'(args ...)))))]

        [other #'other])))

  (define (analyse-match-pattern stx)
    (let walk ((stx stx))
      (define disarmed-stx (syntax-disarm stx orig-insp))
      (syntax-case disarmed-stx ($ quasiquote unquote quote)
        [expander
         (pattern-expander-form? #'expander)
         (pattern-expander-transform disarmed-stx
                                     (lambda (result) (walk (syntax-rearm result stx))))]

        [(ctor args ...)
         (constructor-registered? #'ctor)
         ((free-id-table-ref preserves-pattern-registry #'ctor) 'match-pattern disarmed-stx)]

        ;; Extremely limited support for quasiquoting and quoting
        [(quasiquote (unquote p)) (walk #'p)]
        [(quasiquote (p ...)) (walk #'(list (quasiquote p) ...))]
        [(quasiquote p) #''p]
        [(quote p) #''p]

        [(unquote p) (raise-syntax-error #f "Out-of-place unquote in match-pattern")]

        [(ctor piece ...)
         (struct-info? (id-value #'ctor))
         #`(ctor (:parse piece) ...)]

        [(list-stx piece ...)
         (list-id? #'list-stx)
         #`(list-stx (:parse piece) ...)]

        [(hash-stx piece ...)
         (hash-or-hasheqv-id? #'hash-stx)
         #`(hash-table #,@(let loop ((pieces (syntax->list #'(piece ...))))
                            (match pieces
                              ['() '()]
                              [(list* k v more) (list* k #`(:parse #,v) (loop more))]))
                       [_ _] ___)]

        [(or-stx piece ...)
         (and (identifier? #'or-stx) (free-identifier=? #'or #'or-stx))
         #`(or (:parse piece) ...)]

        [(and-stx piece ...)
         (and (identifier? #'and-stx) (free-identifier=? #'and #'and-stx))
         #`(and (:parse piece) ...)]

        [other
         #`other])))

  (define (expand-:pattern stx)
    (syntax-case stx ()
      [(_ pat-stx atomic-literal-transformer)
       (analyse-pattern #'pat-stx
                        #:check-destructuring
                        (lambda (stx) (raise-syntax-error #f "Attempt to destructure known-atomic"))
                        #:wrap-literal
                        (lambda (stx) #`(atomic-literal-transformer #,stx)))]
      [(_ pat-stx)
       (analyse-pattern #'pat-stx)])))

(define-pattern-expander :pattern
  (lambda (stx)
    ;; This effectively quasiquotes the pattern; a naked unquote in
    ;; the pattern returns to the "outer" context. This is useful for
    ;; observing observers:
    ;;
    ;;   (Observe (:pattern (some-pattern ...)) _)
    ;;
    (expand-:pattern stx))
  (lambda (stx)
    (expand-:pattern stx)))

(define-match-expander :parse
  (lambda (stx)
    (syntax-case stx ()
      [(_ pat-stx) (analyse-match-pattern #'pat-stx)])))

;;---------------------------------------------------------------------------

(define (pattern->shape desc)
  (let walk-node ((key-rev '()) (desc desc))
    (match desc
      [(Pattern-group type entries)
       (append* (list (cons (reverse key-rev) type))
                (for/list [((key subdesc) (in-hash entries))]
                  (walk-node (cons key key-rev) subdesc)))]
      [(Pattern-bind subdesc)
       (walk-node key-rev subdesc)]
      [(Pattern-discard)
       '()]
      [(Pattern-lit value)
       '()])))

(define (select-pattern-leaves desc capture-fn discard-fn lit-fn)
  (let walk-node ((key-rev '()) (desc desc))
    (match desc
      [(Pattern-group _type entries)
       (append* (for/list [((key subdesc) (in-hash entries))]
                  (walk-node (cons key key-rev) subdesc)))]
      [(Pattern-bind subdesc)
       (append (capture-fn key-rev) (walk-node key-rev subdesc))]
      [(Pattern-discard)
       (discard-fn key-rev)]
      [(Pattern-lit value)
       (lit-fn key-rev (->preserve value))])))

(define (pattern->constant-values desc)
  (select-pattern-leaves desc
                         (lambda (key-rev) (list))
                         (lambda (key-rev) (list))
                         (lambda (key-rev value) (list value))))

(define (pattern->constant-paths desc)
  (select-pattern-leaves desc
                         (lambda (key-rev) (list))
                         (lambda (key-rev) (list))
                         (lambda (key-rev value) (list (reverse key-rev)))))

(define (pattern->check-paths desc)
  (select-pattern-leaves desc
                         (lambda (key-rev) (list))
                         (lambda (key-rev) (list (reverse key-rev)))
                         (lambda (key-rev value) (list))))

(define (pattern->capture-paths desc)
  (select-pattern-leaves desc
                         (lambda (key-rev) (list (reverse key-rev)))
                         (lambda (key-rev) (list))
                         (lambda (key-rev value) (list))))

(define-syntax (!dump-registered-preserves-patterns! stx)
  (syntax-case stx ()
    [(_)
     (let ()
       (local-require racket/pretty)
       (for [(k (in-list (free-id-table-keys preserves-pattern-registry)))]
         (printf " - ~v\n" k))
       #'(void))]))

;;---------------------------------------------------------------------------
;;; Local Variables:
;;; eval: (put 'pattern-binding-let 'racket-indent-function 1)
;;; eval: (put 'current-pattern-binding-let 'racket-indent-function 1)
;;; End: