type 'a node_t = | Node of int * 'a node_t * 'a * 'a node_t | Leaf ;; type 'a t = int * 'a node_t;; let empty = 0, Leaf;; let length (l, _) = l;; let get (len, root) idx = if (idx < 0) || (idx >= len) then invalid_arg "index out of bounds" else let rec loop idx = function | Leaf -> assert false | Node(leftweight, left, elem, right) -> if idx < leftweight then loop idx left else if idx == leftweight then elem else loop (idx - leftweight) right in loop idx root ;; let set (len, root) idx x = if (idx < 0) || (idx >= len) then invalid_arg "index out of bounds" else let rec loop idx = function | Leaf -> assert false | Node(leftweight, left, elem, right) -> if idx < leftweight then Node(leftweight, (loop idx left), elem, right) else if idx == leftweight then Node(leftweight, left, x, right) else Node(leftweight, left, elem, (loop (idx - leftweight) right)) in len, (loop idx root) ;; let makenode totalweight leftweight left elem right = let rightweight = totalweight - 1 - leftweight in if (leftweight * 2) <= rightweight then (* rotate left *) match right with | Leaf -> assert false | Node(rw, rl, re, rr) -> Node(leftweight + 1 + rw, Node(leftweight, left, elem, rl), re, rr) else if (rightweight * 2) <= leftweight then (* rotate right *) match left with | Leaf -> assert false | Node(lw, ll, le, lr) -> Node(lw, ll, le, Node(leftweight - 1 - lw, lr, elem, right)) else Node(leftweight, left, elem, right) ;; let rec join totalweight leftweight left right = match right with | Leaf -> left | Node(lw, l, e, r) -> makenode totalweight (leftweight + lw) (join (leftweight + lw) lw left l) e r ;; let insert (len, root) idx elem = if (idx < 0) || (idx > len) then invalid_arg "index out of bounds" else let rec loop totalweight idx = function | Leaf -> Node(0, Leaf, elem, Leaf) | Node(lw, left, x, right) -> if idx <= lw then makenode (totalweight + 1) (lw + 1) (loop lw idx left) x right else makenode (totalweight + 1) lw left x (loop (totalweight - 1 - lw) (idx - lw) right) in (len + 1), (loop len idx root) ;; let remove (len, root) idx = if (idx < 0) || (idx >= len) then invalid_arg "index out of bounds" else let rec loop totalweight idx = function | Leaf -> assert false | Node(lw, left, x, right) -> if idx < lw then makenode (totalweight - 1) (lw - 1) (loop lw idx left) x right else if idx == lw then join (totalweight - 1) lw left right else makenode (totalweight - 1) lw left x (loop (totalweight - 1 - lw) (idx - lw) right) in (len - 1), (loop len idx root) ;; let make len x = if (len < 0) then invalid_arg "Funarr.make" else let rec loop base len = if (len == 0) then Leaf else let lw = (len - 1)/2 in Node(lw, (loop base lw), x, (loop (base + lw + 1) (len - lw - 1))) in len, (loop 0 len) ;; let init len f = if (len < 0) then invalid_arg "Funarr.make" else let rec loop base len = if (len == 0) then Leaf else let lw = (len - 1)/2 in let left = loop base lw in let x = f (base + lw) in let right = loop (base + lw + 1) (len - lw - 1) in Node(lw, left, x, right) in len, (loop 0 len) ;; let iter f (_, root) = let rec loop = function | Leaf -> () | Node(_, left, x, right) -> let () = loop left in let () = f x in loop right in loop root ;; let map f (len, root) = let rec loop = function | Leaf -> Leaf | Node(lw, l, x, r) -> let left = loop l in let y = f x in let right = loop r in Node(lw, left, y, right) in len, (loop root) ;; let fold_left f init (_, root) = let rec loop accu = function | Leaf -> accu | Node(_, l, x, r) -> let accu = loop accu l in let accu = f accu x in loop accu r in loop init root ;; let fold_right f (_, root) init = let rec loop accu = function | Leaf -> accu | Node(_, l, x, r) -> let accu = loop accu r in let accu = f x accu in loop accu l in loop init root ;;