From ff956b9cc2e08540ffde95a95895d6ac50e1b2e0 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 30 Mar 2015 21:05:58 +0200
Subject: [PATCH] add `containers_misc.Puf.iter`

---
 src/misc/puf.ml  | 15 ++++++++++++++-
 src/misc/puf.mli |  4 ++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/misc/puf.ml b/src/misc/puf.ml
index 2a338fd2..ef4ca0c8 100644
--- a/src/misc/puf.ml
+++ b/src/misc/puf.ml
@@ -58,6 +58,8 @@ module PArray = struct
         a
       end
 
+  let iteri f t = Array.iteri f (reroot t)
+
   let get t i =
     match !t with
     | Array a -> a.(i)
@@ -204,6 +206,9 @@ module type S = sig
     (** [iter_equiv_class uf a f] calls [f] on every element of [uf] that
         is congruent to [a], including [a] itself. *)
 
+  val iter : _ t -> (elt -> unit) -> unit
+    (** Iterate on all root values *)
+
   val inconsistent : _ t -> (elt * elt * elt * elt) option
     (** Check whether the UF is inconsistent. It returns [Some (a, b, a', b')]
         in case of inconsistency, where a = b, a = a' and b = b' by congruence,
@@ -222,7 +227,8 @@ module type S = sig
 
   val explain_distinct : 'e t -> elt -> elt -> elt * elt
     (** [explain_distinct uf a b] gives the original pair [a', b'] that
-        made [a] and [b] distinct by calling [distinct a' b'] *)
+        made [a] and [b] distinct by calling [distinct a' b']. The
+        terms must be distinct, otherwise Failure is raised. *)
 end
 
 module IH = Hashtbl.Make(struct type t = int let equal i j = i = j let hash i = i end)
@@ -446,6 +452,13 @@ module Make(X : ID) : S with type elt = X.t = struct
         in
         traverse ia
 
+  let iter uf f =
+    PArray.iteri
+      (fun i e -> match e with
+         | None -> ()
+         | Some d -> if d.next = i then f d.elt
+      ) uf.data
+
   let inconsistent uf = uf.inconsistent
 
   (** Closest common ancestor of the two elements in the proof forest *)
diff --git a/src/misc/puf.mli b/src/misc/puf.mli
index c44f4e2b..6ae10d5e 100644
--- a/src/misc/puf.mli
+++ b/src/misc/puf.mli
@@ -113,6 +113,10 @@ module type S = sig
     (** [iter_equiv_class uf a f] calls [f] on every element of [uf] that
         is congruent to [a], including [a] itself. *)
 
+  val iter : _ t -> (elt -> unit) -> unit
+  (** Iterate on all root values
+      @since NExT_RELEASE *)
+
   val inconsistent : _ t -> (elt * elt * elt * elt) option
     (** Check whether the UF is inconsistent. It returns [Some (a, b, a', b')]
         in case of inconsistency, where a = b, a = a' and b = b' by congruence,