#lang racket/base

(provide (all-defined-out))

(require syndicate/support/struct)
(require racket/match)
(require racket/set)
(require racket/list)
(require racket/hash)

(require (for-syntax racket/base))
(require (for-syntax syntax/stx))

(module+ test (require rackunit))

;; A `SkProj` is a *skeleton projection*, a specification of loci
;; within a tree-shaped datum to collect into a flat list.
;;
;;     SkProj = (Listof (Listof Nat))
;;
;; The outer list specifies elements of the flat list; the inner lists
;; specify paths via zero-indexed links to child nodes in the
;; tree-shaped datum being examined. A precondition for use of a
;; `SkProj` is that the datum being examined has been checked for
;; conformance to the skeleton being projected.

;; A `SkKey` is the result of running a `SkProj` over a term,
;; extracting the values at the denoted locations.

;; A `SkCont` is a *skeleton continuation*, a collection of "next
;; steps" after a `Skeleton` has matched the general outline of a
;; datum.
;;
;;     SkCont = (MutableHash SkProj (MutableHash SkKey (MutableHash SkProj (Setof (... -> Any)))))
;;
;; The outer `SkProj` selects *constant* portions of the term for more
;; matching against the `SkKey`s in the hash table. The inner
;; `SkProj`, if any, selects *variable* portions of the term to be
;; given to the handler function.

