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

#lang racket/base

(provide (for-syntax preserves-pattern-registry
                     register-preserves-pattern!
                     analyse-pattern
                     analyse-pattern-bindings)
         define-preserves-pattern
         :pattern
         :template

         pattern->constant-values
         pattern->constant-paths
         pattern->capture-paths
         pattern->capture-names

         !dump-registered-preserves-patterns!

         (all-from-out "schemas/gen/dataspace-patterns.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-schema)
(require racket/match)
(require racket/list)
(require "pattern-expander.rkt")
(require "schemas/gen/dataspace-patterns.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 (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")])]
                     ['template
                      (syntax-case s ()
                        [(_ field-stxs ...)
                         (syntax (->preserve (ctor-stx field-stxs ...)))])]
                     ['bindings
                      (syntax-case s ()
                        [(_ field-stxs ...) bindings-stx]
                        [_ (raise-syntax-error 'ctor-stx "Invalid binding-pattern")])]))))
              (void))]))

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

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

(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)))))

  (define (discard-id? stx)
    (and (identifier? stx)
         (free-identifier=? #'_ stx)))

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

  (define (list-id? stx)
    (and (identifier? stx)
         (free-identifier=? #'list stx)))

  (define (hash-or-hasheqv-id? stx)
    (and (identifier? stx)
         (or (free-identifier=? #'hash stx)
             (free-identifier=? #'hasheqv stx))))

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

  (define (member-entry key-stx pat-stx)
    (define analysed (analyse-pattern pat-stx))
    (syntax-case analysed (Pattern-DDiscard DDiscard)
      [(Pattern-DDiscard (DDiscard)) (list)]
      [_ (list key-stx analysed)]))

  (define (struct-label-and-arity 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)))
    (values (syntax-e ctor-stx)
            expected-count))

  (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 (analyse-pattern stx)
    (define disarmed-stx (syntax-disarm stx orig-insp))
    (syntax-case disarmed-stx ($ quasiquote unquote quote)
      [(ctor args ...)
       (constructor-registered? #'ctor)
       ((free-id-table-ref preserves-pattern-registry #'ctor) 'pattern disarmed-stx)]

      [(expander args ...)
       (pattern-expander-id? #'expander)
       (pattern-expander-transform disarmed-stx
                                   (lambda (result)
                                     (analyse-pattern (syntax-rearm result stx))))]

      ;; Extremely limited support for quasiquoting and quoting
      [(quasiquote (unquote p)) (analyse-pattern #'p)]
      [(quasiquote (p ...)) (analyse-pattern #'(list (quasiquote p) ...))]
      [(quasiquote p) #`(Pattern-DLit (DLit 'p))]
      [(quote p) #`(Pattern-DLit (DLit 'p))]

      [(unquote p) #'p]

      [(ctor piece ...)
       (struct-info? (id-value #'ctor))
       (let-values (((label arity) (struct-label-and-arity #'ctor (length (syntax->list #'(piece ...))))))
         #`(Pattern-DCompound
            (DCompound-rec (CRec '#,label #,arity)
                           (hasheqv #,@(append*
                                        (for/list [(n (in-naturals))
                                                   (piece (in-list (syntax->list #'(piece ...))))]
                                          (member-entry n piece)))))))]

      [(list-stx piece ...)
       (list-id? #'list-stx)
       #`(Pattern-DCompound
          (DCompound-arr (CArr #,(length (syntax->list #'(piece ...))))
                         (hasheqv #,@(append*
                                      (for/list [(n (in-naturals))
                                                 (piece (in-list (syntax->list #'(piece ...))))]
                                        (member-entry n piece))))))]

      [(hash-stx piece ...)
       (hash-or-hasheqv-id? #'hash-stx)
       #`(Pattern-DCompound
          (DCompound-dict (CDict)
                          (hash #,@(append-map-pairs member-entry (syntax->list #'(piece ...))))))]

      [id
       (dollar-id? #'id)
       #`(Pattern-DBind (DBind '#,(undollar #'id) (Pattern-DDiscard (DDiscard))))]

      [($ (unquote bp) p)
       #`(Pattern-DBind (DBind bp #,(analyse-pattern #'p)))]

      [($ id p)
       #`(Pattern-DBind (DBind 'id #,(analyse-pattern #'p)))]

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

      [other
       #`(Pattern-DLit (DLit other))]))

  (define (analyse-pattern-bindings stx)
    (let walk ((stx stx))
      (define disarmed-stx (syntax-disarm stx orig-insp))
      (syntax-case disarmed-stx ($ quasiquote unquote quote)
        [(ctor args ...)
         (constructor-registered? #'ctor)
         ((free-id-table-ref preserves-pattern-registry #'ctor) 'bindings disarmed-stx)]

        [(expander args ...)
         (pattern-expander-id? #'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) '()]

        [(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)
         '()]

        [other
         '()])))

  (define (analyse-template stx)
    (syntax-case stx ($ quasiquote unquote quote)
      [(ctor args ...)
       (constructor-registered? #'ctor)
       ((free-id-table-ref preserves-pattern-registry #'ctor) 'template stx)]

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

      [(ctor piece ...)
       (struct-info? (id-value #'ctor))
       #`(ctor #,@(map analyse-template (syntax->list #'(piece ...))))]

      [(list-stx piece ...)
       (list-id? #'list-stx)
       #`(list-stx #,@(map analyse-template (syntax->list #'(piece ...))))]

      [(hash-stx piece ...)
       (hash-or-hasheqv-id? #'hash-stx)
       #`(hash-stx #,@(append-map-pairs (lambda (k v) (list k (analyse-template v)))
                                        (syntax->list #'(piece ...))))]

      [other #'other])))

(define-pattern-expander :pattern
  (lambda (stx)
    (syntax-case stx ()
      [(_ pat-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 ...)) _)
       (analyse-pattern #'pat-stx)]))
  (lambda (stx)
    (syntax-case stx ()
      [(_ pat-stx)
       (analyse-pattern #'pat-stx)])))

(define-syntax (:template stx)
  (syntax-case stx ()
    [(_ template-stx)
     (analyse-template #'template-stx)]))

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

(define (select-pattern-leaves desc capture-fn lit-fn)
  (let walk-node ((key-rev '()) (desc desc))
    (match desc
      [(Pattern-DCompound (or (DCompound-rec _ members)
                              (DCompound-arr _ members)
                              (DCompound-dict _ members)))
       (append* (for/list [((key subdesc) (in-hash members))]
                  (walk-node (cons key key-rev) subdesc)))]
      [(Pattern-DBind (DBind name subdesc))
       (append (capture-fn key-rev name) (walk-node key-rev subdesc))]
      [(Pattern-DDiscard (DDiscard))
       '()]
      [(Pattern-DLit (DLit value))
       (lit-fn key-rev value)])))

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

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

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

(define (pattern->capture-names desc)
  (select-pattern-leaves desc
                         (lambda (key-rev name-stx) (list name-stx))
                         (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))]))