From 41d8a7a968d5dfe3b371a101245fedc916ac154d Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 16 Jan 2024 14:20:09 -0500
Subject: [PATCH] add `Pvec.equal`

---
 src/pvec/containers_pvec.ml  | 20 +++++++++++++++++++
 src/pvec/containers_pvec.mli |  1 +
 tests/pvec/t_pvec.ml         | 38 +++++++++++++++++++++++++++++++++++-
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/pvec/containers_pvec.ml b/src/pvec/containers_pvec.ml
index d622583f..8a9b9c4b 100644
--- a/src/pvec/containers_pvec.ml
+++ b/src/pvec/containers_pvec.ml
@@ -20,6 +20,16 @@ module A = struct
   let[@inline] return self = [| self |]
   let[@inline] is_full self = length self = branching_factor
 
+  let equal eq a b =
+    length a = length b
+    &&
+    try
+      for i = 0 to length a - 1 do
+        if not (eq (unsafe_get a i) (unsafe_get b i)) then raise_notrace Exit
+      done;
+      true
+    with Exit -> false
+
   let[@inline] push (self : _ t) x =
     let n = length self in
     if n = branching_factor then invalid_arg "Pvec.push";
@@ -339,6 +349,16 @@ let append a b =
   else
     fold_left push a b
 
+let rec equal_tree eq t1 t2 =
+  match t1, t2 with
+  | Empty, Empty -> true
+  | Node a, Node b -> A.equal (equal_tree eq) a b
+  | Leaf a, Leaf b -> A.equal eq a b
+  | (Empty | Leaf _ | Node _), _ -> false
+
+let equal eq (a : _ t) (b : _ t) : bool =
+  a.size = b.size && A.equal eq a.tail b.tail && equal_tree eq a.t b.t
+
 let add_list v l = List.fold_left push v l
 let of_list l = add_list empty l
 let to_list m = fold_rev (fun acc x -> x :: acc) [] m
diff --git a/src/pvec/containers_pvec.mli b/src/pvec/containers_pvec.mli
index 6c40b497..68773a2d 100644
--- a/src/pvec/containers_pvec.mli
+++ b/src/pvec/containers_pvec.mli
@@ -61,6 +61,7 @@ val drop_last : 'a t -> 'a t
 (** Like {!pop_opt} but doesn't return the last element.
   Returns the same vector if it's empty. *)
 
+val equal : ('a -> 'a -> bool) -> 'a t -> 'a t -> bool
 val iter : ('a -> unit) -> 'a t -> unit
 
 val iter_rev : ('a -> unit) -> 'a t -> unit
diff --git a/tests/pvec/t_pvec.ml b/tests/pvec/t_pvec.ml
index babdee8f..e58130aa 100644
--- a/tests/pvec/t_pvec.ml
+++ b/tests/pvec/t_pvec.ml
@@ -71,7 +71,43 @@ q _listuniq (fun l ->
 ;;
 
 t @@ fun () -> choose empty = None;;
-t @@ fun () -> choose (of_list [ 1, 1; 2, 2 ]) <> None
+t @@ fun () -> choose (of_list [ 1, 1; 2, 2 ]) <> None;;
+
+q
+  Q.(pair (small_list int) (small_list int))
+  (fun (l1, l2) -> equal CCInt.equal (of_list l1) (of_list l2) = (l1 = l2))
+;;
+
+q Q.(small_list int) (fun l1 -> equal CCInt.equal (of_list l1) (of_list l1))
+
+let arb_list_with_idx =
+  let open Q in
+  let shrink (l, i) =
+    Iter.(Shrink.(list ~shrink:int l) >|= fun l -> l, min i (List.length l - 1))
+  in
+  let gen =
+    Gen.(
+      let* l = small_list int in
+      let+ i =
+        if l = [] then
+          return 0
+        else
+          0 -- (List.length l - 1)
+      in
+      l, i)
+  in
+  make ~shrink ~print:Print.(pair (list int) int) gen
+;;
+
+q arb_list_with_idx (fun (l1, i) ->
+    if l1 <> [] then (
+      let l2 =
+        let x = List.nth l1 i in
+        CCList.set_at_idx i (x + 1) l1
+      in
+      not (equal CCInt.equal (of_list l1) (of_list l2))
+    ) else
+      true)
 
 module Ref_impl = struct
   type +'a t = 'a list