(module yasos-collections

  (collection? random-access? empty? size gen-keys gen-elts do-elts
   do-keys do-items map-elts map-keys map-items for-each-key for-each-elt
   elt-ref elt-set! elt-take elt-drop elt-slice reduce reduce* reduce-items
   reduce-items* sort! sort make-vector-generator list->generator
   vector->generator string->generator hash-table->generator
   generator->list g-map g-reduce g-find g-filter g-zip
   lseq->list)
  
 (import scheme (chicken base) (chicken format) srfi-69 srfi-127
         (except yasos object object-with-ancestors))

 
;; COLLECTION INTERFACE

;; (collection? obj)  -- predicate
;;
;; (empty? collection)  -- I bet you can guess what these do as well...
;; (size collection)
;;
;; (do-elts proc coll+) -- apply proc element-wise to collections
;; (do-keys proc coll+) -- .. return value is unspecified
;;
;; (map-elts proc coll+) -- as with do-*, but returns collection
;; (map-keys proc coll+) -- e.g. (map-keys + (list 1 2 3) (vector 1 2 3))
;;					-> #( 2 4 6 )
;;
;; (for-each-key coll proc) -- for single collection (more efficient)
;; (for-each-elt coll proc)
;;
;; (reduce proc seed coll+) -- e.g. (reduce + 0 (vector 1 2 3))
;;
;;==============================
;; Collections must implement:
;;  collection?
;;  gen-elts
;;  gen-keys
;;  size
;;  print
;;
;; Collections should implement {typically faster}:
;;  for-each-key
;;  for-each-elt
;;
;; Collections may optionally implement:
;; 1) random access operations:
;; elt-ref
;; elt-set!
;; 2) selector operations:
;; elt-take
;; elt-drop
;; elt-slice

 
 (define *eof-object* (read (open-input-string "")))
 (define (eof-object) *eof-object*)

 (define (list-any pred lis)
   (and (not (null? lis))
        (let lp ((head (car lis)) (tail (cdr lis)))
          (if (null? tail)
              (pred head)		; Last PRED app is tail call.
              (or (pred head) (lp (car tail) (cdr tail)))))))

 (define (list-zip list1 . more-lists) (apply map list list1 more-lists))

 (define (list-take lis k)
  (let recur ((lis lis) (k k))
    (if (eq? 0 k) '()
	(cons (car lis)
	      (recur (cdr lis) (- k 1))))))

 (define (list-drop lis k)
   (let iter ((lis lis) (k k))
    (if (eq? 0 k) lis (iter (cdr lis) (- k 1)))))
 
 
 (define-operation (collection? obj)
  ;; default
   (cond
    ((or (list? obj) (vector? obj) (string? obj) (hash-table? obj)) #t)
    (else #f)
    ))

 
 (define-operation (random-access? obj)
  ;; default
   (cond
    ((or (list? obj) (vector? obj) (string? obj) (hash-table? obj)) #t)
    (else #f)
    ))

 
 (define (empty? collection) (zero? (size collection)))
 
 (define-operation (elt-ref <collection> i);; random access collection
  ;; default behavior
   (cond                     
    ((vector? <collection>) (vector-ref <collection> i)) 
    ((list?   <collection>) (list-ref  <collection> i))
    ((string? <collection>) (string-ref  <collection> i))
    ((hash-table? <collection>) (hash-table-ref <collection> i))
    (else 
     (error "operation not supported: elt-ref"))
    ))
 
 (define-operation (elt-set! <collection> i v);; random access collection
  ;; default behavior
   (cond                     
    ((vector? <collection>) (vector-set! <collection> i v)) 
    ((list?   <collection>) (list-set!  <collection> i v))
    ((string? <collection>) (string-set!  <collection> i v))
    ((hash-table? <collection>) (hash-table-set! <collection> i v))
    (else 
     (error "operation not supported: elt-set!"))
    ))

 (define-operation (elt-take <collection> n);; random access collection
  ;; default behavior
   (cond                     
    ((vector? <collection>) (subvector <collection> 0 n))
    ((list?   <collection>) (list-take  <collection> n))
    ((string? <collection>) (substring  <collection> 0 n))
    ((hash-table? <collection>)
     (let ((keys (hash-table-keys <collection>))
           (result (make-hash-table)))
       (for-each
        (lambda (k) (hash-table-set! result k (hash-table-ref <collection> k)))
        (list-take keys n))
       result))
    (else 
     (error "operation not supported: elt-take"))
    ))


 (define-operation (elt-drop <collection> n);; random access collection
  ;; default behavior
   (cond                     
    ((vector? <collection>) (subvector <collection> n))
    ((list?   <collection>) (list-drop  <collection> n))
    ((string? <collection>) (substring  <collection> n))
    ((hash-table? <collection>)
     (let ((keys (hash-table-keys <collection>))
           (result (make-hash-table)))
       (for-each
        (lambda (k) (hash-table-set! result k (hash-table-ref <collection> k)))
        (list-drop keys n))
       result))
    (else 
     (error "operation not supported: elt-take"))
    ))
 
 (define-operation (elt-slice <collection> m n);; random access collection
  ;; default behavior
   (cond                     
    ((vector? <collection>) (subvector <collection> m n))
    ((list?   <collection>) (list-take (list-drop  <collection> m) (- n m)))
    ((string? <collection>) (substring  <collection> m n))
    ((hash-table? <collection>)
     (let ((keys (hash-table-keys <collection>))
           (result (make-hash-table)))
       (for-each
        (lambda (k) (hash-table-set! result k (hash-table-ref <collection> k)))
        (list-take (list-drop keys m) (- n m)))
       result))
    (else 
     (error "operation not supported: elt-take"))
    ))

 (define-operation (gen-elts <collection>);; return SRFI-121 element generator
  ;; default behavior
   (cond                     
    ((vector? <collection>) (vector->generator <collection>)) 
    ((list?   <collection>) (list->generator   <collection>))
    ((string? <collection>) (string->generator <collection>))
    ((hash-table? <collection>) (hash-table->generator <collection>))
    (else 
     (error "operation not supported: gen-elts "))
    ))


 (define-operation (gen-keys collection)
   (cond
    ((or (vector? collection) (list? collection) (string? collection))
     (let ( (max+1 (size collection)) (index (make-parameter 0) ))
       (lambda ()
         (let ((i (index)))
           (cond
            ((< i max+1)
             (index (add1 i))
             i)
            (else (eof-object))
            ))
         ))
     )
    ((hash-table? collection)
     (list->generator (hash-table-keys collection)))
    (else
     (error "operation not handled: gen-keys " collection))
    ))

 (define (do-elts <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (generators (map gen-elts <collections>))
          )
     (let loop ( (counter 0) )
       (cond
        ((< counter max+1)
         (apply <proc> (map (lambda (g) (g)) generators))
         (loop (add1 counter))
         )
        (else 'unspecific)  ; done
        )  )
     ) )

 (define (do-keys <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (generators (map gen-keys <collections>))
          )
     (let loop ( (counter 0) )
       (cond
        ((< counter max+1)
         (apply <proc> (map (lambda (g) (g)) generators))
         (loop (add1 counter))
         )
        (else 'unspecific)  ; done
        )  )
     ) )

 (define (do-items <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (elt-generators (map gen-elts <collections>))
          (key-generators (map gen-keys <collections>)) )
     (let loop ( (counter 0) )
       (cond
        ((< counter max+1)
         (apply <proc> (list-zip (map (lambda (g) (g)) key-generators)
                                 (map (lambda (g) (g)) elt-generators)))
         (loop (add1 counter))
         )
        (else 'unspecific)  ; done
        )  )
     ) )

 (define (map-elts <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (generators (map gen-elts <collections>))
          (vec (make-vector (size (car <collections>))))
          )
     (let loop ( (index 0) )
       (cond
        ((< index max+1)
         (vector-set! vec index (apply <proc> (map (lambda (g) (g)) generators)))
         (loop (add1 index))
         )
        (else vec)  ; done
        )  )
     ) )

 (define (map-keys <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (generators (map gen-keys <collections>))
          (vec (make-vector (size (car <collections>))))
          )
     (let loop ( (index 0) )
       (cond
        ((< index max+1)
         (vector-set! vec index (apply <proc> (map (lambda (g) (g)) generators)))
         (loop (add1 index))
         )
        (else vec)  ; done
        )  )
     ) )

 (define (map-items <proc> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (key-generators (map gen-keys <collections>))
          (elt-generators (map gen-elts <collections>))
          (vec (make-vector (size (car <collections>))))
          )
     (let loop ( (index 0) )
       (cond
        ((< index max+1)
         (vector-set! vec index (apply <proc> (list-zip (map (lambda (g) (g)) key-generators)
                                                        (map (lambda (g) (g)) elt-generators))))
         (loop (add1 index))
         )
        (else vec)  ; done
        )  )
     ) )

 (define-operation (for-each-key <proc> <collection>)
  ;; default
   (do-keys <proc> <collection>) ;; talk about lazy!
   )

 (define-operation (for-each-elt <proc> <collection>)
   (do-elts <proc> <collection>)
   )

 (define (reduce <proc> <seed> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (elt-generators (map gen-elts <collections>))
          (ax (make-parameter <seed>))
          )
     (let loop ( (count 0) )
       (cond
        ((< count max+1)
         (ax (apply <proc> (append (map (lambda (g) (g)) elt-generators) (list (ax)))))
         (loop (add1 count))
         )
        (else (ax))
        ) )
     )  )
 

 (define (reduce-items <proc> <seed> . <collections>)
   (let ( (max+1 (size (car <collections>)))
          (key-generators (map gen-keys <collections>))
          (elt-generators (map gen-elts <collections>))
          (ax (make-parameter <seed>))
          )
     (let loop ( (count 0) )
       (cond
        ((< count max+1)
         (ax (apply <proc> (append (list-zip (map (lambda (g) (g)) key-generators)
                                             (map (lambda (g) (g)) elt-generators))
                                   (list (ax)))))
         (loop (add1 count))
         )
        (else (ax))
        ) )
     )  )
 
