counting_partition.clj

;;Not an amazing name
(defn partition*
  "Given a function that takes an input and returns a count n, will take an item from a sequence,
  call f on it to get count n, add that input to a partition along with the next n items from the coll"
  ([f]
   (fn [rf]
     (let [part (volatile! [])
           size (volatile! ::not-set)]
       (fn
         ([] (rf) )
         ([result] (rf result))
         ([result input]

          (vswap! part conj input)

          (when (identical? @size ::not-set)
            (vreset! size (inc (f input))))

          (if (= (count @part) @size)
            (let [v @part]
              (vswap! part empty)
              (vreset! size ::not-set)
              (rf result v))
            result))))))
  ([f coll]
   (lazy-seq
    (when-let [n (f (first coll))]
      (cons (take (inc n) coll)
            (when (not-empty coll)
              (partition* f (drop (inc n) coll))))))))

(deftest test-partition*
  (let [lookup    {5 3, 8 1, 0 5}
        stream    [5 9 0 7, 8 1, 0 4 5 8 9 5 ]
        c (chan 1 (partition* lookup))
        ]
    (is (= '((5 9 0 7) (8 1) (0 4 5 8 9 5))
              (partition* lookup stream)))
    (is (= [[5 9 0 7] [8 1] [0 4 5 8 9 5]]
           (transduce (partition* lookup) conj stream)))

    (a/onto-chan c stream)
    (is (= [[5 9 0 7] [8 1] [0 4 5 8 9 5]]
           (<!! (a/into [] c))))))