#lang racket/base
;; "Skeletons" describe the indexed structure of a dataspace.
;; In particular, they efficiently connect assertions to matching interests.

(provide (struct-out skeleton-interest)
         (struct-out visibility-restriction)
         make-empty-skeleton
         add-interest!
         remove-interest!
         add-assertion!
         remove-assertion!
         send-assertion!)

(module+ for-test
  (provide make-empty-skeleton/cache
           extend-skeleton!
           path-cmp))

(require racket/match)
(require racket/hash)
(require racket/list)

(require "support/struct.rkt")
(require "bag.rkt")
(require "pattern.rkt")
(require "assertions.rkt")

;; A VisibilityRestriction describes ... TODO
;; (visibility-restriction SkProj Assertion)
(struct visibility-restriction (proj term) #:transparent)

;; A ScopedAssertion is a VisibilityRestriction or an Assertion.
;; (Corollary: Instances of `visibility-restriction` can never be assertions.)

;; A `Skeleton` is a structural guard on an assertion: essentially,
;; specification of (the outline of) its shape; its silhouette.
;; Following a skeleton's structure leads to zero or more `SkCont`s.
;;
;;       Skeleton = (skeleton-node SkCont (AListof SkSelector (MutableHash SkClass SkNode)))
;;     SkSelector = (skeleton-selector Nat Nat)
;;        SkClass = StructType | (list-type Nat) | (vector-type Nat)
;;
(struct skeleton-node (continuation [edges #:mutable]) #:transparent)
(struct skeleton-selector (pop-count index) #:transparent)
(struct list-type (arity) #:transparent)
(struct vector-type (arity) #:transparent)
;;
;; A `SkDesc` is a single assertion silhouette, usually the
;; evaluation-result of `desc->skeleton-stx` from `pattern.rkt`.
;;
;; A `SkCont` is a *skeleton continuation*, a collection of "next
;; steps" after a `Skeleton` has matched the general outline of an
;; assertion.
;;
;; INVARIANT: At each level, the caches correspond to the
;; appropriately filtered and projected contents of the dataspace
;; containing the structures.
;;
;;      SkCont = (skeleton-continuation
;;                 (MutableSet ScopedAssertion)
;;                 (MutableHash SkProj (MutableHash SkKey SkConst)))
;;     SkConst = (skeleton-matched-constant
;;                 (MutableSet ScopedAssertion)
;;                 (MutableHash SkProj SkAcc))
;;       SkAcc = (skeleton-accumulator
;;                 (MutableBag SkKey)
;;                 (MutableSeteq (... -> Any)))
;;
(struct skeleton-continuation (cache table) #:transparent)
(struct skeleton-matched-constant (cache table) #:transparent)
(struct skeleton-accumulator (cache handlers) #:transparent)
;;
;; A `SkProj` is a *skeleton projection*, a specification of loci
;; within a tree-shaped assertion 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 assertion being examined. A precondition for use of a
;; `SkProj` is that the assertion 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.
;;
;;     SkKey = (Listof Any)
;;
;; Each `SkProj` in `SkCont` selects *constant* portions of the term
;; for more matching against the `SkKey`s in the table associated with
;; the `SkProj`. Each `SkProj` in `SkConst`, if any, selects
;; *variable* portions of the term to be given to the handler
;; functions in the associated `SkAcc`.

;; A `SkInterest` is a specification for an addition to or removal
;; from an existing `Skeleton`.
;;
;;     SkInterest = (skeleton-interest SkDesc
;;                                     SkProj
;;                                     SkKey
;;                                     SkProj
;;                                     (... -> Any)
;;                                     (Option ((MutableBag SkKey) -> Any)))
;;
;; The `SkDesc` gives the silhouette. The first `SkProj` is the
;; constant-portion selector, to be matched against the `SkKey`. The
;; second `SkProj` is used on matching assertions to extract the
;; variable portions, to be passed to the handler function.
;;
(struct skeleton-interest (desc
                           const-selector
                           const-value
                           var-selector
                           handler
                           cleanup
                           ) #:transparent)

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

(define (make-empty-skeleton/cache cache)
  (skeleton-node (skeleton-continuation cache
                                        (make-hash))
                 '()))

(define (make-empty-skeleton)
  (make-empty-skeleton/cache (make-hash)))

(define (make-empty-matched-constant)
  (skeleton-matched-constant (make-hash) (make-hash)))

(define (skcont-add! c i)
  (match-define (skeleton-interest _desc cs cv vs h _cleanup) i)
  (define (classify-assertions)
    (define cvt (make-hash))
    (hash-for-each (skeleton-continuation-cache c)
                   (lambda (a _)
                     (define avs (apply-projection (unscope-assertion a) cs))
                     (define sc (hash-ref! cvt avs make-empty-matched-constant))
                     (hash-set! (skeleton-matched-constant-cache sc) a #t)))
    cvt)
  (define cvt (hash-ref! (skeleton-continuation-table c) cs classify-assertions))
  (define sc (hash-ref! cvt cv make-empty-matched-constant))
  (define (make-accumulator)
    (define cache (make-bag))
    (hash-for-each (skeleton-matched-constant-cache sc)
                   (lambda (a _)
                     (unpack-scoped-assertion [restriction-path term] a)
                     (when (unrestricted? vs restriction-path)
                       (bag-change! cache (apply-projection term vs) 1))))
    (skeleton-accumulator cache (make-hasheq)))
  (define acc (hash-ref! (skeleton-matched-constant-table sc) vs make-accumulator))
  (hash-set! (skeleton-accumulator-handlers acc) h #t)
  (for [(vars (in-bag (skeleton-accumulator-cache acc)))] (apply h '+ vars)))

(define (skeleton-matched-constant-empty? sc)
  (match-define (skeleton-matched-constant cache table) sc)
  (and (hash-empty? cache) (hash-empty? table)))

(define (skcont-remove! c i)
  (match-define (skeleton-interest _desc cs cv vs h cleanup) i)
  (define cvt (hash-ref (skeleton-continuation-table c) cs #f))
  (when cvt
    (define sc (hash-ref cvt cv #f))
    (when sc
      (define acc (hash-ref (skeleton-matched-constant-table sc) vs #f))
      (when acc
        (when (and cleanup (hash-has-key? (skeleton-accumulator-handlers acc) h))
          (cleanup (skeleton-accumulator-cache acc)))
        (hash-remove! (skeleton-accumulator-handlers acc) h)
        (when (hash-empty? (skeleton-accumulator-handlers acc))
          (hash-remove! (skeleton-matched-constant-table sc) vs)))
      (when (skeleton-matched-constant-empty? sc)
        (hash-remove! cvt cv)))
    (when (hash-empty? cvt)
      (hash-remove! (skeleton-continuation-table c) cs))))

(define (term-matches-class? term class)
  (cond
    [(list-type? class) (and (list? term) (= (length term) (list-type-arity class)))]
    [(vector-type? class) (and (vector? term) (= (vector-length term) (vector-type-arity class)))]
    [(struct-type? class) (and (non-object-struct? term) (eq? (struct->struct-type term) class))]
    [else (error 'term-matches-class? "Invalid class: ~v" class)]))

(define (subterm-matches-class? term path class)
  (term-matches-class? (apply-projection-path (unscope-assertion term) path) class))

(define (unscope-assertion scoped-assertion)
  (match scoped-assertion
    [(visibility-restriction _ term) term]
    [term term]))

(define-syntax-rule (unpack-scoped-assertion [path term] expr)
  (define-values (path term)
    (match expr
      [(visibility-restriction p t) (values p t)]
      [other (values #f other)])))

(define (path-cmp a b)
  (match* (a b)
    [('() '()) '=]
    [('() _) '<]
    [(_ '()) '>]
    [((cons av a) (cons bv b))
     (cond [(< av bv) '<]
           [(> av bv) '>]
           [else (path-cmp a b)])]))

(define (unrestricted? capture-paths restriction-paths)
  ;; We are "unrestricted" if Set(capture-paths) ⊆ Set(restriction-paths). Since both variables
  ;; really hold lists, we operate with awareness of the order the lists are built here. We
  ;; know that the lists are built in fringe order; that is, they are sorted wrt `path-cmp`.
  (or (not restriction-paths)
      (let outer ((cs capture-paths) (rs restriction-paths))
        (match cs
          ['() #t]
          [(cons c cs)
           (let inner ((rs rs))
             (match rs
               ['() #f] ;; `c` does not exist in `restriction-paths` ∴ restricted.
               [(cons r rs)
                (match (path-cmp c r)
                  ['< #f] ;; `c` < `r` ==> `c` not in `restriction-paths` ∴ restricted.
                  ['= (outer cs rs)]
                  ['> (inner rs)])]))]))))

(define (extend-skeleton! sk desc)
  (define (walk-node! rev-path sk pop-count index desc)
    (match desc
      [(list class-desc pieces ...)
       (define class
         (cond [(struct-type? class-desc) class-desc]
               [(eq? class-desc 'list) (list-type (length pieces))]
               [(eq? class-desc 'vector) (vector-type (length pieces))]
               [else (error 'extend-skeleton! "Invalid class-desc: ~v" class-desc)]))
       (define selector (skeleton-selector pop-count index))
       (define table
         (match (assoc selector (skeleton-node-edges sk))
           [#f (let ((table (make-hash)))
                 (set-skeleton-node-edges! sk (cons (cons selector table) (skeleton-node-edges sk)))
                 table)]
           [(cons _selector table) table]))
       (define (make-skeleton-node-with-cache)
         (define unfiltered (skeleton-continuation-cache (skeleton-node-continuation sk)))
         (define filtered (make-hash))
         (define path (reverse rev-path))
         (hash-for-each unfiltered
                        (lambda (a _)
                          (when (subterm-matches-class? a path class)
                            (hash-set! filtered a #t))))
         (make-empty-skeleton/cache filtered))
       (define next (hash-ref! table class make-skeleton-node-with-cache))
       (walk-edge! (cons 0 rev-path) next 0 0 pieces)]
      [_
       (values pop-count sk)]))
  (define (walk-edge! rev-path sk pop-count index pieces)
    (match pieces
      ['()
       (values (+ pop-count 1) sk)]
      [(cons p pieces)
       (let-values (((pop-count sk) (walk-node! rev-path sk pop-count index p)))
         (walk-edge! (cons (+ index 1) (cdr rev-path)) sk pop-count (+ index 1) pieces))]))
  (let-values (((_pop-count sk) (walk-node! '() sk 0 0 desc)))
    sk))

(define (add-interest! sk i)
  (let ((sk (extend-skeleton! sk (skeleton-interest-desc i))))
    (skcont-add! (skeleton-node-continuation sk) i)))

(define (remove-interest! sk i)
  (let ((sk (extend-skeleton! sk (skeleton-interest-desc i))))
    (skcont-remove! (skeleton-node-continuation sk) i)))

(define (skeleton-modify! sk term0 modify-skcont! on-missing-skconst modify-skconst! modify-skacc!)
  (unpack-scoped-assertion [restriction-path term0-term] term0)

  (define (walk-node! sk term-stack)
    (match-define (skeleton-node continuation edges) sk)

    (modify-skcont! continuation term0)
    (let ((sct (skeleton-continuation-table continuation))
          (constant-projections-to-remove '()))
      (hash-for-each sct
                     (lambda (constant-proj key-proj-handler)
                       (define constants (apply-projection term0-term constant-proj))
                       (define proj-handler
                         (hash-ref key-proj-handler
                                   constants
                                   (lambda () (on-missing-skconst key-proj-handler constants))))
                       (when proj-handler
                         (when (eq? (modify-skconst! proj-handler term0) 'remove-check)
                           (when (skeleton-matched-constant-empty? proj-handler)
                             (hash-remove! key-proj-handler constants)
                             (when (hash-empty? key-proj-handler)
                               (set! constant-projections-to-remove
                                     (cons constant-proj constant-projections-to-remove)))))
                         (hash-for-each (skeleton-matched-constant-table proj-handler)
                                        (lambda (variable-proj acc)
                                          (when (unrestricted? variable-proj restriction-path)
                                            (define vars (apply-projection term0-term variable-proj))
                                            (modify-skacc! acc vars term0)))))))
      (for-each (lambda (constant-proj) (hash-remove! sct constant-proj))
                constant-projections-to-remove))

    (for [(edge (in-list edges))]
      (match-define (cons (skeleton-selector pop-count index) table) edge)
      (define popped-stack (drop term-stack pop-count))
      (define old-top (car popped-stack))
      (define new-top
        (cond [(non-object-struct? old-top) (vector-ref (struct->vector old-top) (+ index 1))]
              [(list? old-top) (list-ref old-top index)]
              [(vector? old-top) (vector-ref old-top index)]))
      (define entry (hash-ref table
                              (cond [(non-object-struct? new-top) (struct->struct-type new-top)]
                                    [(list? new-top) (list-type (length new-top))]
                                    [(vector? new-top) (vector-type (vector-length new-top))]
                                    [else #f])
                              #f))
      (when entry (walk-node! entry (cons new-top popped-stack)))))

  (walk-node! sk (list (list term0-term))))

(define (add-term-to-skcont! skcont term)
  (hash-set! (skeleton-continuation-cache skcont) term #t))
(define (add-term-to-skconst! skconst term)
  (hash-set! (skeleton-matched-constant-cache skconst) term #t))
(define (add-term-to-skacc! skacc vars _term)
  ;; (log-info ">>>>>> At addition time for ~v, cache has ~v"
  ;;           _term
  ;;           (hash-ref (skeleton-accumulator-cache skacc) vars 0))
  (match (bag-change! (skeleton-accumulator-cache skacc) vars 1)
    ['absent->present
     (hash-for-each (skeleton-accumulator-handlers skacc)
                    (lambda (handler _) (apply handler '+ vars)))]
    ;; 'present->absent and 'absent->absent absurd
    ['present->present
     (void)]))

(define (add-assertion! sk term)
  (skeleton-modify! sk
                    term
                    add-term-to-skcont!
                    (lambda (cv-table cv)
                      (let ((sc (make-empty-matched-constant)))
                        (hash-set! cv-table cv sc)
                        sc))
                    add-term-to-skconst!
                    add-term-to-skacc!))

(define (remove-term-from-skcont! skcont term)
  (hash-remove! (skeleton-continuation-cache skcont) term))
(define (remove-term-from-skconst! skconst term)
  (hash-remove! (skeleton-matched-constant-cache skconst) term)
  'remove-check)
(define (remove-term-from-skacc! skacc vars _term)
  (define cache (skeleton-accumulator-cache skacc))
  ;; (log-info ">>>>>> At removal time for ~v, cache has ~v" _term (hash-ref cache vars 0))
  (if (bag-member? cache vars)
      (match (bag-change! cache vars -1)
        ['present->absent
         (hash-for-each (skeleton-accumulator-handlers skacc)
                        (lambda (handler _) (apply handler '- vars)))]
        ;; 'absent->absent and 'absent->present absurd
        ['present->present
         (void)])
      (log-warning "Removing assertion not previously added: ~v" _term)))

(define (remove-assertion! sk term)
  (skeleton-modify! sk
                    term
                    remove-term-from-skcont!
                    (lambda (_cv-table _cv) #f)
                    remove-term-from-skconst!
                    remove-term-from-skacc!))

(define (send-assertion! sk term)
  (skeleton-modify! sk
                    term
                    void
                    (lambda (_cv-table _cv) #f)
                    void
                    (lambda (skacc vars _term)
                      (hash-for-each (skeleton-accumulator-handlers skacc)
                                     (lambda (handler _) (apply handler '! vars))))))

;; TODO: avoid repeated descent into `term` by factoring out prefixes of paths in `proj`
(define (apply-projection term proj)
  (for/list [(path (in-list proj))]
    (apply-projection-path term path)))

(define (apply-projection-path term path)
  (for/fold [(term term)] [(index (in-list path))]
    (cond [(non-object-struct? term) (vector-ref (struct->vector term) (+ index 1))]
          [(list? term) (list-ref term index)]
          [(vector? term) (vector-ref term index)]
          [else (error 'apply-projection "Term representation not supported: ~v" term)])))