;; reduce operation where the first element of the collection is the seed

 (define (reduce* <proc> . <collections>)
   (let* ( (max+1 (- (size (car <collections>)) 1))
           (elt-generators (map gen-elts <collections>))
           (ax (make-parameter (map (lambda (g) (g)) elt-generators)))
           )
     (let loop ( (count 0) )
       (cond
        ((< count max+1)
         (let ((args (append (map (lambda (g) (g)) elt-generators) (ax))))
           (ax (list (apply <proc> args))))
         (loop (add1 count))
         )
        (else (car (ax)))
        ) )
     )  )

 (define (reduce-items* <proc> . <collections>)
   (let* ( (max+1 (- (size (car <collections>)) 1))
           (key-generators (map gen-keys <collections>))
           (elt-generators (map gen-elts <collections>))
           (ax (make-parameter (list-zip (map (lambda (g) (g)) key-generators)
                                         (map (lambda (g) (g)) elt-generators))))
           )
     (let loop ( (count 0) )
       (cond
        ((< count max+1)
         (ax (list (apply <proc> (append (list-zip (map (lambda (g) (g)) key-generators)
                                                   (map (lambda (g) (g)) elt-generators))
                                         (ax)))))
         (loop (add1 count))
         )
        (else (car (ax)))
        ) )
     )  )

 
