[Uf] move union move for repr and added to egraph

src/egraph_simple.ml

@@ -52,6 +52,17 @@ let has_repr env t =
 
 let print_repr = UnionFind.print_repr
 
+let union env cl1 cl2 =
+  let prio1,prio2,env = UnionFind.union env cl1 cl2 in
+  let env = if Cl.equal prio2 cl2 then env else
+      let env = UnionFind.set_repr env prio2 (repr env cl2) in
+      let env = if UnionFind.is_added env cl2
+        then UnionFind.added env prio2
+        else env in
+      env
+  in
+  prio1,prio2,env
+
 (** Normalize the term and update the use
     precondition cl is_findable
 *)
@@ -193,7 +204,7 @@ let rec equal_aux nb queue env cl1 cl2 =
       end
     else
       begin
-        let cl1,cl2,env = UnionFind.union env cl1 cl2 in
+        let cl1,cl2,env = union env cl1 cl2 in
         cl1,cl2,env
       end
   in

src/term.ml

@@ -125,15 +125,6 @@ struct
     let prio1,prio2 = Shuffle.shufflep (cl1,cl2) in
     Debug.dprintf debug "[Uf]@[ Union@ @[%a@]@ @[%a@]@] -> %a@\n"
       (print_repr t) cl1 (print_repr t) cl2 Cl.print prio2;
-    let t = if Cl.equal prio2 cl2 then t else
-        let move m =
-          try
-            Cl.M.add prio2 (Cl.M.find cl2 m) m
-          with Not_found -> m in
-        {t with
-          repr = move t.repr;
-          added = move t.added}
-    in
     prio1,prio2,{ t with eq = Cl.M.add prio1 prio2 t.eq }
 
   let union_force t cl1 cl2 =
@@ -208,7 +199,6 @@ struct
           let m = Cl.M.remove cl m in
           if Cl.M.is_empty m then None else Some m) arg t.use}
 
-
   let is_added t cl = Cl.S.mem cl t.added
 
   let added t cl = {t with added = Cl.S.add cl t.added}