(module lazy-seq

(lazy-seq make-lazy-seq lazy-null
 lazy-seq? lazy-seq-realized? lazy-null? 
 lazy-seq->list list->lazy-seq lazy-list
 lazy-head lazy-tail lazy-length
 lazy-append lazy-reverse
 lazy-take lazy-drop lazy-ref
 lazy-take-while lazy-drop-while
 lazy-map lazy-filter lazy-each
 lazy-iterate lazy-repeat lazy-repeatedly
 lazy-numbers input-port->lazy-seq
 lazy-cycle lazy-append-map)

(import chicken scheme)
(use srfi-1 extras)

(define-record lazy-seq
  body value)

(define %make-lazy-seq
  make-lazy-seq)

(define (make-lazy-seq body)
  (%make-lazy-seq body #f))

(define-syntax lazy-seq
  (syntax-rules ()
    ((_ body ...)
     (make-lazy-seq
      (lambda () body ...)))))

(define-record-printer (lazy-seq seq out)
  (display "#<lazy-" out)
  (cond ((not (lazy-seq-realized? seq))
         (display "seq ...>" out))
        ((lazy-null? seq)
         (display "null>" out))
        (else
         (display "seq" out)
         (let loop ((seq seq))
           (if (lazy-seq-realized? seq)
               (if (lazy-null? seq)
                   (display ">" out)
                   (begin
                     (display " " out)
                     (write (lazy-head seq) out)
                     (loop (lazy-tail seq))))
               (display " ...>" out))))))

(define (lazy-seq-realized? seq)
  (not (lazy-seq-body seq)))

(define lazy-null
  (lazy-seq '()))

(define (lazy-null? seq)
  (null? (realized-lazy-seq seq)))

(define (realized-lazy-seq seq)
  (or (lazy-seq-value seq)
      (let ((value ((lazy-seq-body seq))))
        (lazy-seq-body-set! seq #f)
        (let loop ((value value))
          (if (or (null? value) (pair? value))
              (begin
                (lazy-seq-value-set! seq value)
                value)
              (loop (or (lazy-seq-value value)
                        ((lazy-seq-body value)))))))))

(define (lazy-head seq)
  (car (realized-lazy-seq seq)))

(define (lazy-tail seq)
  (cdr (realized-lazy-seq seq)))

(define (lazy-seq->list seq)
  (if (lazy-null? seq)
      '()
      (cons (lazy-head seq)
            (lazy-seq->list
             (lazy-tail seq)))))

(define (lazy-list . elements)
  (list->lazy-seq elements))

(define (list->lazy-seq list)
  (if (null? list)
      (%make-lazy-seq #f '())
      (%make-lazy-seq
       #f (cons (car list)
                (list->lazy-seq (cdr list))))))

(define (lazy-length seq)
  (let loop ((count 0) (seq seq))
    (if (lazy-null? seq)
        count
        (loop (+ count 1) (lazy-tail seq)))))

(define (lazy-take n seq)
  (lazy-seq
    (if (or (zero? n) (lazy-null? seq))
        '()
        (cons (lazy-head seq)
              (lazy-take (- n 1) (lazy-tail seq))))))

(define (lazy-drop n seq)
  (lazy-seq
    (if (or (zero? n) (lazy-null? seq))
        seq
        (lazy-drop (- n 1) (lazy-tail seq)))))

(define (lazy-take-while pred? seq)
  (let loop ((seq seq))
    (lazy-seq
      (cond ((lazy-null? seq) '())
            ((pred? (lazy-head seq))
             (cons (lazy-head seq)
                   (loop (lazy-tail seq))))
            (else '())))))

(define (lazy-drop-while pred? seq)
  (let loop ((seq seq))
    (lazy-seq
      (cond ((lazy-null? seq) '())
            ((pred? (lazy-head seq))
             (loop (lazy-tail seq)))
            (else seq)))))

(define (lazy-numbers #!key (step 1) (start 0) count)
  (let loop ((count count) (start start) (step step))
    (lazy-seq
      (if (and count (zero? count))
          '()
          (cons start
                (loop (and count (- count 1))
                      (+ start step)
                      step))))))

(define (lazy-append . seqs)
  (let loop ((seqs seqs))
    (lazy-seq
      (if (null? seqs)
          '()
          (let loop2 ((seq (car seqs)))
            (lazy-seq
              (if (lazy-null? seq)
                  (loop (cdr seqs))
                  (cons (lazy-head seq)
                        (loop2 (lazy-tail seq))))))))))

(define (make-lazy-mapping-proc append-result)
  (case-lambda
   ((proc seq)
    (let loop ((seq seq))
      (lazy-seq
        (if (lazy-null? seq)
            '()
            (append-result
             (proc (lazy-head seq))
             (loop (lazy-tail seq)))))))
   ((proc seq . seqs)
    (let loop ((seqs (cons seq seqs)))
      (lazy-seq
        (if (any lazy-null? seqs)
            '()
            (append-result
             (apply proc (map lazy-head seqs))
             (loop (map lazy-tail seqs)))))))))

(define lazy-map
  (make-lazy-mapping-proc cons))

(define lazy-append-map
  (make-lazy-mapping-proc lazy-append))

(define (lazy-filter pred? seq)
  (let loop ((seq seq))
    (lazy-seq
      (if (lazy-null? seq)
          '()
          (let ((head (lazy-head seq))
                (tail (loop (lazy-tail seq))))
            (if (pred? head)
                (cons head tail)
                tail))))))

(define (lazy-ref n seq)
  (if (zero? n)
      (lazy-head seq)
      (lazy-ref (- n 1) (lazy-tail seq))))

(define (lazy-each proc . seqs)
  (unless (any lazy-null? seqs)
    (apply proc (map lazy-head seqs))
    (apply lazy-each proc (map lazy-tail seqs))))

(define (input-port->lazy-seq port read)
  (let loop ()
    (lazy-seq
      (let ((datum (read port)))
        (if (eof-object? datum)
            '()
            (cons datum (loop)))))))

(define (lazy-repeat x)
  (lazy-seq (cons x (lazy-repeat x))))

(define (lazy-repeatedly f)
  (lazy-seq (cons (f) (lazy-repeatedly f))))

(define (lazy-iterate f x)
  (lazy-seq
    (cons x (lazy-iterate f (f x)))))

(define (lazy-reverse seq)
  (let loop ((seq seq) (rev-seq (lazy-seq '())))
    (lazy-seq
      (if (lazy-null? seq)
          rev-seq
          (loop (lazy-tail seq)
                (%make-lazy-seq #f (cons (lazy-head seq) rev-seq)))))))

(define (lazy-cycle seq)
  (lazy-seq
    (if (lazy-null? seq)
        '()
        (let loop ((rest seq))
          (lazy-seq
            (if (lazy-null? rest)
                (loop seq)
                (cons (lazy-head rest)
                      (loop (lazy-tail rest)))))))))

)