;; 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
(define-struct tree        (val)        #:prefab)
(define-struct (leaf tree) ()           #:prefab)
(define-struct (node tree) (left right) #:prefab)

;; Sequential Lists
(define-struct s:mt   ()
  #:property prop:equal+hash
  (list 
   (lambda (mt1 mt2 _) (eq? mt1 mt2))
   (lambda (mt _) 1984)
   (lambda (mt _) 4891))
  
  #:property prop:sequence
  (lambda (mt)
    (make-do-sequence
     (lambda ()
       (values (lambda _ 'ignore)
               (lambda _ 'ignore)
               'ignore
               (lambda _ false)
               (lambda _ true)
               (lambda _ true))))))

(define-struct s:kons (first rest)
  #:property prop:equal+hash
  (list 
   (lambda (sc1 sc2 equal?)
     (and (equal? (s:first sc1) (s:first sc2))
          (equal? (s:rest sc1) (s:rest sc2))))
   
   ;; I'm just guessing here...
   (lambda (sc equal-hash-code)
     (+ (bitwise-bit-field (equal-hash-code (s:first sc)) 0 14)
        (arithmetic-shift 
         (bitwise-bit-field (equal-hash-code (s:rest sc)) 0 14) 14)))
   
   (lambda (sc equal-hash-code)
     (+ (bitwise-bit-field (equal-hash-code (s:first sc)) 14 28)
        (arithmetic-shift 
         (bitwise-bit-field (equal-hash-code (s:rest sc)) 14 28) 14))))
  
  #:property prop:sequence 
  (lambda (rac)
    (make-do-sequence
     (lambda ()
       (values ra:first   ;; Traverse s:cons as ra:cons.
               ra:rest
               rac
               ra:cons?
               (lambda _ #t)
               (lambda _ #t))))))

(define the-s:mt (make-s:mt))
(define s:first s:kons-first)
(define s:rest  s:kons-rest)
(define s:cons? s:kons?)
(define s:empty? (lambda (x) (eq? x the-s:mt)))

(define-match-expander s:empty
  (syntax-rules ()
    [(s:empty) (struct s:mt ())])
  the-s:mt)

(define-match-expander s:cons 
  (syntax-rules ()
    [(s:cons x y) (struct s:kons (x y))])
  make-s:kons)

(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]]].

;; Consumes 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 ([(t* v*) 
                             (tree-ref/update s* (node-left t) (- i 1) f)])
                 (values v* (make-node (tree-val t) t* (node-right t))))
               (let-values ([(t* v*) 
                             (tree-ref/update s* (node-right t) (- i 1 s*) f)])
                 (values v* (make-node (tree-val t) (node-left t) t*)))))]))

;; [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* (s:cons (cons s t*) r)))]
             [else 
              (let-values ([(v* r*) (loop r (- j s))])
                (values v* (s:cons (s:first xs) r*)))])])))
  
;; [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 -> X
(define (ra:list-ref ls i)
  (let-values ([(v* _) (ra:list-ref/set ls i '_)]) 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*))

;; Fold combinators
(define (make-foldl empty? first rest)
  (letrec ((f (lambda (cons empty ls)
                (cond [(empty? ls) empty]
                      [else (f cons
                               (cons (first ls) empty) 
                               (rest ls))]))))
    f))

(define (make-foldr empty? first rest)
  (letrec ((f (lambda (cons empty ls)
                (cond [(empty? ls) empty]
                      [else (cons (first ls)
                                  (f cons empty (rest ls)))]))))
    f))

;; X [RaListof X] -> [RaListof X]
(define (ra:cons x ls)
 (match ls 
   [(s:cons (cons s t1) (s:cons (cons s t2) r))
    (s:cons (cons (+ 1 s s) (make-node x t1 t2)) r)]
   [else 
    (s:cons (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)))
       (s:cons (cons s* t1) (s:cons (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 (make-foldl ra:empty? ra:first ra:rest))
(define ra:foldr (make-foldr ra:empty? ra:first ra:rest))

;; [X Y -> Y] Y [SListof X] -> Y
(define s:foldl (make-foldl s:empty? s:first s:rest))
(define s:foldr (make-foldr s:empty? s:first s:rest))

;; [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]
(define (ra:build-list n f)
  (let loop ((i n) (a ra:empty))
    (cond [(zero? i) a]
          [else (loop (sub1 i) 
                      (ra:cons (f (sub1 i)) a))])))

;; [X -> Y] [Tree X] -> [Tree Y]
(define (tree-map f t)
  (cond [(leaf? t) (make-leaf (f (tree-val t)))]
        [(node? t) (make-node (f (tree-val t))
                              (tree-map f (node-left t))
                              (tree-map f (node-right t)))]))

;; [X -> Y] [SListof X] -> [SListof Y]
(define (s:map f ls)
  (s:foldr (lambda (x r) (s:cons (f x) r)) s:empty ls))

;; [X -> y] [RaListof X] -> [RaListof Y]
;; Takes advantage of the fact that map produces a list of equal size.
(define (ra:map f ls)
  (s:map (lambda (p) (cons (car p) (tree-map f (cdr p))))
         ls))

;; [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] -> [RaListof X]
(define (ra:append ls1 ls2)
  (ra:foldr ra:cons ls2 ls1))

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

(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:map       map       (-> (-> any/c any) ra:list? ra:list?))
 (rename ra:foldr     foldr     (-> (-> any/c any/c any) any/c ra:list? any))
 (rename ra:foldl     foldl     (-> (-> any/c any/c any) any/c ra:list? any))
 (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    (-> ra:list? 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: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?)))