Clojure: Using a stream without holding head. Is this correct?

Question

I want to put a lazy sequence somewhere to provide data as needed. I know I have to avoid keeping hold of the head of the sequence. I came up with the follwoing solution, did I miss something?

(defn headless [s] 
  (let [a (atom s)] 
    (fn 
      ([]  (let [s @a r (first s)] 
             (swap! a rest) r)) 
      ([n] (let [s @a rs (take n s)] 
             (swap! a #(drop n %)) rs))))) 

Usage example, this simple generator just gives the natural numbers.

(def nums (headless (iterate inc 0)))

(nums 5)
; (0 1 2 3 4)

(nums)
;5

Update: The "test" should use dorun, not doall. See solution by lgrapenthin

A (not overly realistic) test with

(doall (map #(nums %) (repeat 20)))

crashed after 5 min usage of all 4 cores with an Exception (OutOfMemoryError Java heap space)

Answer 1

Your code works.

This form:

(doall (map #(nums %) (repeat 20)))

will generate an infinite amount of (nums 20) and never return. You can use dorun instead, to drop the generated (nums 20) and not keep them in memory. However, it will not return because (repeat 20) generates an infinite lazy seq.

Slightly more readable version of headless

(defn headless [s] 
  (let [a (atom s)] 
    (fn 
      ([]  (let [s @a]
             (swap! a rest)
             (first s))
      ([n] (let [s @a]
             (swap! a (partial drop n))
             (take n s)))))))