3 files changed, 33 insertions(+), 22 deletions(-)

M lib/vector.ml
M lib/vector.mli
M test/vectorTest.ml

@@ 12,6 12,8 @@ let bits_per_level = 5
 let branch_factor = 1 lsl bits_per_level
 (* inner_mask is used to access items within a single level of the tree *)
 let inner_mask = branch_factor - 1
+(* For debugging in-development *)
+let verbose = false
 
 type 'a node = {
     (* Note: This array is always [branch_factor]-sized. *)


@@ 101,7 103,6 @@ let tail_offset (t : 'a t) : int =
     [i] is not in bounds.
  *)
 let array_for
-      ?(verbose : bool = false)
       (t : 'a t) (i : int) : 'a node_item array option =
   if verbose then
     Printf.eprintf "array_for %d:" i;


@@ 122,8 123,8 @@ let array_for
                    Printf.eprintf " -> %d" idx in
         match node.array.(idx) with
         | Inner node -> descend node (height - bits_per_level)
-        | Leaf _
-          | Empty -> failwith "Invariant violation, should only have inner nodes in array_for"
+        | Leaf _ -> failwith "Invariant violation, found Leaf in array_for"
+        | Empty -> failwith "Invariant violation, found Empty in array_for"
     in
     Some (descend t.root t.height)
 


@@ 139,6 140,7 @@ let verify_invariants (t : 'a t) : unit =
   assert (t == empty () || (Array.length t.tail) > 0);
   (* Tail should only have leaves *)
   assert (is_all_leaves t.tail);
+  assert (Int.equal 0 (t.height mod bits_per_level));
   let rec check_level node height =
     if Int.equal 0 height then
       assert (is_all_leaves node.array)


@@ 272,11 274,6 @@ let pop_opt (t : 'a t) : 'a t option =
      let last_index = count - 1 in
      if last_index > (tail_offset t) then
        (* Shrink tail *)
-       let () = Printf.eprintf "last_index %d, tail_offset %d, tail length %d\n"
-                  last_index
-                  (tail_offset t)
-                  (Array.length t.tail)
-       in
        let tail = Array.sub t.tail 0 (Array.length t.tail - 1) in
        Some { t with tail; count = count - 1 }
      else


@@ 284,23 281,25 @@ let pop_opt (t : 'a t) : 'a t option =
           go into tree and pluck out array to be the new tail.
         *)
        let new_last_index = last_index - 1 in
-       let new_tail = array_for ~verbose:true t new_last_index |> Option.get in
+       let new_tail = array_for t new_last_index |> Option.get in
        assert (Int.equal branch_factor (Array.length new_tail)
-                 && is_all_leaves new_tail);
+               && is_all_leaves new_tail);
        (* Remove the new tail from the subtree rooted at [parent]
           Returns the new version of [parent]
           Inverse of [append_node] in [append]. *)
        let rec remove parent height =
          let child_index = (new_last_index lsr height) land inner_mask in
-         if Int.equal bits_per_level height then
+         if height < bits_per_level then
+           failwith "Invariant violation"
+         else if Int.equal bits_per_level height then
            let new_node = { array = Array.copy parent.array } in
            let () = new_node.array.(child_index) <- Empty in
            new_node
          else
            let old_child = match parent.array.(child_index) with
              | Inner node -> node
-             | Leaf _
-               | Empty -> failwith "Invariant violation, only expecting Inner"
+             | Leaf _ -> failwith (Printf.sprintf "Invariant violation at height %d, expecting Inner, got Leaf" height)
+             | Empty -> failwith "Invariant violation, expecting Inner, got Empty"
            in
            let child = remove old_child (height - bits_per_level) in
            let new_node = { array = Array.copy parent.array } in


@@ 316,8 315,10 @@ let pop_opt (t : 'a t) : 'a t option =
        let (root, height) =
          (* TODO should we use bit logic to correspond with append_node? *)
          match root.array.(0), root.array.(1) with
-         (* Tree should shrink *)
-         | Inner inner, Empty -> inner, t.height - bits_per_level
+         (* Tree should shrink. The `when` is so that we keep an empty tree node
+            when dropping below [branch-factor] elements *)
+         | Inner inner, Empty when t.height > bits_per_level ->
+            inner, t.height - bits_per_level
          (* Tree doesn't need to shrink *)
          | _ -> root, t.height
        in

@@ 9,14 9,14 @@ val empty : unit -> 'a t
 (** Gets the element at index, or None if out of bounds *)
 val nth_opt : 'a t -> int -> 'a option
 
-(** Gets the element at index, or raises Invalid_argument if out of bounds *)
+(** Gets the element at index, or raises [Invalid_argument] if out of bounds *)
 val nth : 'a t -> int -> 'a
 
 (** Appends an element to the vector *)
 val append : 'a t -> 'a -> 'a t
 
 (** Returns a new vector with index i associated to a. If i is equal to the current vector size, a is appended to the vector, expanding it.
-    If i is less than zero or greater than the current vector size, Invalid_argument is raised.
+    If i is less than zero or greater than the current vector size, [Invalid_argument] is raised.
  *)
 val update : 'a t -> int -> 'a -> 'a t
 

@@ 10,19 10,29 @@ let test_insert_seq_sanity () =
   let actual = List.of_seq (Vector.to_seq v) in
   assert (List.equal Int.equal actual expected)
 
+let array_equal a b =
+  let rec check idx =
+    if idx >= (Array.length a) then
+      true
+    else
+      (Int.equal a.(idx) b.(idx)) && check (idx + 1)
+  in
+  (Int.equal (Array.length a) (Array.length b)) &&
+    check 0
+
 let test_insert_remove () =
   let num = 1 lsl 14 in
   let v = Seq.init num Fun.id |> Vector.of_seq in
+  let arr = Array.init num Fun.id in
   let f (acc : int Vector.t) (i : int) : int Vector.t =
-    Printf.eprintf "Iteration %d\n" i;
     let next = Vector.pop acc in
-    (* let actual = List.of_seq (Vector.to_seq v) in
-    let expected = List.init (num - i - 1) Fun.id in
-    assert (List.equal Int.equal actual expected); *)
+    let actual = Array.of_seq (Vector.to_seq next) in
+    let expected = Array.sub arr 0 (num - i - 1) in
+    assert (array_equal actual expected);
     next
   in
   let v =
-    List.fold_left f v (List.init num Fun.id)
+    Seq.fold_left f v (Seq.init num Fun.id)
   in
   assert (Int.equal 0 (Vector.length v))