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 iteri f (_, root) = let rec loop idx = function | Leaf -> idx | Node(_, left, x, right) -> let idx = loop idx left in let () = f idx x in loop (idx+1) right in let _ = loop 0 root in () ;; 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 mapi f (len, root) = let rec loop idx = function | Leaf -> idx, Leaf | Node(lw, l, x, r) -> let idx, left = loop idx l in let y = f idx x in let idx, right = loop (idx+1) r in idx, Node(lw, left, y, right) in let _, root = loop 0 root in len, 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 ;; let search f (_, root) = let rec loop = function | Leaf -> raise Not_found | Node(_, l, x, r) -> let rval = f x in if rval < 0 then loop l else if rval > 0 then loop r else x in loop root ;; let search_index f (_, root) = let rec loop idx = function | Leaf -> raise Not_found | Node(lw, l, x, r) -> let rval = f x in if rval < 0 then loop idx l else if rval > 0 then loop (idx + lw + 1) r else (idx + lw) in loop 0 root ;; let append (awght, aroot) (bwght, broot) = let rec loop aw ar bw br = match ar, br with | Leaf, _ -> br | _, Leaf -> ar | Node(alw, al, ax, ar), Node(blw, bl, bx, br) -> if alw >= (bw - blw - 1) then let ar = loop (aw - alw - 1) ar bw br in makenode (aw + bw) alw al ax ar else let bl = loop aw ar blw bl in makenode (aw + bw) (aw + blw) bl bx br in (awght + bwght), (loop awght aroot bwght broot) ;; let sub (size, root) idx len = let rec loop idx len = function | Leaf -> assert false | Node(lw, l, x, r) -> let last = idx+len in let lw1 = lw + 1 in let l' = if idx < lw then if last >= lw then loop idx len l else loop idx (lw - idx) l else (0, Leaf) in let l' = if (idx <= lw) && (last >= lw1) then insert l' 0 x else l' in let r' = if (idx + len) > lw1 then if (idx > lw1) then loop (idx - lw1) (len - lw1) r else loop 0 (len - lw1) r else (0, Leaf) in append l' r' in if (idx < 0) || (len < 1) || (idx >= size) || (len >= size) || ((idx + len) > size) then invalid_arg "IdxTree.sub" else loop idx len root ;; let rev (len, root) = let rec loop tw = function | Leaf -> Leaf | Node(lw, l, x, r) -> let rw = tw - lw - 1 in Node(rw, (loop rw r), x, (loop lw l)) in len, (loop len root) ;; let of_list lst = let rec loop totwt lst = if totwt == 0 then lst, Leaf else let lw = (totwt - 1)/2 in let lst, l = loop lw lst in match lst with | [] -> assert false | h :: t -> let lst, r = loop (totwt - 1 - lw) lst in lst, Node(lw, l, h, r) in let totwt = List.length lst in let _, root = loop totwt lst in (totwt, root) ;; let to_list (_, root) = let rec loop accum = function | Leaf -> accum | Node(_, l, x, r) -> loop (x :: (loop accum r)) l in loop [] root ;; let of_array arr = let rec loop idx len = if len == 0 then Leaf else if len == 1 then Node(0, Leaf, arr.(idx), Leaf) else let lw = (len - 1)/2 in let rw = len - 1 - lw in Node(lw, (loop idx lw), arr.(idx + lw), (loop (idx + lw + 1) rw)) in let len = Array.length arr in len, (loop 0 len) ;; let to_array (len, root) = if (len == 0) then [| |] else let arr = Array.make len ( match root with | Node(_, _, x, _) -> x | Leaf -> assert false ) in let rec loop idx = function | Leaf -> () | Node(lw, l, x, r) -> loop idx l; arr.(idx+lw) <- x; loop (idx+lw+1) r in loop 0 root; arr ;; let enum (len, root) = let rec pushnode t node = match node with | Leaf -> () | Node(_, l, _, _) -> t := (node :: !t); pushnode t l in let popnode r = match !r with | [] -> raise Enum.No_more_elements | h :: t -> r := t; h in let next t c () = match popnode t with | Leaf -> assert false | Node(_, _, x, r) -> pushnode t r; decr c; x and count c () = !c in let rec clone t c () = let t = ref (!t) and c = ref (!c) in Enum.make ~next:(next t c) ~count:(count c) ~clone:(clone t c) in let t = ref [] and c = ref len in pushnode t root; Enum.make ~next:(next t c) ~count:(count c) ~clone:(clone t c) ;; let of_enum e = let rec loop cnt = let lw = (cnt - 1)/2 in let rw = cnt - 1 - lw in let l = loop lw in let x = match Enum.get e with | None -> assert false | Some x -> x in let r = loop rw in Node(lw, l, x, r) in let len = Enum.count e in len, (loop len) ;;