;; A `Skeleton` is a structural guard on a datum: essentially,
;; specification of (the outline of) its shape; its silhouette.
;;
;;     Skeleton = (skeleton-node SkCont (AListof SkLabel SkNode))
;;      SkLabel = (skeleton-edge Nat Nat SkClass Nat)
;;      SkClass = StructType | 'list
;;
(struct skeleton-node (continuations [edges #:mutable]) #:transparent)
(struct skeleton-edge (pop-count index class arity) #:transparent)

(define (make-empty-skeleton)
  (skeleton-node (make-hash) '()))

(define (select-pattern-leaves stx capture-fn atom-fn)
  (define (walk-node key-rev stx)
    (match stx
      [(list pieces ...) (walk-edge 0 key-rev pieces)]
      ['$ (capture-fn key-rev)]
      ['_ (list)]
      [atom (atom-fn key-rev atom)]))
  (define (walk-edge index key-rev pieces)
    (match pieces
      ['() '()]
      [(cons p pieces) (append (walk-node (cons index key-rev) p)
                               (walk-edge (+ index 1) key-rev pieces))]))
  (walk-node '(0) stx))

(define (pattern-stx->key stx)
  (select-pattern-leaves stx
                         (lambda (_key-rev) (list))
                         (lambda (_key-rev atom) (list atom))))

(define (pattern-stx->skeleton-proj stx)
  (select-pattern-leaves stx
                         (lambda (_key-rev) (list))
                         (lambda (key-rev _atom) (list (reverse key-rev)))))

(define (pattern-stx->capture-proj stx)
  (select-pattern-leaves stx
                         (lambda (key-rev) (list (reverse key-rev)))
                         (lambda (_key-rev _atom) (list))))

(define merge-skcont!
  (let ()
    (define (merge-proj-handler old new)
      (hash-union! old new #:combine set-union))
    (define (merge-key-proj-handler old new)
      (hash-union! old new #:combine merge-proj-handler))
    (lambda (old new)
      (hash-union! old new #:combine merge-key-proj-handler))))

;; Imperatively extends `sk` to include the pattern `stx` terminating
;; in `skcont`.
(define (extend-skeleton! sk skcont stx)
  (define (walk-node! sk pop-count index stx)
    (match stx
      [(list pieces ...)
       (define edge (skeleton-edge pop-count index 'list (length pieces)))
       (define next
         (match (assoc edge (skeleton-node-edges sk))
           [#f (let ((next (make-empty-skeleton)))
                 (set-skeleton-node-edges! sk (cons (cons edge next) (skeleton-node-edges sk)))
                 next)]
           [(cons _edge next) next]))
       (walk-edge! next 0 0 pieces)]
      [_
       (values pop-count sk)]))
  (define (walk-edge! sk pop-count index pieces)
    (match pieces
      ['()
       (values (+ pop-count 1) sk)]
      [(cons p pieces)
       (let-values (((pop-count sk) (walk-node! sk pop-count index p)))
         (walk-edge! sk pop-count (+ index 1) pieces))]))
  (let-values (((_pop-count sk) (walk-edge! sk 0 0 (list stx))))
    (merge-skcont! (skeleton-node-continuations sk) skcont)
    sk))

;; TODO: avoid repeated descent into `term` by factoring out prefixes of paths in `proj`
(define (apply-projection term proj)
  (for/list [(path proj)]
    (for/fold [(term (list term))] [(index path)]
      (cond [(list? term) (list-ref term index)]
            [else (error 'apply-projection "Non-lists not supported: ~v" term)]))))

(module+ test
  (define stx0 '$)
  (define stx1 '3)
  (define stxA '($ 3))
  (define stxB '(4 $))
  (define stxC '(_ ($ $)))
  (define stxD '(_ (1 2)))
  (define stxE '((6 6) _))
  (define stxF '(($ $) (3 9)))
  (define stxG '((_ _) (1 2)))
  (define stxH '((_ _) ((_ _) ($ _))))
  (define stxI '(((_ _) _) (_ _)))
  (define stxJ '(((_ _) (_ _)) (_ _)))

  (define (summarise-skeleton sk)
    (define (walk-node sk)
      (match-define (skeleton-node continuations edges) sk)
      (append (if (hash-empty? continuations) '() (list continuations))
              (map walk-edge edges)))
    (define (walk-edge e)
      (match-define (cons (skeleton-edge pop-count index class arity) sk) e)
      (for/fold [(acc (list* (list index class arity) (walk-node sk)))]
                [(n pop-count)]
        (list 'POP acc)))
    (walk-node sk))

  (define (skcont . ids)
    (define acc (make-hash))
    ;; Not quite the right shape! Just a dummy placeholder for testing
    (for [(id ids)]
      (merge-skcont! acc
                     (make-hash
                      (list (cons id (make-hash
                                      (list (cons id (make-hash
                                                      (list (cons id (set))))))))))))
    acc)

  (define (skeleton-stx->skeleton id pat-stx)
    (define sk (make-empty-skeleton))
    (extend-skeleton! sk (skcont id) pat-stx)
    sk)

  (check-equal? `(,(skcont 0))
                (summarise-skeleton (skeleton-stx->skeleton '0 stx0)))
  (check-equal? `(,(skcont 1))
                (summarise-skeleton (skeleton-stx->skeleton '1 stx1)))
  (check-equal? `(((0 list 2) ,(skcont 'A)))
                (summarise-skeleton (skeleton-stx->skeleton 'A stxA)))
  (check-equal? `(((0 list 2) ,(skcont 'B)))
                (summarise-skeleton (skeleton-stx->skeleton 'B stxB)))
  (check-equal? `(((0 list 2) ((1 list 2) ,(skcont 'C))))
                (summarise-skeleton (skeleton-stx->skeleton 'C stxC)))
  (check-equal? `(((0 list 2) ((1 list 2) ,(skcont 'D))))
                (summarise-skeleton (skeleton-stx->skeleton 'D stxD)))
  (check-equal? `(((0 list 2) ((0 list 2) ,(skcont 'E))))
                (summarise-skeleton (skeleton-stx->skeleton 'E stxE)))
  (check-equal? `(((0 list 2) ((0 list 2) (POP ((1 list 2) ,(skcont 'F))))))
                (summarise-skeleton (skeleton-stx->skeleton 'F stxF)))
  (check-equal? `(((0 list 2) ((0 list 2) (POP ((1 list 2) ,(skcont 'G))))))
                (summarise-skeleton (skeleton-stx->skeleton 'G stxG)))
  (check-equal? `(((0 list 2)
                   ((0 list 2) (POP ((1 list 2) ((0 list 2) (POP ((1 list 2) ,(skcont 'H)))))))))
                (summarise-skeleton (skeleton-stx->skeleton 'H stxH)))
  (check-equal? `(((0 list 2) ((0 list 2) ((0 list 2) (POP (POP ((1 list 2) ,(skcont 'I))))))))
                (summarise-skeleton (skeleton-stx->skeleton 'I stxI)))
  (check-equal? `(((0 list 2)
                   ((0 list 2)
                    ((0 list 2) (POP ((1 list 2) (POP (POP ((1 list 2) ,(skcont 'J))))))))))
                (summarise-skeleton (skeleton-stx->skeleton 'J stxJ)))

  (check-equal? `(,(skcont 0 1)
                  ((0 list 2)
                   ,(skcont 'A 'B)
                   ((0 list 2)
                    ,(skcont 'E)
                    ((0 list 2)
                     (POP ((1 list 2) (POP (POP ((1 list 2) ,(skcont 'J))))))
                     (POP (POP ((1 list 2) ,(skcont 'I)))))
                    (POP ((1 list 2) ,(skcont 'F 'G) ((0 list 2) (POP ((1 list 2) ,(skcont 'H)))))))
                   ((1 list 2) ,(skcont 'C 'D))))
                (let ((sk (make-empty-skeleton)))
                  (extend-skeleton! sk (skcont '0) stx0)
                  (extend-skeleton! sk (skcont '1) stx1)
                  (extend-skeleton! sk (skcont 'A) stxA)
                  (extend-skeleton! sk (skcont 'B) stxB)
                  (extend-skeleton! sk (skcont 'C) stxC)
                  (extend-skeleton! sk (skcont 'D) stxD)
                  (extend-skeleton! sk (skcont 'E) stxE)
                  (extend-skeleton! sk (skcont 'F) stxF)
                  (extend-skeleton! sk (skcont 'G) stxG)
                  (extend-skeleton! sk (skcont 'H) stxH)
                  (extend-skeleton! sk (skcont 'I) stxI)
                  (extend-skeleton! sk (skcont 'J) stxJ)
                  (summarise-skeleton sk)))

  (check-equal? '() (pattern-stx->skeleton-proj stx0))
  (check-equal? '((0)) (pattern-stx->skeleton-proj stx1))
  (check-equal? '((0 1)) (pattern-stx->skeleton-proj stxA))
  (check-equal? '((0 0)) (pattern-stx->skeleton-proj stxB))
  (check-equal? '() (pattern-stx->skeleton-proj stxC))
  (check-equal? '((0 1 0) (0 1 1)) (pattern-stx->skeleton-proj stxD))
  (check-equal? '((0 0 0) (0 0 1)) (pattern-stx->skeleton-proj stxE))
  (check-equal? '((0 1 0) (0 1 1)) (pattern-stx->skeleton-proj stxF))
  (check-equal? '((0 1 0) (0 1 1)) (pattern-stx->skeleton-proj stxG))
  (check-equal? '() (pattern-stx->skeleton-proj stxH))
  (check-equal? '() (pattern-stx->skeleton-proj stxI))
  (check-equal? '() (pattern-stx->skeleton-proj stxJ))

  (check-equal? '() (pattern-stx->key stx0))
  (check-equal? '(3) (pattern-stx->key stx1))
  (check-equal? '(3) (pattern-stx->key stxA))
  (check-equal? '(4) (pattern-stx->key stxB))
  (check-equal? '() (pattern-stx->key stxC))
  (check-equal? '(1 2) (pattern-stx->key stxD))
  (check-equal? '(6 6) (pattern-stx->key stxE))
  (check-equal? '(3 9) (pattern-stx->key stxF))
  (check-equal? '(1 2) (pattern-stx->key stxG))
  (check-equal? '() (pattern-stx->key stxH))
  (check-equal? '() (pattern-stx->key stxI))
  (check-equal? '() (pattern-stx->key stxJ))

  (check-equal? '((0)) (pattern-stx->capture-proj stx0))
  (check-equal? '() (pattern-stx->capture-proj stx1))
  (check-equal? '((0 0)) (pattern-stx->capture-proj stxA))
  (check-equal? '((0 1)) (pattern-stx->capture-proj stxB))
  (check-equal? '((0 1 0) (0 1 1)) (pattern-stx->capture-proj stxC))
  (check-equal? '() (pattern-stx->capture-proj stxD))
  (check-equal? '() (pattern-stx->capture-proj stxE))
  (check-equal? '((0 0 0) (0 0 1)) (pattern-stx->capture-proj stxF))
  (check-equal? '() (pattern-stx->capture-proj stxG))
  (check-equal? '((0 1 1 0)) (pattern-stx->capture-proj stxH))
  (check-equal? '() (pattern-stx->capture-proj stxI))
  (check-equal? '() (pattern-stx->capture-proj stxJ))

  (check-equal? '(goodbye hello)
                (apply-projection '((goodbye hello) (3 9)) (pattern-stx->capture-proj stxF)))
  (check-equal? '(99)
                (apply-projection '(4 99) (pattern-stx->capture-proj stxB)))
  (check-equal? '((4 99))
                (apply-projection '(4 99) (pattern-stx->capture-proj stx0)))
  )
