;; Purely Functional Random-Access Lists.

;; Copyright (c) 2007, 2008 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)

;; 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 [Listof [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 -> (values X [Tree X])
(define (tree-ref/set s t i v)
  (cond [(zero? i)
         (values (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/set s* (node-left t) (- i 1) v)])
                 (values v* (make-node (tree-val t) t* (node-right t))))
               (let-values ([(t* v*) (tree-ref/set s* (node-right t) (- i 1 s*) v)])
                 (values v* (make-node (tree-val t) (node-left t) t*)))))]))

;; [RaListof X] Nat X -> (values X [RaListof X])
(define (ra:list-ref/set ls i v)
  (match ls
    [(list) (error (object-name ra:list-ref/set) "index out of bounds: ~a" i)]
    [(cons (cons s t) r)
     (cond [(< i s) 
            (let-values ([(v* t*) (tree-ref/set s t i v)])
              (values v* (cons (cons s t*) r)))]
           [else 
            (let-values ([(v* r*) (ra:list-ref/set r (- i s) v)])
              (values v* (cons (first ls) r*)))])]))

;; [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*))

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

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

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

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

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

;; [Any -> Boolean]
(define (ra:cons? x)
  (match x
    [(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))

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

;; [X Y -> Y] Y [RaListof X] -> Y
(define (ra:foldl f b ls)
  (cond [(ra:empty? ls) b]
        [else (ra:foldl f (f (ra:first ls) b) (ra:rest ls))]))

;; [X Y -> Y] Y [RaListof X] -> Y
(define (ra:foldr f b ls)
  (cond [(ra:empty? ls) b]
        [else (f (ra:first ls) (ra:foldr f b (ra:rest 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-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?)))