;; lib/minikanren/stream.sx — Phase 2 piece A: lazy streams of substitutions.
;;
;; SX has no improper pairs (cons requires a list cdr), so we use a
;; tagged stream-cell shape for mature stream elements:
;;
;;   stream  ::= mzero                        empty (the SX empty list)
;;             | (:s HEAD TAIL)               mature cell, TAIL is a stream
;;             | thunk                        (fn () ...) → stream when forced
;;
;; HEAD is a substitution dict. TAIL is again a stream (possibly a thunk),
;; which is what gives us laziness — mk-mplus can return a mature head with
;; a thunk in the tail, deferring the rest of the search.

(define mzero (list))

(define s-cons (fn (h t) (list :s h t)))

(define
  s-cons?
  (fn (s) (and (list? s) (not (empty? s)) (= (first s) :s))))

(define s-car (fn (s) (nth s 1)))
(define s-cdr (fn (s) (nth s 2)))

(define unit (fn (s) (s-cons s mzero)))

(define stream-pause? (fn (s) (and (not (list? s)) (callable? s))))

;; mk-mplus — interleave two streams. If s1 is paused we suspend and
;; swap (Reasoned Schemer "interleave"); otherwise mature-cons head with
;; mk-mplus of the rest.
(define
  mk-mplus
  (fn
    (s1 s2)
    (cond
      ((empty? s1) s2)
      ((stream-pause? s1) (fn () (mk-mplus s2 (s1))))
      (:else (s-cons (s-car s1) (mk-mplus (s-cdr s1) s2))))))

;; mk-bind — apply goal g to every substitution in stream s, mk-mplus-ing.
(define
  mk-bind
  (fn
    (s g)
    (cond
      ((empty? s) mzero)
      ((stream-pause? s) (fn () (mk-bind (s) g)))
      (:else (mk-mplus (g (s-car s)) (mk-bind (s-cdr s) g))))))

;; stream-take — force up to n results out of a (possibly lazy) stream
;; into a flat SX list of substitutions. n = -1 means take all.
(define
  stream-take
  (fn
    (n s)
    (cond
      ((= n 0) (list))
      ((empty? s) (list))
      ((stream-pause? s) (stream-take n (s)))
      (:else
        (cons
          (s-car s)
          (stream-take
            (if (= n -1) -1 (- n 1))
            (s-cdr s)))))))