;; Purely Functional Random-Access Lists.

;; Copyright (c) 2009 David Van Horn
;; Licensed under the Academic Free License version 3.0

;; (at dvanhorn (dot ccs neu edu))

;; Implementation based on Okasaki, FPCA '95.
;; Provisions and contracts at bottom.
#lang scheme
(require (planet cce/scheme:4:1/planet)
         (this-package-in private/tree)
         (this-package-in private/fold)
         (this-package-in private/scons))

(define indx-msg "index ~a too large for list: ~a")

;; A [Tree X] is one of
;; - (make-leaf X)
;; - (make-node X [Tree X] [Tree X]),
;; where height of both subtrees is equal,
;; ie. Tree is a complete binary tree.

;; A [RaListof X] is a [SListof [Pair Nat [Tree X]]].

(define-struct (ra:kons s:kons) ()
  #:property prop:custom-write
  ;; [RaListof X] Port Boolean -> Void
  (lambda (ra p write?)
    (let ((print (if write? write display)))
      (let ((curly? (print-pair-curly-braces)))
        (display (if curly? "{" "(") p)
        (let loop ((ls ra))
          (unless (ra:empty? ls)
            (print (ra:first ls) p)
            (unless (ra:empty? (ra:rest ls))
              (display " " p))
            (loop (ra:rest ls))))
        (display (if curly? "}" ")") p))))
    
  #:property prop:sequence
  ;; [RaListof X] -> [Seq X]
  (lambda (ra)
    ;; Incurs logarithmic overhead at sequence construction time,
    ;; but keeps you from having to dispatch at each position.
    (let ((init (s:foldr (lambda (p r) (cons (cdr p) r)) empty ra)))
      (make-do-sequence
       (lambda ()
         (values
          (lambda (x) (tree-val (car x)))
          (lambda (p)
            (let ((tr (car p)))
              (cond [(leaf? tr) (cdr p)]
                    [else
                     (cons (node-left tr)
                           (cons (node-right tr)
                                 (cdr p)))])))
          init cons? void void))))))

;; Consumes n = 2^i-1 and produces 2^(i-1)-1.
;; Nat -> Nat
(define (half n)
  (arithmetic-shift n -1))

;; Nat [Tree X] Nat [X -> X] -> (values X [Tree X])
(define (tree-ref/update s t i f)
  (cond [(zero? i)
         (values (tree-val t)
                 (let ((v* (f (tree-val t))))
                   (cond [(leaf? t) (make-leaf v*)]
                         [else        
                          (make-node v* (node-left t) (node-right t))])))]
        [else
         (let ((s* (half s)))
           (if (<= i s*)
               (let-values ([(v* t*) 
                             (tree-ref/update s* (node-left t) (- i 1) f)])
                 (values v* (make-node (tree-val t) t* (node-right t))))
               (let-values ([(v* t*) 
                             (tree-ref/update s* (node-right t) (- i 1 s*) f)])
                 (values v* (make-node (tree-val t) (node-left t) t*)))))]))

;; Nat [Tree X] Nat -> X
;; Special-cased above to avoid logarathmic amount of cons'ing
;; and any multi-values overhead.  Operates in constant space.
(define (tree-ref s t i)
  (cond [(zero? i) (tree-val t)]
        [else
         (let ((s* (half s)))
           (if (<= i s*)
               (tree-ref s* (node-left t)  (- i 1))
               (tree-ref s* (node-right t) (- i 1 s*))))]))

