; Author: Juergen Lorenz
; ju (at) jugilo (dot) de
;
; Copyright (c) 2012-2014, Juergen Lorenz, Moritz Heidkamp
; All rights reserved.
; 
; Redistribution and use in source and binary forms, with or without
; modification, are permitted provided that the following conditions are
; met:
; 
; Redistributions of source code must retain the above copyright
; notice, this list of conditions and the following disclaimer.
; 
; Redistributions in binary form must reproduce the above copyright
; notice, this list of conditions and the following disclaimer in the
; documentation and/or other materials provided with the distribution.
; 
; Neither the name of the author nor the names of its contributors may be
; used to endorse or promote products derived from this software without
; specific prior written permission. 
; 
; THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
; IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
; TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
; PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
; HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
; SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
; TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
; PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
; LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
; NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
; SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
; 
; Last update: May 03, 2014
;
(module lazy-lists
  (Lazy Make-lazy
   assume-in
   List->list list->List input->List
   First Rest Car Cdr Length Length-min Append Reverse
   List? Null? Realized? Reverse*
   List-infinite? Realize
   List-not-null? List-finite? Lists-one-finite?
   Take Drop Ref Take-while Drop-while Count-while
   Memp Member Memq Memv
   Equ? Equal? Eq? Eqv?
   Assp Assoc Assq Assv
   Map Filter Sieve For-each
   Iterate Repeat Repeatedly
   Cardinals Primes Cycle Range
   Nil Cons Merge Sort Sorted? Split-at Split-with
   vector->List List->vector
   Fold-right Fold-left Fold-right* Fold-left*
   Zip Every? Some? List
   Admissible?)

(import scheme
        (only data-structures o compress list-of?)
        (only chicken
              cond-expand
              define-record-type
              define-record-printer
              cut when case-lambda
              nth-value unless
              make-parameter error void add1 sub1
              fixnum? fx+ fx= fx>= fx< fx- fx/))

;;; Assert-checked muss als interpreter oder compiler-switch
;;; uebergeben werden:
;;; -feature Assert-checked
;(define-syntax Assert
;  (syntax-rules ()
;    ((_ xpr sym . args)
;     ;; this is done at runtime
;     ;; the whole cond-expand is compiled
;     (cond-expand
;       (Assert-checked
;         (unless xpr
;           (error sym "Assertion failed:" 'xpr . args)))
;       (else)))))
;;  (ir-macro-transformer
;;    (lambda (form inject compare?)
;;      (let ((xpr (cadr form))
;;            (sym (caddr form))
;;            (args (cdddr form)))
;;        ;; this is done at compile time
;;        ;; only one branch is compiled
;;        (cond-expand
;;          (Assert-checked
;;            `(if ,xpr
;;               (if #f #f)
;;               (error ,sym "Assertion failed:" ',xpr ,@args)))
;;          (else
;;            `(if #f #f)))))))
;; internal
(define-syntax assume
  (syntax-rules ()
    ((_ sym) sym)
    ((_ sym test)
     (if test sym 'test))
    ((_ sym test . tests)
     (if test
       (assume sym . tests)
       'test))))

;;; assumptions-checked should be provided by interpreter- or
;;; compiler-switch:
;;; -feature assumptions-checked
;;; or via function call
;;; (register-feature! 'assumptions-checked)
(define-syntax assume-in
  (syntax-rules ()
    ((_ sym test . tests)
     (cond-expand
       (assumptions-checked
         (let ((res (assume sym test . tests)))
           (unless (eq? res sym)
             (error sym "Assertion failed" res))))
       (else)))))

;; all defined operators hidden
(define lazy-list
  (define-record-type lazy-list
    (make-lazy-list length body value)
    lazy-list?
    (length lazy-list-length lazy-list-length-set!)
    (body lazy-list-body lazy-list-body-set!)
    (value lazy-list-value lazy-list-value-set!)))

(define-syntax Lazy
  (syntax-rules ()
    ((_ len xpr . xprs)
     (Make-lazy len (lambda () xpr . xprs)))))


(define (Make-lazy len thunk)
  (make-lazy-list len thunk #f))

(define (Cons var Lst)
  (assume-in 'Cons
    (List? Lst))
  (let ((len (lazy-list-length Lst)))
    (Lazy (if len (+ 1 len) #f)
      (cons var Lst))))

(define Length lazy-list-length)

(define (Length-min . Lsts)
  (assume-in 'Length-min
    ((list-of? List?) Lsts))
  (let* ((lens (map Length Lsts))
         (finites (compress lens lens)))
    (if (null? finites)
      #f
      (apply min finites))))

(define List? lazy-list?)

(define Nil
  (make-lazy-list 0 (lambda () '()) #f))

(define-record-printer (lazy-list Lst out)
  (assume-in 'define-record-printer
    (List? Lst))
  (display "#<List[" out)
  (display (lazy-list-length Lst) out)
  (display "]" out)
  (cond ((not (Realized? Lst))
         (display " ...>" out))
        ((Null? Lst)
         (display "]>" out))
        (else
         (let loop ((Lst Lst))
           (if (Realized? Lst)
               (if (Null? Lst)
                   (display ">" out)
                   (begin
                     (display " " out)
                     (write (First Lst) out)
                     (loop (Rest Lst))))
               (display " ...>" out))))))

;; hidden, does most of the dirty work
(define (realize Lst)
  (or (lazy-list-value Lst)
      (let ((value ((lazy-list-body Lst))))
        (lazy-list-body-set! Lst #f)
        (let loop ((value value))
          (if (or (null? value) (pair? value))
              (begin
                (lazy-list-value-set! Lst value)
                value)
              (loop (or (lazy-list-value value)
                        ((lazy-list-body value)))))))))

(define (Realize Lst)
  (assume-in 'Realize
    (List-finite? Lst))
  (let ((len (lazy-list-length Lst)))
    (when len
      (Ref (- len 1) Lst)
      Lst)))
  
(define (Realized? Lst)
  (assume-in 'Realized
    (List? Lst))
  (not (lazy-list-body Lst)))

(define (Null? Lst)
  (assume-in 'Null?
    (List? Lst))
  (null? (realize Lst)))

(define (First Lst)
  (assume-in 'First
    (List? Lst))
  (car (realize Lst)))

(define (Car Lst)
  (assume-in 'Car
    (List? Lst))
  (First Lst))

;; to speed up cdring for lists with preknown length
(define (rest Lst)
  (assume-in 'Car (List? Lst))
  (cdr (realize Lst)))

(define (Rest Lst)
  (assume-in 'Rest (List? Lst))
  (let (
    (len (lazy-list-length Lst))
    (Result (cdr (realize Lst)))
    )
    (lazy-list-length-set! Result (if len (fx- len 1) #f))
    Result))

(define (Cdr Lst)
  (assume-in 'Cdr (List? Lst))
  (Rest Lst))

(define (Admissible? n Lst)
  (assume-in 'Admissible?
    (List? Lst)
    (fixnum? n)
    (fx>= n 0))
  (let ((len (lazy-list-length Lst)))
    (or (not len) (fx< n len))))
  
(define (Ref n Lst)
  (assume-in 'Ref
    (Admissible? n Lst))
  (let ((len (lazy-list-length Lst)))
    (let loop ((n n) (Lst Lst))
      (if (fx= n 0)
        (First Lst)
        (loop (fx- n 1) (Rest Lst))))))

(define (List->list Lst)
  (assume-in 'List->list
    (List-finite? Lst))
  (let loop ((lst '()) (Lst Lst))
    (if (Null? Lst)
      (reverse lst)
      (loop (cons (First Lst) lst)
            (Rest Lst)))))

(define (list->List lst)
  (assume-in 'list->List
    (list? lst))
  (let loop ((lst (reverse lst)) (Lst Nil))
    (if (null? lst)
      Lst
      (loop (cdr lst) (Cons (car lst) Lst)))))

(define (List . args)
  (list->List args))

(define (Take n Lst)
  (assume-in 'Take
    (List? Lst)
    (fixnum? n)
    (fx>= n 0))
  (call-with-values
    (lambda () (Split-at n Lst))
    (lambda (a b) a)))

(define (Drop n Lst)
  (assume-in 'Drop
    (Admissible? n Lst))
  (call-with-values
    (lambda () (Split-at n Lst))
    (lambda (a b) b)))

(define (Take-while ok? Lst)
  (assume-in 'Take-while
    (List-finite? Lst)
    (procedure? ok?))
  (nth-value 0 (Split-with ok? Lst)))

(define (Count-while ok? Lst)
  (assume-in 'Count-while
    (List-finite? Lst)
    (procedure? ok?))
  (nth-value 1 (Split-with ok? Lst)))

(define (Drop-while ok? Lst)
  (assume-in 'Drop-while
    (List-finite? Lst)
    (procedure? ok?))
  (nth-value 2 (Split-with ok? Lst)))

(define (Memp ok? Lst)
  (assume-in 'Memp
    (List-finite? Lst)
    (procedure? ok?))
  (Drop-while (o not ok?) Lst))

(define (Memq var Lst)
  (assume-in 'Memq
    (List-finite? Lst))
  (Memp (cut eq? <> var) Lst))

(define (Memv var Lst)
  (assume-in 'Memv
    (List-finite? Lst))
  (Memp (cut eqv? <> var) Lst))

(define (Member var Lst)
  (assume-in 'Member
    (List-finite? Lst))
  (Memp (cut equal? <> var) Lst))

(define (Equ? =? Lst1 Lst2)
  (assume-in 'Equ?
    (procedure? =?)
    (List? Lst1)
    (List? Lst2))
  (if (eqv? (lazy-list-length Lst1) (lazy-list-length Lst2))
    (if (lazy-list-length Lst1)
      ;; both finite
      (let loop ((Lst1 Lst1) (Lst2 Lst2))
        (cond
          ((Null? Lst1) #t)
          ((=? (First Lst1) (First Lst2))
           (loop (Rest Lst1) (Rest Lst2)))))
      ;; both infinite
      (eq? Lst1 Lst2))
    #f))

(define (Eq? Lst1 Lst2)
  (assume-in 'Eq?
    (List? Lst1)
    (List? Lst2))
  (Equ? eq? Lst1 Lst2))

(define (Eqv? Lst1 Lst2)
  (assume-in 'Eqv?
    (List? Lst1)
    (List? Lst2))
  (Equ? eqv? Lst1 Lst2))

(define (Equal? Lst1 Lst2)
  (assume-in 'Equal?
    (List? Lst1)
    (List? Lst2))
  (Equ? equal? Lst1 Lst2))

(define (Assp ok? al)
  (assume-in 'Assp
    (procedure? ok?)
    ((list-of? pair?) al))
  (let (
    (Lst (Drop-while (lambda (pair) (not (ok? (car pair)))) al))
    )
    (if (Null? Lst) #f (First Lst))))

(define (Assq key al)
  (assume-in 'Assq
    (symbol? key)
    ((list-of? pair?) al))
  (Assp (cut eq? <> key) al))

(define (Assv key al)
  (assume-in 'Assv
    ((list-of? pair?) al))
  (Assp (cut eqv? <> key) al))

(define (Assoc key al)
  (assume-in 'Assoc
    ((list-of? pair?) al))
  (Assp (cut equal? <> key) al))

;; internal
(define (memp ok? lst)
  (let loop ((lst lst))
    (cond
      ((null? lst) #f)
      ((ok? (car lst)) lst)
      (else
        (loop (cdr lst))))))

(define (Map proc . Lsts)
  (assume-in 'Map
    (procedure? proc)
    ((list-of? List?) Lsts))
  (if (null? Lsts)
    Nil
    (let loop ((Lsts Lsts))
      (Lazy (apply Length-min Lsts)
        (if (memp Null? Lsts)
          '()
          (cons (apply proc (map Car Lsts))
                (loop (map Cdr Lsts))))))))

(define (For-each proc . Lsts)
  (assume-in 'For-each
    (procedure? proc)
    (apply Lists-one-finite? Lsts))
  (if (null? Lsts)
    (void)
    (let ((len (apply Length-min Lsts)))
      (do ((Lsts Lsts (map Cdr Lsts)) (k 0 (fx+ k 1)))
        ((fx= k len))
        (apply proc (map Car Lsts))))))

(define (Filter ok? Lst)
  (assume-in 'Filter
    (List? Lst))
  (let loop ((Lst Lst))
    (if (Null? Lst)
      Nil
      (let ((first (First Lst)))
        (if (lazy-list-length Lst)
          ;; compute new length via Cons
          (if (ok? first)
            (Cons first (loop (Rest Lst)))
            (loop (Rest Lst)))
          (Lazy #f
            (if (ok? first)
              (cons first (loop (rest Lst)))
              (loop (rest Lst)))))))))

(define (Zip Lst1 Lst2)
  (assume-in
    (List? Lst1)
    (List? Lst2))
  (if (Null? Lst1)
    Lst2
    (if (and (lazy-list-length Lst1) (lazy-list-length Lst2))
      ;; both finite, compute new length with Cons
      (Cons (First Lst1) (Zip Lst2 (Rest Lst1)))
      ;; new length infinite
      (Lazy #f 
        (cons (First Lst1) (Zip Lst2 (Rest Lst1)))))))
   
(define (input->List port read)
  (let loop ()
    ;(Lazy #f
      (let ((datum (read port)))
        (if (eof-object? datum)
            Nil
            ;'()
            (Cons datum (loop))))));)

(define Repeat
  (case-lambda
    ((x) (Lazy #f (cons x (Repeat x))))
    ((n x)
     (assume-in 'Repeat
       (fixnum? n)
       (fx>= n 0))
     (Take n (Repeat x)))))

(define Repeatedly
  (case-lambda
    ((thunk)
     (assume-in 'Repeatedly
       (procedure? thunk))
     (Lazy #f (cons (thunk) (Repeatedly thunk))))
    ((n thunk)
     (assume-in 'Repeatedly
       (procedure? thunk)
       (fixnum? n)
       (fx>= n 0))
     (Take n (Repeatedly thunk)))))

(define Iterate
  (case-lambda
    ((f x)
     (assume-in 'Iterate
       (procedure? f))
     (Lazy #f (cons x (Iterate f (f x)))))
    ((n f x)
     (assume-in 'Iterate
       (procedure? f)
       (fixnum? n)
       (fx>= n 0))
     (Take n (Iterate f x)))))

(define Cycle
  (case-lambda
    ((Lst)
     (assume-in 'Cycle
       (List? Lst))
     (if (Null? Lst)
       Nil
       (let loop ((tail Lst))
         (Lazy #f
           (if (Null? tail)
             (loop Lst)
             (cons (First tail)
                   (loop (rest tail))))))))
    ((n Lst)
     (assume-in 'Cycle
       (List? Lst)
       (fixnum? n)
       (fx>= n 0))
     (Take n (Cycle Lst)))))

(define Range
  (case-lambda
    ((upto)
     (Iterate (abs upto)
              (if (fx>= upto 0)
                (cut fx+ <> 1)
                (cut fx- <> 1))
              0))
    ((from upto)
     (Iterate (abs (fx- upto from))
              (if (fx>= upto from)
                      (cut fx+ <> 1)
                      (cut fx- <> 1))
              from))
    ((from upto step)
     (Iterate (fx/ (fx+ (fx- step 1) (abs (fx- upto from))) step)
              (if (fx>= upto from)
                      (cut fx+ <> step)
                      (cut fx- <> step))
              from))))

(define (Append2 Lst1 Lst2)
  (if (not (lazy-list-length Lst1))
    Lst1
    (let loop ((Lst Lst1))
      (Lazy (if (lazy-list-length Lst2)
              (+ (lazy-list-length Lst1)
                 (lazy-list-length Lst2))
              #f)
        (if (Null? Lst)
          Lst2
          (cons (First Lst) (loop (rest Lst))))))))

(define (Append . Lsts)
  (assume-in 'Append
    ((list-of? List-finite?) (butlast Lsts))
    (List? (last Lsts)))
  (cond
    ((null? Lsts) Nil)
    ((null? (cdr Lsts)) (car Lsts))
    (else
      (Append2 (car Lsts) (apply Append (cdr Lsts))))))

(define (Reverse Lst)
  (assume-in 'Reverse
    (List-finite? Lst))
  (let loop ((Lst Lst) (reverse Nil))
    (if (Null? Lst)
      reverse
      (Lazy (lazy-list-length Lst)
        (loop (rest Lst)
              (Cons (First Lst) reverse))))))

(define (Reverse* Lst)
  (assume-in 'Reverse*
    (List? Lst))
  (letrec (
    (result 
      (Cons Nil
        (Map Cons
             Lst
             (Lazy (lazy-list-length Lst) result))))
    )
    (Rest result)))

(define (Merge <? Lst1 Lst2)
  (assume-in 'Merge
    (procedure? <?)
    (List-finite? Lst1)
    (List-finite? Lst2))
  (let ((len (+ (lazy-list-length Lst1) (lazy-list-length Lst2))))
    (let loop ((Lst1 Lst1) (Lst2 Lst2))
      (cond
        ((Null? Lst1) Lst2)
        ((Null? Lst2) Lst1)
        ((<? (First Lst1) (First Lst2))
         (Lazy len (cons (First Lst1) (loop (rest Lst1) Lst2))))
        (else 
         (Lazy len
               (cons (First Lst2) (loop Lst1 (rest Lst2)))))))))

(define (Sort <? Lst)
  (assume-in 'Sort
    (procedure? <?)
    (List-finite? Lst))
  (let ((len (lazy-list-length Lst)))
    (if (< len 2)
      Lst
      (let ((halflen (quotient len 2)))
        (Merge <?
          (Sort <? (Take halflen Lst))
          (Sort <? (Drop halflen Lst)))))))

(define (Sorted? <? Lst)
  (assume-in 'Sorted?
    (procedure? <?)
    (List-finite? Lst))
  (let loop ((Lst Lst))
    (cond
      ((Null? Lst) #t)
      ((Null? (Cdr Lst)) #t)
      ((<? (Car Lst) (Car (Cdr Lst)))
       (loop (Cdr Lst)))
      (else #f))))

(define (vector->List vec)
  (assume-in 'vector->List
    (vector? vec))
  (let loop ((res Nil) (n (vector-length vec)))
    (if (zero? n)
      res
      (loop (Cons (vector-ref vec (- n 1)) res) (- n 1)))))

;; see comment to List->list
(define (List->vector Lst)
  (assume-in 'List->vector
    (List-finite? Lst))
  (let ((vec (make-vector (lazy-list-length Lst) #f)))
    (let loop ((k 0) (Lst Lst))
      (cond
        ((Null? Lst)
         vec)
        (else
          (vector-set! vec k (First Lst))
          (loop (+ k 1) (rest Lst)))))))

(define (Split-at n Lst)
  (assume-in 'Split-at
    (Admissible? n Lst))
  (let loop ((n n) (head Nil) (tail Lst))
    (if (or (Null? tail) (zero? n))
      (values (Reverse head) tail)
      (loop (- n 1)
            (Cons (First tail) head)
            (Rest tail)))))

(define (Split-with ok? Lst)
  (assume-in 'Split-with
    (procedure? ok?)
    (List-finite? Lst))
  (let loop ((head Nil) (index 0) (tail Lst))
    (if (or (Null? tail) (not (ok? (First tail))))
    ;(if (or (Null? tail) (ok? (First tail)))
      (values (Reverse head) index tail)
      (loop (Cons (First tail) head)
            (+ index 1)
            (Rest tail)))))

(define (Sieve =? Lst)
  (assume-in 'Sieve
    (procedure? =?)
    (List? Lst))
  (let loop ((Lst Lst))
    (if (Null? Lst)
      Nil
      (let (
        (first (First Lst))
        (tail
          (Filter
            (lambda (x)
              (not (=? x (First Lst))))
            (Rest Lst)))
        )
        (if (lazy-list-length Lst)
          (Cons first (loop tail))
          (Lazy #f
            (cons first (loop tail))))))))

(define (Fold-left op base . Lsts)
  (assume-in 'Fold-left
    (procedure? op)
    (apply Lists-one-finite? Lsts))
  (let loop ((base base)
             (Lsts Lsts)
             (len (apply Length-min Lsts)))
    (if (zero? len)
      base
      (loop (apply op base (map First Lsts))
            (map Rest Lsts)
            (fx- len 1)))))

(define (Fold-right op base . Lsts)
  (assume-in 'Fold-right
    (procedure? op)
    (apply Lists-one-finite? Lsts))
  (let loop ((Lsts Lsts)
             (len (apply Length-min Lsts)))
    (if (zero? len)
      base
      (apply op 
             (append (map First Lsts)
                     (list (loop (map Rest Lsts) (fx- len 1))))))))

;;; The following two routines return Lists
(define (Fold-left* op base . Lsts)
  (assume-in 'Fold-left*
    (procedure? op)
    ((list-of? List?) Lsts))
  (letrec (
    (fold
      (Cons base
        (apply Map op
                   (Lazy (apply Length-min Lsts) fold)
                   Lsts)))
    )
    (Rest fold)))

(define (Fold-right* op base . Lsts) ; changes order of List items
  (assume-in 'Fold-right*
    (procedure? op)
    ((list-of? List?) Lsts))
  (letrec (
    (fold
      (Cons base
        (apply Map op
               (append Lsts
                       (list
                         (Lazy (apply Length-min Lsts) fold))))))
    )
    (Rest fold)))

(define (Every? ok? Lst)
  (assume-in 'Every?
    (procedure? ok?)
    (List-finite? Lst))
  (let loop ((Lst Lst))
    (cond
      ((Null? Lst) #t)
      ((ok? (First Lst))
       (loop (Rest Lst)))
      (else #f))))

(define (Some? ok? Lst)
  (assume-in 'Some?
    (procedure? ok?)
    (List-finite? Lst))
  (let loop ((Lst Lst))
    (cond
      ((Null? Lst) #f)
      ((ok? (First Lst)) #t)
      (else
        (loop (Rest Lst))))))

(define (List-not-null? xpr)
  (and (List? xpr)
       (not (Null? xpr))))

(define (List-finite? xpr)
  (and (List? xpr) (if (Length xpr) #t #f)))

(define (List-infinite? xpr)
  (and (List? xpr) (if (Length xpr) #f #t)))

(define (Lists-one-finite? . Lsts)
  (assume-in 'List-one-finite?
    (not (null? Lsts))
    ((list-of? List?) Lsts))
  (if (apply Length-min Lsts) #t #f))

;;; two examples
(define (Cardinals)
  (let loop ((n 0)) (Lazy #f (cons n (loop (+ n 1))))))

(define (Primes)
  (Sieve (lambda (x y) (zero? (remainder x y)))
         (Drop 2 (Cardinals))))

) ; module lazy-lists

;(use simple-tests)
;(import lazy-lists)