;; generator for list elements
 (define (list->generator <list>)
   (let ((l (make-parameter <list>)))
     (lambda ()
       (if (null? (l))
           (eof-object)
           (let ( (elt (car (l))) )
             (l (cdr (l)))
             elt))
       ))
   )

 (define (make-vector-generator <accessor>)
   (lambda (vec)
     (let ( (max+1 (size vec))
            (index (make-parameter 0))
            )
       (lambda ()
         (let ((i (index)))
           (cond ((< i max+1)
                  (index (add1 i))
                  (<accessor> vec i)
                  )
                 (else (eof-object))
                 ))
         ))
     ))

 (define vector->generator (make-vector-generator vector-ref))

 (define string->generator (make-vector-generator string-ref))

 (define (hash-table->generator table)
   (let ((keys (make-parameter (hash-table-keys table))))
     (lambda () 
       (cond ((null? keys) (eof-object))
             (else (let ((res (hash-table-ref table (car (keys)))))
                     (keys (cdr (keys)))
                     res))
             ))
     ))


;; nota bene:  list-set! is bogus for element 0

 (define (list-set! <list> <index> <value>)

   (define (set-loop last this idx)
     (cond
      ((zero? idx) 
       (set-cdr! last (cons <value> (cdr this)))
       <list>
       )
      (else (set-loop (cdr last) (cdr this) (sub1 idx)))
      )  )

  ;; main
   (if (zero? <index>)
       (cons <value> (cdr <list>)) ;; return value
       (set-loop <list> (cdr <list>) (sub1 <index>)))
   )

 
 ;;
 ;; In-place quick sort from SRFI-32 reference implementation.
 ;; Modified so that the comparison function uses element indices as
 ;; well as element values:
 ;;
 ;; elt< :: i1 v1 i2 v2 -> boolean
 ;;
 ;; Copyright (c) 1998 by Olin Shivers. You may do as you please with
 ;; this code, as long as you do not delete this notice or hold me
 ;; responsible for any outcome related to its use.
 ;;
 
 (define (sort! elt< v . rest)
   (let-optionals rest ((start 0) (end (size v)))
    (let recur ((l start) (r end))	; Sort the range [l,r).
      (if (< 1 (- r l))
	  
	  ;; Choose the median of V[l], V[r], and V[middle] for the pivot.
	  (let ((median
		 (lambda (i1 i2 i3)
		   (let ((v1 (elt-ref v i1))
			 (v2 (elt-ref v i2))
			 (v3 (elt-ref v i3)))
		     (receive (ilittle little ibig big)
			      (if (elt< i1 v1 i2 v2)
                                  (values i1 v1 i2 v2)
                                  (values i2 v2 i1 v1))
			      (if (elt< ibig big i3 v3) 
				  (values ibig big)
				  (if (elt< ilittle little i3 v3) 
				      (values i3 v3) 
				      (values ilittle little))))))))
	    
	    (let-values (((ipivot pivot) (median l (quotient (+ l r) 2) (- r 1))))
   	      (let loop ((i l) (j (- r 1)))
		(let ((i (let scan ((i i)) (if (elt< i (elt-ref v i) ipivot pivot)
					       (scan (+ i 1))
					       i)))
		      (j (let scan ((j j)) (if (elt< ipivot pivot j (elt-ref v j))
					       (scan (- j 1))
					       j))))
		  (if (< i j)
		      (let ((tmp (elt-ref v j)))		
			(elt-set! v j (elt-ref v i))	; Swap V[I]
			(elt-set! v i tmp)		;  and V[J].
			(loop (+ i 1) (- j 1)))
		      
		      (begin (recur l i) (recur (+ j 1) r)))))))
	  v)
      ))
   )

 ;; Blit FROM[I,END) to TO[J,?].
 
 (define (vector-blit! from i end to j)
    (assert (< i end))
    (let recur ((i i) (j j))
      (if (< i end)
          (let ((vi (elt-ref from i)))
            (vector-set! to j vi)
            (recur (+ i 1) (+ j 1)))
          ))
      )

 
 ;; Given array A and indices p, q, r such that p < q < r, 
 ;; merge subarray A[p..q) and subarray A[q..r) into array B[n..]