;; [RaListof X] Nat [X -> X] -> (values X [RaListof X])
(define (ra:list-ref/update ls i f)
  (let loop ((xs ls) (j i))
    (match xs
      [(s:empty) (error 'ra:list-ref/update indx-msg i ls)]
      [(s:cons (cons s t) r)
       (cond [(< j s) 
              (let-values ([(v* t*) (tree-ref/update s t j f)])
                (values v* (make-ra:kons (cons s t*) r)))]
             [else 
              (let-values ([(v* r*) (loop r (- j s))])
                (values v* (make-ra:kons (s:first xs) r*)))])])))

;; [RaListof X] Nat [X -> X] -> [RaListof X]
(define (ra:list-update ls i f)
  (let-values ([(_ r) (ra:list-ref/update ls i f)]) r))

;; [RaListof X] Nat -> X
;; Special-cased above to avoid logarathmic amount of cons'ing
;; and any multi-values overhead.  Operates in constant space.
(define (ra:list-ref ls i)
  (let loop ((xs ls) (j i))
    (match xs
      [(s:empty) (error 'ra:list-ref indx-msg i ls)]
      [(s:cons (cons s t) r)
       (cond [(< j s) (tree-ref s t j)]
             [else (loop r (- j s))])])))

;; [RaListof X] Nat X -> (values X [RaListof X])
(define (ra:list-ref/set ls i v)
  (ra:list-ref/update ls i (lambda (_) v)))
  
;; [RaListof X] Nat X -> [RaListof X]
(define (ra:list-set ls i v)
  (let-values ([(_ l*) (ra:list-ref/set ls i  v)]) l*))

;; X [RaListof X] -> [RaListof X]
(define (ra:cons x ls)
 (match ls 
   [(s:cons (cons s t1) (s:cons (cons s t2) r))
    (make-ra:kons (cons (+ 1 s s) (make-node x t1 t2)) r)]
   [else 
    (make-ra:kons (cons 1 (make-leaf x)) ls)]))      

;; [RaListof X] -> X
(define (ra:first ls)
  (match ls
    [(s:empty) (error 'ra:first "expected non-empty list")]
    [(s:cons (cons s (struct tree (x))) r) x]))

;; [RaListof X] -> [RaListof X]
(define (ra:rest ls)
  (match ls
    [(s:empty) (error 'ra:rest "expected non-empty list")]
    [(s:cons (cons s (struct leaf (x))) r) r]
    [(s:cons (cons s (struct node (x t1 t2))) r)
     (let ((s* (half s)))
       (make-ra:kons (cons s* t1) (make-ra:kons (cons s* t2) r)))]))

;; [RaListof X]
(define ra:empty s:empty)

;; [Any -> Boolean]
(define ra:empty? s:empty?)

;; [X Y -> Y] Y [RaListof X] -> Y
(define ra:foldl/1 (make-foldl ra:empty? ra:first ra:rest))
(define ra:foldr/1 (make-foldr ra:empty? ra:first ra:rest))

;; [X Y ... Z -> Z] Z [RaListof X] [RaListof Y] ... -> Z
(define ra:foldl
  (case-lambda
    [(f a ls) (ra:foldl/1 f a ls)]
    [(f a . lss)
     (check-nary-loop-args 'ra:foldl add1 f lss)
     (let loop ((lss lss) (a a))
       (cond [(ra:empty? (car lss)) a]
             [else
              (loop (map ra:rest lss)
                    (apply f (append (map ra:first lss)
                                     (list a))))]))]))

(define ra:foldr
  (case-lambda
    [(f b ls) (ra:foldr/1 f b ls)]
    [(f b . lss)
     (check-nary-loop-args 'ra:foldr add1 f lss)
     (let recr ((lss lss))
       (cond [(ra:empty? (car lss)) b]
             [else
              (apply f (append (map ra:first lss)
                               (list (recr (map ra:rest lss)))))]))]))

;; [X -> Any] [RaListof X] -> Any
(define ra:andmap/1 (make-andmap ra:empty? ra:first ra:rest))
(define ra:ormap/1  (make-ormap  ra:empty? ra:first ra:rest))

;; [X Y ... -> Any] [RaListof X] [RaListof Y] ... -> Any
(define ra:andmap
  (case-lambda
    [(f ls) (ra:andmap/1 f ls)]
    [(f . lss)
     (check-nary-loop-args 'ra:andmap (lambda (x) x) f lss)
     (cond [(ra:empty? (car lss)) true]
           [else
            (let loop ((lss lss))
              (cond [(ra:empty? (ra:rest (car lss)))
                     (apply f (map ra:first lss))]
                    [else
                     (and (apply f (map ra:first lss))
                          (loop (map ra:rest lss)))]))])]))

(define ra:ormap
  (case-lambda
    [(f ls) (ra:ormap/1 f ls)]
    [(f . lss)
     (check-nary-loop-args 'ra:ormap (lambda (x) x) f lss)
     (cond [(ra:empty? (car lss)) false]
           [else
            (let loop ((lss lss))
              (cond [(ra:empty? (ra:rest (car lss)))
                     (apply f (map ra:first lss))]
                    [else
                     (or (apply f (map ra:first lss))
                         (loop (map ra:rest lss)))]))])]))
     
;; Symbol [Nat -> Nat] Proc [Listof [RaListof X]] -> Void
(define (check-nary-loop-args name mod f lss)
  (let ((n (ra:length (car lss)))
        (m (mod (length lss))))
    (let loop ((l (cdr lss)))
      (unless (empty? l)
        (unless (= n (ra:length (car l)))
          (error name 
                 "given lists of un-equal size: ~a" lss))
        (loop (cdr l))))
    (unless (procedure-arity-includes? f m)
      (error name
             "arity mismatch for ~a, expects ~a arguments, given ~a"
             f (procedure-arity f) m))))


;; [Any -> Boolean]
(define (ra:cons? x)
  (match x
    [(s:cons (cons (? integer?) (? tree?)) r) true]
    [else false]))

;; [Any -> Boolean]
(define (ra:list? x)
  (or (ra:empty? x)
      (ra:cons? x)))
      
;; X ... -> [RaListof X]
(define (ra:list . xs)
  (foldr ra:cons ra:empty xs))

;; X ... [RaListof X] -> [RaListof X]
(define (ra:list* x . r+t)
  (let loop ((xs+t (cons x r+t)))
    (match xs+t
      [(list (? ra:list? t)) t]
      [(list x) (error 'ra:list* "expected list, given: ~a" x)]
      [(cons x xs+t) (ra:cons x (loop xs+t))])))

;; Nat [Nat -> X] -> [RaListof X]
;; Optimized based on skew decomposition.
(define (ra:build-list n f)
  (let loop ((n n) (a ra:empty))
    (cond [(zero? n) a]
          [else 
           (let ((t (largest-skew-binary n)))
             (let ((n* (- n t)))
               (loop n*
                     (make-ra:kons 
                      (cons t (build-tree t (lambda (i) (f (+ i n*)))))
                      a))))])))

;; Simple build-list
#;
(define (ra:build-list i f)
  (let loop ((i (sub1 i)) (a ra:empty))
    (cond [(< i 0) a]
          [else (loop (sub1 i)
                      (ra:cons (f i) a))])))

;; Nat X -> [RaListof X]
(define (ra:make-list n x)
  (let loop ((n n) (a ra:empty))
    (cond [(zero? n) a]
          [else 
           (let ((t (largest-skew-binary n)))
             (let ((n* (- n t)))
               (loop n*
                     (make-ra:kons 
                      (cons t (tr:make-tree t x))
                      a))))])))

;; Simple make-list
#;
(define (ra:make-list n x)
  (ra:build-list n (lambda (i) x)))

;; A Skew is a Nat 2^k-1 with k > 0.

;; Skew -> Skew
(define (skew-succ t) (add1 (arithmetic-shift t 1)))

;; Computes the largest skew binary term t <= n.
;; Nat -> Skew
(define (largest-skew-binary n)
  (cond [(= 1 n) 1]
        [else
         (let* ((t (largest-skew-binary (half n)))
                (s (skew-succ t)))
           (cond [(> s n) t]
                 [else s]))]))

;; [X -> y] [RaListof X] -> [RaListof Y]
;; Takes advantage of the fact that map produces a list of equal size.
(define ra:map
  (case-lambda 
    [(f ls)
     (s:foldr (lambda (p r)
                (make-ra:kons (cons (car p) (tree-map f (cdr p))) r))
              ra:empty
              ls)]
    
    [(f . lss)
     (check-nary-loop-args 'ra:map (lambda (x) x) f lss)
     (let recr ((lss lss)) ; list of s:list of (cons Nat Tree)
       (cond [(s:empty? (car lss)) ra:empty]
             [else
              (make-ra:kons 
               (cons (car (s:first (car lss)))
                     (tree-map/n f (map (compose cdr s:first) lss)))
               (recr (map s:rest lss)))]))]))

;; [RaListof X] -> Nat
(define (ra:length ls)
  (s:foldl (lambda (p len) (+ len (car p))) 0 ls))

;; [RaListof X] Nat -> [RaListof X]
(define (ra:list-tail ls i)
  (let loop ((xs ls) (j i))
    (cond [(zero? j) xs]
          [(ra:empty? xs) (error 'ra:list-tail indx-msg i ls)]
          [else (loop (ra:rest xs) (sub1 j))])))

;; [RaListof X] ... -> [RaListof X]
(define (ra:append . lss)
  (cond [(empty? lss) ra:empty]
        [lss (let recr ((lss lss))
               (cond [(empty? (cdr lss)) (car lss)]
                     [else (ra:foldr ra:cons
                                     (recr (cdr lss))
                                     (car lss))]))]))

;; [RaListof X] -> [RaListof X]
(define (ra:reverse ls)
  (ra:foldl ra:cons ra:empty ls))

(define-sequence-syntax ra:in-list
  (lambda () #'(lambda (x) x))
  (lambda (stx)
    (syntax-case stx ()
      [((id) (_ ra-list-exp))
       #'[(id)
          (:do-in
           ;; outer bindings
           ([(forest)
             (s:foldr (lambda (p r) (cons (cdr p) r)) empty ra-list-exp)])
           'outer-check
           ;; loop bindings
           ([tree (and (cons? forest)
                       (car forest))] 
            [forest (and (cons? forest)
                         (cdr forest))])
           ;; pos check
           (tree? tree)
           ;; inner bindings
           ([(id) (tree-val tree)]
            [(tree forest)
             (cond [(node? tree)
                    (values (node-left tree)
                            (cons (node-right tree) forest))]
                   [(and (leaf? tree) (cons? forest))
                    (values (car forest)
                            (cdr forest))]
                   [else (values false false)])])
                   
           #t ;; pre guard
           #t ;; post guard
           ;; loop args
           (tree forest))]]
      [_ #f])))

(provide (rename-out [ra:in-list in-list]))
(provide/contract 
 (rename ra:cons      cons      (-> any/c ra:list? ra:cons?))
 (rename ra:empty     empty     ra:empty?)
 (rename ra:list-ref  list-ref  (-> ra:cons? natural-number/c any))
 (rename ra:list-set  list-set  (-> ra:cons? natural-number/c any/c ra:cons?))
 (rename ra:cons?     cons?     (-> any/c boolean?))
 (rename ra:empty?    empty?    (-> any/c boolean?))
 (rename ra:list?     list?     (-> any/c boolean?))
 (rename ra:first     first     (-> ra:cons? any))
 (rename ra:rest      rest      (-> ra:cons? ra:list?)) 
 (rename ra:list      list      (->* () () #:rest (listof any/c) ra:list?))
 (rename ra:list*     list*     (->* (any/c) () #:rest (listof any/c) ra:list?))
 (rename ra:length    length    (-> ra:list? natural-number/c))
 (rename ra:append    append    (->* () () #:rest (listof ra:list?) ra:list?))
 (rename ra:reverse   reverse   (-> ra:list? ra:list?))
 (rename ra:list-tail list-tail (-> ra:list? natural-number/c ra:list?))
 (rename ra:make-list make-list (-> natural-number/c any/c ra:list?))
 (rename ra:list-update list-update
         (-> ra:cons? natural-number/c (-> any/c any) ra:cons?))
 (rename ra:list-ref/update list-ref/update
         (-> ra:cons? natural-number/c (-> any/c any) (values any/c ra:cons?)))
 (rename ra:list-ref/set list-ref/set
         (-> ra:cons? natural-number/c any/c (values any/c ra:cons?)))
 (rename ra:build-list build-list 
         (-> natural-number/c (-> natural-number/c any) ra:list?))
 
 (rename ra:map map
         (case-> (-> (-> any/c any) ra:list? ra:list?)
                 (-> procedure? ra:list? ra:list?
                     #:rest (listof ra:list?)
                     ra:list?)))
  
 (rename ra:andmap andmap
         (case-> (-> (-> any/c any) ra:list? any)
                 (-> procedure? ra:list? ra:list?
                     #:rest (listof ra:list?)
                     any)))
  
 (rename ra:ormap ormap
         (case-> (-> (-> any/c any) ra:list? any)
                 (-> procedure? ra:list? ra:list?
                     #:rest (listof ra:list?)
                     any)))
 
  (rename ra:foldr foldr     
         (case-> (-> (-> any/c any/c any) any/c ra:list? any)
                 (-> procedure? any/c ra:list? ra:list?
                     #:rest (listof ra:list?)
                     any)))
 
  (rename ra:foldl foldl 
         (case-> (-> (-> any/c any/c any) any/c ra:list? any)
                 (-> procedure? any/c ra:list? ra:list?
                     #:rest (listof ra:list?)
                     any))))