;; Closures that outlive the frame they were made in. The environment is
;; allocated by the collector, so a capturing fn may be returned, kept in a
;; Vec, an Option field of a struct or a global, and called after any number
;; of collections. Every line below is called after a forced collection that
;; would have freed an environment nobody rooted.
;;
;; Capture is by value: the environment holds copies of the locals taken when
;; the fn was made. A counter shares state through a captured dyn map, which
;; is a reference to one object on the collector's heap.

(declare gc-collect [] () "flan_gc_collect")

(defn churn [] ()
  (let [i 0]
    (while (< i 2000)
      (let [v (vec-new dyn)] (push v i) (push v "junk"))
      (set i (+ i 1))))
  (gc-collect))

(defn make-adder [n i64] (Fn [i64] i64)
  (fn [x] (+ x n)))

;; Two captures deep: the inner fn captures the outer one's copy of [base],
;; and the returned value captures a function value that captures.
(defn make-scaled [base i64 k i64] (Fn [i64] i64)
  (let [add (make-adder base)]
    (fn [x] (* k (add x)))))

(defn keep [f (Fn [i64] i64)] (Fn [i64] i64) f)

;; A function value read back out of an environment and handed on: [sneak]'s
;; literal captures [g], and [getf] returns the copy.
(defn getf [f (Fn [] (Fn [i64] i64))] (Fn [i64] i64) (f))
(defn sneak [g (Fn [i64] i64)] (Fn [i64] i64) (getf (fn [] g)))

;; Stored through a pointer into a slot of the caller's frame.
(defn stash [p (Ptr (Fn [i64] i64)) f (Fn [i64] i64)] () (set (deref p) f))

(defn make-counter [] (Fn [] i64)
  (let [st {:n 0}]
    (fn []
      (put st :n (+ (get st :n) 1))
      (i64 (get st :n)))))

;; A dyn captured and returned: the text is on the collector's heap and
;; reachable only through the environment.
(defn make-greeter [who] (Fn [] i64)
  (let [msg (vec-new dyn)]
    (push msg who)
    (push msg "and")
    (push msg who)
    (fn [] (i64 (length msg)))))

(defstruct Button [label string on-click (Option (Fn [] i64))])

(defonce handler (Option (Fn [i64] i64)))

(defn call-opt [o (Option (Fn [i64] i64)) x i64] i64
  (match o
    (Some f) (f x)
    None     -1))

(defn double [x i64] i64 (* 2 x))

;; A closure made as an argument and held while the next argument collects —
;; kept by the callee, so it is one the collector owns — and one collected
;; for while the callee runs, before the callee has put it anywhere. A
;; parameter of function type is not rooted by the callee; the caller holds
;; it.
(defn apply-to [f (Fn [i64] i64) x i64] i64 (f x))
(defn churn-1 [] i64 (churn) 1)
(defonce last-fn (Option (Fn [i64] i64)))
(defn remember [f (Fn [i64] i64) x i64] i64
  (churn)
  (set last-fn (Some f))
  (f x))

;; A handler clause keeps its copies on the establishing frame, and may now
;; capture a dyn and a closure like an fn may.
(defstruct Ping [n i64])
(defonce heard i64)
(defn pinged [x i64] i64 (signal (Ping {.n x})) x)

(defn keep-handled [f (Fn [i64] i64)] (Fn [i64] i64)
  (handler-bind [(Ping [c] (set heard 0))] f))

(defn handled [] ()
  (let [names (vec-new dyn)
        f (make-adder 50)]
    (push names "a")
    (push names "b")
    (handler-bind [(Ping [c] (churn) (set heard (+ (f (.n c)) (i64 (length names)))))]
      (pinged 1))
    (println heard)))

;; A data type's case holding a function value, in a Vec, and a Vec of Vecs.
;; The collector names every case's function-value word and follows only the
;; ones that hold an environment it made.
(defdata Action
  [Idle
   (Run  [f (Option (Fn [i64] i64))  n i64])
   (Pair [k i64  g (Option (Fn [i64] i64))])])

