2 files changed, 54 insertions(+), 28 deletions(-)

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

@@ 45,9 45,8 @@ type 'a t = {
     (* The "tail" is an optimization to the vector that allows true-constant time
        access to the end of the vector, in recognition of the fact that vectors
        are often pushed and popped at the end. The tail array is not stored in the tree.
-       
-       Typed as 'a node_item to match node type, but is logically 'a option array.
-       i.e., this should never hold Inner, only Leaf or Empty
+       Typed as 'a node_item to match node type, but is logically 'a array.
+       i.e., this should never hold Inner or Empty, only Leaf
        
        This array will grow and shrink as the vector is updated, unlike the array
        inside 'a node, but it will always be nonempty (other than for the empty vector)


@@ 128,6 127,31 @@ let array_for
     in
     Some (descend t.root t.height)
 
+let is_leaf = function
+  | Leaf _ -> true
+  | _ -> false
+
+let is_all_leaves array =
+  Array.for_all is_leaf array
+
+let verify_invariants (t : 'a t) : unit =
+  (* Except for the empty vector, tail should always be nonempty *)
+  assert (t == empty () || (Array.length t.tail) > 0);
+  (* Tail should only have leaves *)
+  assert (is_all_leaves t.tail);
+  let rec check_level node height =
+    if Int.equal 0 height then
+      assert (is_all_leaves node.array)
+    else
+      Array.iter
+        (function
+         | Inner child -> check_level child (height - bits_per_level)
+         | Empty -> ()
+         | Leaf _ -> assert false)
+        node.array
+  in
+  check_level t.root t.height
+
 let nth_opt (t : 'a t) (index : int) : 'a option =
   let last_level_index = index land inner_mask in
   match array_for t index with


@@ 153,6 177,7 @@ let append (t : 'a t) (value : 'a) : 'a t =
     Array.blit t.tail 0 new_tail 0 old_tail_length;
     { t with tail = new_tail; count = old_count + 1 }
   else
+    let () = assert (Int.equal branch_factor (Array.length t.tail)) in
     (* No more room in tail, insert the tail into the tree and make a new tail *)
     (* Creates a path of Inner nodes [height] deep and places [node] there, returning
        the parent of the new subtree. *)


@@ 196,12 221,15 @@ let append (t : 'a t) (value : 'a) : 'a t =
         new_root.array.(1) <- Inner (create_subtree t.height node_to_insert);
         new_root, t.height + bits_per_level
     in
-    {
+    let result = {
       root;
       height;
       count = old_count + 1;
       tail = [|(Leaf value)|];
-    }
+      }
+    in
+    let () = verify_invariants result in
+    result
 
 let update (t : 'a t) (i : int) (value : 'a) : 'a t =
   let count = length t in


@@ 232,7 260,9 @@ let update (t : 'a t) (i : int) (value : 'a) : 'a t =
         let () = new_node.array.(child_index) <- Inner child in
         new_node
     in
-    { t with root = rebuild t.root t.height }
+    let result = { t with root = rebuild t.root t.height } in
+    let () = verify_invariants result in
+    result
 
 let pop_opt (t : 'a t) : 'a t option =
   match t.count with


@@ 251,21 281,12 @@ let pop_opt (t : 'a t) : 'a t option =
        Some { t with tail; count = count - 1 }
      else
        (* Popped last element of tail,
-          go into tree and pluck out an array to be the new tail.
+          go into tree and pluck out array to be the new tail.
         *)
        let new_last_index = last_index - 1 in
-       let array = array_for ~verbose:true t new_last_index |> Option.get in
-       (* Resize the array to be snug *)
-       let new_tail_len =
-         let rec f index =
-           match array.(index) with
-           | Leaf _ -> index + 1
-           | Empty -> f (index + 1)
-           | Inner _ -> failwith "Invariant violation, not expecting Inner in last-level array"
-         in
-         f 0
-       in
-       let new_tail = Array.sub array 0 new_tail_len in
+       let new_tail = array_for ~verbose:true t new_last_index |> Option.get in
+       assert (Int.equal branch_factor (Array.length 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]. *)


@@ 278,7 299,8 @@ let pop_opt (t : 'a t) : 'a t option =
          else
            let old_child = match parent.array.(child_index) with
              | Inner node -> node
-             | Leaf _ | Empty -> failwith "Invariant violation, only expecting Inner"
+             | Leaf _
+               | Empty -> failwith "Invariant violation, only expecting Inner"
            in
            let child = remove old_child (height - bits_per_level) in
            let new_node = { array = Array.copy parent.array } in


@@ 290,14 312,18 @@ let pop_opt (t : 'a t) : 'a t option =
            else
              new_node
        in
-       let new_root = remove t.root t.height in
-       (* Handle the case where the root is now redundant *)
-       let new_root =
-         match new_root.array.(0), new_root.array.(1) with
-         | Inner inner, Empty -> inner
-         | _ -> new_root
+       let root = remove t.root t.height in
+       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 doesn't need to shrink *)
+         | _ -> root, t.height
        in
-       Some { t with root = new_root; tail = new_tail; count = count - 1; }
+       let result = { root; height; tail = new_tail; count = count - 1; } in
+       let () = verify_invariants result in
+       Some result
 
 let pop (t : 'a t) : 'a t =
   match pop_opt t with

@@ 4,7 4,7 @@ module Vector = Persistent.Vector
    Should cover pretty much all of the insertion logic and also of_seq/to_seq
  *)
 let test_insert_seq_sanity () =
-  let num = 1 lsl 20 in
+  let num = 1 lsl 14 in
   let v = Seq.init num Fun.id |> Vector.of_seq in
   let expected = List.init num Fun.id in
   let actual = List.of_seq (Vector.to_seq v) in