(define (vector-merge! elt< a p q r b n)
    (assert (and (< p q) (< q r)))
    (let recur ((i p) (j q) (k n))
      (if (and (< i q) (< j r))
          (let ((ai (elt-ref a i))
                (aj (elt-ref a j)))
            (if (elt< i ai j aj)
                (begin
                  (vector-set! b k ai)
                  (recur (+ 1 i) j (+ 1 k)))
                (begin
                  (vector-set! b k aj)
                  (recur i (+ 1 j) (+ 1 k)))
                ))
          (if (< i q)
              (vector-blit! a i q b k)
              (if (< j r)
                  (vector-blit! a j r b k))))
      )
    b)

 
 ;; Collection merge sort
 (define (sort elt< x)
   (let* ((n (size x))
          (a (make-vector n)))
     (do-items (lambda (item) (vector-set! a (car item) (cadr item))) x)
     (if (< n 2)
         a
         (let ((b (make-vector n)))
           (let recur ((m 1))
             (if (< m n)
                 (let inner-recur ((p 0))
                   (if (< p (- n m))
                       (let ((q (+ p m))
                             (r (min (+ p (* 2 m)) n)))
                         (vector-merge! elt< a p q r b p)
                         (vector-blit! b p r a p)
                         (inner-recur (+ p (* 2 m)))
                         )
                       (recur (* m 2))))
                 b))
           ))
     ))
            

  ;; Generator combinators
  (define (g-map f . gs)
    (lambda ()
      (let ((vs (map (lambda (g) (g)) gs)))
        (if (list-any eof-object? vs)
            (eof-object)
            (apply f vs))
        ))
    )


  (define (g-reduce f seed . gs)
    (define (inner-fold seed)
      (let ((vs (map (lambda (g) (g)) gs)))
        (if (list-any eof-object? vs)
            seed
            (inner-fold (apply f (append vs (list seed)))))))
    (inner-fold seed))


  (define (g-find pred g)
    (let loop ((v (g)))
      (if (or (pred v) (eof-object? v))
          v
          (loop (g)))
      ))

  (define (g-filter pred g)
    (let loop ((v (g)) (res '()))
      (cond ((eof-object? v) res)
            ((pred v) (loop (g) (cons v res)))
            (else (loop (g) res)))
      ))

  (define (g-zip gen . gs)
    (lambda ()
      (let ((value (gen)))
        (if (eof-object? value)
            (eof-object)
            (cons value (map (lambda (g) (g)) gs))
            ))
      ))

  (define (generator->list gen)
    (let recur ((ax '()))
      (let ((value (gen)))
        (if (eof-object? value)
            (reverse ax)
            (recur (cons value ax))
            ))
      ))


  (define (lseq->list lseq)
    (let recur ((lseq lseq) (ax '()))
      (if (null? lseq)
          (reverse ax)
          (recur (lseq-rest lseq) (cons (lseq-first lseq) ax)))
      ))


 )