;; lib/minikanren/tabling.sx — Phase 7 piece A: naive memoization.
;;
;; A `table-2` wrapper for 2-arg relations (input, output). Caches by
;; ground input (walked at call time). On hit, replays the cached output
;; values; on miss, runs the relation, collects all output values from
;; the answer stream, stores, then replays.
;;
;; Limitations of naive memoization (vs proper SLG / producer-consumer
;; tabling):
;;   - Each call must terminate before its result enters the cache —
;;     so cyclic recursive calls with the SAME ground input would still
;;     diverge (not addressed here).
;;   - Caching by full ground walk only; partially-ground args fall
;;     through to the underlying relation.
;;
;; Despite the limitations, naive memoization is enough for the
;; canonical demo: Fibonacci goes from exponential to linear because
;; each fib(k) result is computed at most once.
;;
;; Cache lifetime: a single global mk-tab-cache. Use `(mk-tab-clear!)`
;; between independent queries.

(define mk-tab-cache {})

(define mk-tab-clear! (fn () (set! mk-tab-cache {})))

(define
  mk-tab-lookup
  (fn
    (key)
    (cond
      ((has-key? mk-tab-cache key) (get mk-tab-cache key))
      (:else :miss))))

(define
  mk-tab-store!
  (fn (key vals) (set! mk-tab-cache (assoc mk-tab-cache key vals))))

(define
  mk-tab-ground-term?
  (fn
    (t)
    (cond
      ((is-var? t) false)
      ((mk-cons-cell? t)
        (and
          (mk-tab-ground-term? (mk-cons-head t))
          (mk-tab-ground-term? (mk-cons-tail t))))
      ((mk-list-pair? t) (every? mk-tab-ground-term? t))
      (:else true))))

(define
  mk-tab-replay-vals
  (fn
    (vals output s)
    (cond
      ((empty? vals) mzero)
      (:else
        (let
          ((sp (mk-unify output (first vals) s)))
          (let
            ((this-stream (cond ((= sp nil) mzero) (:else (unit sp)))))
            (mk-mplus this-stream (mk-tab-replay-vals (rest vals) output s))))))))

(define
  table-2
  (fn
    (name rel-fn)
    (fn
      (input output)
      (fn
        (s)
        (let
          ((winput (mk-walk* input s)))
          (cond
            ((mk-tab-ground-term? winput)
              (let
                ((key (str name "@" winput)))
                (let
                  ((cached (mk-tab-lookup key)))
                  (cond
                    ((= cached :miss)
                      (let
                        ((all-substs (stream-take -1 ((rel-fn input output) s))))
                        (let
                          ((vals (map (fn (s2) (mk-walk* output s2)) all-substs)))
                          (begin
                            (mk-tab-store! key vals)
                            (mk-tab-replay-vals vals output s)))))
                    (:else (mk-tab-replay-vals cached output s))))))
            (:else ((rel-fn input output) s))))))))

;; --- table-1: 1-arg relation (one input, no output to cache) ---
;; The relation is a predicate `(p input)` that succeeds or fails.
;; Cache stores either :ok or :no.

(define
  table-1
  (fn
    (name rel-fn)
    (fn
      (input)
      (fn
        (s)
        (let
          ((winput (mk-walk* input s)))
          (cond
            ((mk-tab-ground-term? winput)
              (let
                ((key (str name "@1@" winput)))
                (let
                  ((cached (mk-tab-lookup key)))
                  (cond
                    ((= cached :miss)
                      (let
                        ((stream ((rel-fn input) s)))
                        (let
                          ((peek (stream-take 1 stream)))
                          (cond
                            ((empty? peek)
                              (begin (mk-tab-store! key :no) mzero))
                            (:else (begin (mk-tab-store! key :ok) stream))))))
                    ((= cached :ok) (unit s))
                    ((= cached :no) mzero)
                    (:else mzero)))))
            (:else ((rel-fn input) s))))))))

;; --- table-3: 3-arg relation (input1 input2 output) ---
;; Cache keyed by (input1, input2). Output values cached as a list.

(define
  table-3
  (fn
    (name rel-fn)
    (fn
      (i1 i2 output)
      (fn
        (s)
        (let
          ((wi1 (mk-walk* i1 s)) (wi2 (mk-walk* i2 s)))
          (cond
            ((and (mk-tab-ground-term? wi1) (mk-tab-ground-term? wi2))
              (let
                ((key (str name "@3@" wi1 "/" wi2)))
                (let
                  ((cached (mk-tab-lookup key)))
                  (cond
                    ((= cached :miss)
                      (let
                        ((all-substs (stream-take -1 ((rel-fn i1 i2 output) s))))
                        (let
                          ((vals (map (fn (s2) (mk-walk* output s2)) all-substs)))
                          (begin
                            (mk-tab-store! key vals)
                            (mk-tab-replay-vals vals output s)))))
                    (:else (mk-tab-replay-vals cached output s))))))
            (:else ((rel-fn i1 i2 output) s))))))))