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

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

(provide make-empty-skeleton
         add-interest!
         remove-interest!
         add-assertion!
         remove-assertion!
         send-assertion!)

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

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

(require preserves)

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

;; 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 SkMove (MutableHash ConstructorSpec Skeleton)))
;;    SkMove = (skeleton-move Nat (Listof Any))
;;
;; A `ConstructorSpec` specifies a record and its label, or a sequence,
;; or a dictionary.
;;
;;  ConstructorSpec = (U (GroupType-rec any) (GroupType-arr) (GroupType-dict))
;;
(struct skeleton-node (continuation [edges #:mutable]) #:transparent)
(struct skeleton-move (pop-count path) #:transparent)
;;
;; A `Pattern` is a pattern over assertions, following the schema of
;; the same name in schemas/dataspacePatterns.prs. Instances of
;; `Pattern` are usually produced by the `:pattern` macro and related
;; tooling.
;;
;; 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
;;                 (MutableHash Assertion #t)
;;                 (MutableHash (constant-positions SkProj SkProj)
;;                              (MutableHash SkKey SkConst)))
;;     SkConst = (skeleton-matched-constant
;;                 (MutableHash Assertion #t)
;;                 (MutableHash SkProj SkAcc))
;;       SkAcc = (skeleton-accumulator
;;                 (MutableBag SkKey)
;;                 (MutableHasheq EntityRef (MutableHash SkKey Handle)))
;;
(struct skeleton-continuation (cache table) #:transparent)
(struct constant-positions (with-values required-to-exist) #:transparent)
(struct skeleton-matched-constant (cache table) #:transparent)
(struct skeleton-accumulator (cache peers) #: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 Any))
;;
;; The outer list specifies elements of the flat list; the inner lists
;; specify paths via 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 entity-refs in
;; the associated `SkAcc`.

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

(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! turn c pat ref)
  (define cs (pattern->constant-paths pat))
  (define ce (pattern->check-paths pat))
  (define (classify-assertions)
    (define cvt (make-hash))
    (hash-for-each (skeleton-continuation-cache c)
                   (lambda (a _)
                     (unless (void? (apply-projection a ce))
                       (define avs (apply-projection a cs))
                       (unless (void? avs)
                         (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) (constant-positions cs ce) classify-assertions))
  (define sc (hash-ref! cvt (pattern->constant-values pat) make-empty-matched-constant))
  (define vs (pattern->capture-paths pat))
  (define (make-accumulator)
    (define cache (make-bag))
    (hash-for-each (skeleton-matched-constant-cache sc)
                   (lambda (term _) (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))
  (define peer-entry (make-hash))
  (hash-set! (skeleton-accumulator-peers acc) ref peer-entry)
  (for [(vars (in-bag (skeleton-accumulator-cache acc)))]
    (hash-set! peer-entry vars (turn-assert! turn ref 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! turn c pat ref)
  (define ck (constant-positions (pattern->constant-paths pat)
                                 (pattern->check-paths pat)))
  (define cvt (hash-ref (skeleton-continuation-table c) ck #f))
  (when cvt
    (define cv (pattern->constant-values pat))
    (define sc (hash-ref cvt cv #f))
    (when sc
      (define vs (pattern->capture-paths pat))
      (define acc (hash-ref (skeleton-matched-constant-table sc) vs #f))
      (when acc
        (for [(handle (in-hash-values (hash-ref (skeleton-accumulator-peers acc) ref)))]
          (turn-retract! turn handle))
        (hash-remove! (skeleton-accumulator-peers acc) ref)
        (when (hash-empty? (skeleton-accumulator-peers 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) ck))))

(define (⊖ h1 h0)
  (define-values (h1tail h0tail) (drop-common-prefix h1 h0))
  (skeleton-move (length h0tail) h1tail))

(define (⊕ h0 move)
  (match-define (skeleton-move n h) move)
  (append (drop-right h0 n) h))

(define (shape->visit s)
  (let walk ((path '()) (s s))
    (match s
      ['() '()]
      [(cons (cons h ctor-spec) more)
       (cons (cons (⊖ h path) ctor-spec) (walk h more))])))

(define (term-matches-ctor-spec? term ctor-spec)
  (match ctor-spec
    [(GroupType-rec label)
     (and (non-object-struct? term)
          (equal? (struct-type-name (struct->struct-type term)) label))]
    [(GroupType-arr)
     (list? term)]
    [(GroupType-dict)
     (hash? term)]))

(define (subterm-matches-ctor-spec? term path ctor-spec)
  (term-matches-ctor-spec? (apply-projection-path term path) ctor-spec))

(define (skeleton-node-edge-table! sk move)
  (match (assoc move (skeleton-node-edges sk))
    [#f (let ((table (make-hash)))
          (set-skeleton-node-edges! sk (cons (cons move table) (skeleton-node-edges sk)))
          table)]
    [(cons _move table) table]))

(define (extend-skeleton! sk pat)
  (let visit ((h '()) (sk sk) (moves (shape->visit (pattern->shape pat))))
    (match moves
      ['() sk]
      [(cons (cons move ctor-spec) moves)
       (define table (skeleton-node-edge-table! sk move))
       (define path (⊕ h move))
       (define (make-skeleton-node-with-cache)
         (define unfiltered (skeleton-continuation-cache (skeleton-node-continuation sk)))
         (define filtered (make-hash))
         (hash-for-each unfiltered
                        (lambda (a _)
                          (when (subterm-matches-ctor-spec? a path ctor-spec)
                            (hash-set! filtered a #t))))
         (make-empty-skeleton/cache filtered))
       (define next (hash-ref! table ctor-spec make-skeleton-node-with-cache))
       (visit path next moves)])))

(define (add-interest! turn sk pat ref)
  (skcont-add! turn (skeleton-node-continuation (extend-skeleton! sk pat)) pat ref))

(define (remove-interest! turn sk pat ref)
  (skcont-remove! turn (skeleton-node-continuation (extend-skeleton! sk pat)) pat ref))

(define (skeleton-modify! turn sk term0 modify-skcont! on-missing-skconst modify-skconst! modify-skacc!)
  (let walk-node! ((sk sk) (term-stack (list term0)))
    (match-define (skeleton-node continuation edges) sk)
    (modify-skcont! continuation term0)

    (let ((sct (skeleton-continuation-table continuation))
          (constant-keys-to-remove '()))
      (for [((constant-key key-proj-handler) (in-hash sct))]
        (unless (void? (apply-projection term0 (constant-positions-required-to-exist constant-key)))
          (define constants (apply-projection term0 (constant-positions-with-values constant-key)))
          (unless (void? constants)
            (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-keys-to-remove (cons constant-key constant-keys-to-remove)))))
              (hash-for-each (skeleton-matched-constant-table proj-handler)
                             (lambda (variable-proj acc)
                               (define vars (apply-projection term0 variable-proj))
                               (modify-skacc! turn acc vars term0)))))))
      (for [(constant-key (in-list constant-keys-to-remove))]
        (hash-remove! sct constant-key)))

    (for [(edge (in-list edges))]
      (match-define (cons (skeleton-move pop-count path) table) edge)
      (define popped-stack (drop term-stack pop-count))
      (define old-top (car popped-stack))
      (define new-top (apply-projection-path old-top path))
      (define ctor-spec (cond [(non-object-struct? new-top)
                               (GroupType-rec (struct-type-name
                                               (struct->struct-type new-top)))]
                              [(list? new-top) (GroupType-arr)]
                              [(hash? new-top) (GroupType-dict)]
                              [else #f]))
      (define entry (hash-ref table ctor-spec #f))
      (when entry (walk-node! entry (cons new-top popped-stack))))))

(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! turn 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-peers skacc)
                    (lambda (ref peer-entry)
                      (hash-set! peer-entry vars (turn-assert! turn ref vars))))]
    ;; 'present->absent and 'absent->absent absurd
    ['present->present
     (void)]))

(define (add-assertion! turn sk term)
  (skeleton-modify! turn
                    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! turn 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))
  (match (bag-change! cache vars -1)
    ['present->absent
     (hash-for-each (skeleton-accumulator-peers skacc)
                    (lambda (_ref peer-entry)
                      (turn-retract! turn (hash-ref peer-entry vars))
                      (hash-remove! peer-entry vars)))]
    ['present->present
     (void)]
    ;; 'absent->absent absurd
    ['absent->present
     ;; 'absent->present should be absurd, but may be a programming error
     (bag-change! cache vars 1) ;; undo the change
     (log-warning "Removing assertion not previously added: ~v" _term)]))

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

(define (send-assertion! turn sk term)
  (skeleton-modify! turn
                    sk
                    term
                    void
                    (lambda (_cv-table _cv) #f)
                    void
                    (lambda (turn skacc vars _term)
                      (hash-for-each (skeleton-accumulator-peers skacc)
                                     (lambda (ref _peer-entry) (turn-message! turn ref vars))))))

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

(define (apply-projection-path term path)
  (for/fold [(term term)]
            [(key (in-list path))]
    (step-term term key)))

(define (index? k) (exact-nonnegative-integer? k))

(define (step-term term key)
  (cond [(non-object-struct? term)
         (define v (struct->vector term))
         (define k (+ key 1)) ;; skip the label
         (when (and (index? k) (< k (vector-length v))) (vector-ref v k))]
        [(list? term)
         (when (and (index? key) (< key (length term))) (list-ref term key))]
        [(hash? term)
         (hash-ref term key (void))]
        [else
         (error 'apply-projection "Term representation not supported: ~v" term)]))