honeylisp/stream.fnl

(var make-stream nil)
(fn stream [table]
  (make-stream
    {:table (or table [])
     :i 0
     :n (length table)
     :step 1
     :next 
     (fn [self]
       (when (~= self.i self.n)
         (set self.i (+ self.i self.step))
         true))
     :curr (fn [self] (. self.table self.i))
     :reverse
     (fn [self]
       (local prev-i self.i)
       (set self.i (+ self.n self.step))
       (set self.n (+ prev-i self.step))
       (set self.step (* self.step -1))
       self)}))

(fn kvstream [table]
  (make-stream
    {:table (or table {})
     :curr-key nil
     :curr-val :start
     :curr (fn [self] (values self.curr-key self.curr-val))
     :keys (fn [self] (self:map (fn [k v] k)))
     :values (fn [self] (self:map (fn [k v] v)))
     :next
     (fn [self]
       (when self.curr-val
         (set (self.curr-key self.curr-val) (next self.table self.curr-key)))
       (~= self.curr-key nil))}))

(fn one [...]
  (make-stream
    {:vals [...]
     :advanced false
     :curr (fn [self] (unpack self.vals))
     :next
     (fn [self]
       (if self.advanced
         false
         (do
           (set self.advanced true)
           true)))}))

(fn iter [self]
  (values
    (fn [self _]
      (when (self:next)
        (self:curr)))
    self
    nil))

(fn first [self]
  (if (self:next) (self:curr) nil))

(fn map [stream f]
  (make-stream
    {: stream : f
     :curr (fn [self] (self.f (self.stream:curr)))
     :next (fn [self] (self.stream:next))
     :reverse
     (fn [self]
       (set self.stream (self.stream:reverse))
       self)}))

(fn filter [stream f]
  (make-stream
    {: stream : f :curr-val nil
     :curr (fn [self] self.curr-val)
     :next
     (fn [self]
       (set self.curr-val nil)
       (var has-more (self.stream:next))
       (while has-more
         (let [curr (self.stream:curr)
               include? (self.f curr)]
           (when include?
             (set self.curr-val curr))
           (set has-more (if include? false (self.stream:next)))))
       self.curr-val)}))

(fn reduce [stream f init]
  (var val init)
  (each [v (stream:iter)]
    (set val (f val v)))
  val)

(fn flatten [stream]
  (make-stream
    {: stream
     :curr-stream nil
     :curr (fn [self] (self.curr-stream:curr))
     :next
     (fn [self]
       (var reached-next 
         (if self.curr-stream (self.curr-stream:next) false))
       (while (not reached-next)
         (set self.curr-stream 
           (if (self.stream:next) (self.stream:curr) nil))
         (set reached-next
           (if self.curr-stream (self.curr-stream:next) true)))
       (~= self.curr-stream nil))}))

(fn concat [s ...]
  (: (stream [s ...]) :flatten))

(fn skip [stream n]
  (make-stream
    {: stream
     : n
     :curr (fn [self] (self.stream:curr))
     :next
     (fn [self]
       (for [_ 1 self.n]
         (self.stream:next))
       (set self.n 0)
       (self.stream:next))}))

(fn take [stream n]
  (make-stream
    {: stream
     : next
     :curr (fn [self] (self.stream:curr))
     :next
     (fn [self]
       (if (> self.n 0)
         (do
           (set self.n (- self.n 1))
           (self.stream:next))
         false))}))

(fn tolist [stream]
  (let [l []]
    (each [v (stream:iter)]
      (table.insert l v))
    l))

(fn tomap [stream]
  (let [m {}]
    (each [k v (stream:iter)]
      (tset m k v))
    m))

(set make-stream 
  (fn [stream]
    (set stream.iter iter)
    (set stream.map map)
    (set stream.filter filter)
    (set stream.first first)
    (set stream.reduce reduce)
    (set stream.flatten flatten)
    (set stream.concat concat)
    (set stream.skip skip)
    (set stream.take take)
    (set stream.tolist tolist)
    (set stream.tomap tomap)
    stream))

{: stream : kvstream : one}