(defn act [a Action x i64] i64
  (match a
    Idle 0
    (Run f n) (+ n (call-opt f x))
    (Pair k g) (+ k (call-opt g x))))

(defn shapes [] ()
  (let [acts (vec-new Action)
        a (arena-new 65536)
        outer (vec-new (Vec (Fn [i64] i64)) a)]
    (push acts (Action.Run {.f (Some (make-adder 1)) .n 10}))
    (push acts (Action.Pair {.k 20 .g (Some (make-adder 2))}))
    (push acts Action.Idle)
    (let [inner (vec-new (Fn [i64] i64) a)]
      (push inner (make-adder 3))
      (push outer inner))
    (churn)
    (println (+ (act (at acts 0) 1) (+ (act (at acts 1) 1) (act (at acts 2) 1)))
             ((at (at outer 0) 0) 1))))

;; A data type holding a Vec of itself: the collector walks as deep as the
;; data goes, and the function value two levels down is still found.
(defdata Tree [(Node [f (Option (Fn [i64] i64))  kids (Vec Tree)])])

(defn tree-sum [t Tree x i64] i64
  (match t
    (Node f kids)
    (let [s (call-opt f x)
          i 0]
      (while (< i (length kids))
        (set s (+ s (tree-sum (at kids i) x)))
        (set i (+ i 1)))
      s)))

(defn trees [] ()
  (let [a (arena-new 65536)
        leaf-kids (vec-new Tree a)
        kids (vec-new Tree a)]
    (push kids (Tree.Node {.f (Some (make-adder 5)) .kids leaf-kids}))
    (let [root (Tree.Node {.f None .kids kids})]
      (churn)
      (println (tree-sum root 1)))))

;; Environments nobody holds are collected: a hundred thousand of them made
;; and dropped leave the heap as small as it was.
(declare gc-live-bytes [] i64 "flan_gc_live_bytes")

(defn many [] ()
  (let [sum (i64 0)]
    (dotimes [i 100000]
      (set sum (+ sum ((make-adder 1) 0))))
    (gc-collect)
    (println sum (< (gc-live-bytes) 1000000))))

(defn main [] i32
  ;; Returned, then called after a collection.
  (let [add5 (make-adder 5)]
    (churn)
    (println (add5 10)))
  ;; Nested, and passed through a function that hands its parameter back.
  (let [f (keep (make-scaled 1 3))]
    (churn)
    (println (f 4)))
  ;; Out of an environment, through a handler-bind's value, and through a
  ;; pointer.
  (let [f (keep-handled (sneak (make-adder 20)))
        slot (make-adder 0)]
    (stash (addr slot) (make-adder 7))
    (churn)
    (println (f 1) (slot 1)))
  ;; Held beside a sibling operand that collects.
  (let [n 40]
    (println (apply-to (fn [x] (+ x n)) (churn-1))
             (remember (fn [x] (+ x n)) 1)))
  ;; A Vec of function values, one capture per iteration plus a widened name.
  (let [fs (vec-new (Fn [i64] i64))]
    (let [i 0]
      (while (< i 5)
        (push fs (make-adder (* i 100)))
        (set i (+ i 1))))
    (push fs double)
    (churn)
    (let [total (i64 0)]
      (let [i 0]
        (while (< i (length fs))
          (set total (+ total ((at fs i) 1)))
          (set i (+ i 1))))
      (println total)))
  ;; A counter over a captured dyn map: two calls on one value see one state.
  (let [c (make-counter)]
    (c)
    (churn)
    (c)
    (println (c)))
  ;; A captured dyn.
  (let [g (make-greeter "world")]
    (churn)
    (println (g)))
  ;; A struct field and a global, both through Option.
  (let [b (Button {.label "ok" .on-click (Some (make-counter))})]
    (churn)
    (match (.on-click b)
      (Some f) (do (f) (println (f)))
      None     (println "none")))
  (set handler (Some (make-adder 1000)))
  (churn)
  (println (call-opt handler 7))
  (handled)
  (shapes)
  (trees)
  (many)
  0)