From 2083e3ee8659d553e734cf42a566a3432448eae2 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Sun, 25 Oct 2020 16:01:34 -0700
Subject: [PATCH] [sledge] Minor simplifications using Set and Map iterators

Reviewed By: jvillard

Differential Revision: D24313893

fbshipit-source-id: 8275b3127
---
 sledge/cli/sledge_cli.ml     | 4 ++--
 sledge/nonstdlib/iter.ml     | 1 +
 sledge/nonstdlib/iter.mli    | 1 +
 sledge/nonstdlib/map.ml      | 3 +++
 sledge/nonstdlib/map_intf.ml | 3 +++
 sledge/nonstdlib/multiset.ml | 6 ++----
 sledge/nonstdlib/set.ml      | 1 +
 sledge/nonstdlib/set_intf.ml | 1 +
 sledge/src/fol.ml            | 7 ++-----
 sledge/src/llair_to_Fol.ml   | 4 +---
 sledge/src/ses/var0.ml       | 8 ++------
 11 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/sledge/cli/sledge_cli.ml b/sledge/cli/sledge_cli.ml
index 725e76717..3575b2d74 100644
--- a/sledge/cli/sledge_cli.ml
+++ b/sledge/cli/sledge_cli.ml
@@ -95,8 +95,8 @@ let used_globals pgm preanalyze : Domain_used_globals.r =
       (Llair.Reg.Map.map summary_table ~f:Llair.Reg.Set.union_list)
   else
     Declared
-      (IArray.fold pgm.globals Llair.Reg.Set.empty ~f:(fun g ->
-           Llair.Reg.Set.add g.reg ))
+      (Llair.Reg.Set.of_iter
+         (Iter.map ~f:(fun g -> g.reg) (IArray.to_iter pgm.globals)))
 
 let analyze =
   let%map_open bound =
diff --git a/sledge/nonstdlib/iter.ml b/sledge/nonstdlib/iter.ml
index 66be7f802..5b7adae07 100644
--- a/sledge/nonstdlib/iter.ml
+++ b/sledge/nonstdlib/iter.ml
@@ -12,6 +12,7 @@ module Import = struct
 end
 
 let mem elt seq ~eq = mem ~eq ~x:elt seq
+let map seq ~f = map ~f seq
 let sort seq ~cmp = sort ~cmp seq
 let sort_uniq seq ~cmp = sort_uniq ~cmp seq
 let sorted seq ~cmp = sorted ~cmp seq
diff --git a/sledge/nonstdlib/iter.mli b/sledge/nonstdlib/iter.mli
index d52c3645b..06d866a74 100644
--- a/sledge/nonstdlib/iter.mli
+++ b/sledge/nonstdlib/iter.mli
@@ -12,6 +12,7 @@ module Import : sig
 end
 
 val mem : 'a -> 'a t -> eq:('a -> 'a -> bool) -> bool
+val map : 'a t -> f:('a -> 'b) -> 'b t
 val sort : 'a t -> cmp:('a -> 'a -> int) -> 'a t
 val sort_uniq : 'a t -> cmp:('a -> 'a -> int) -> 'a t
 val sorted : 'a t -> cmp:('a -> 'a -> int) -> bool
diff --git a/sledge/nonstdlib/map.ml b/sledge/nonstdlib/map.ml
index 9c4a1807a..aafa00a18 100644
--- a/sledge/nonstdlib/map.ml
+++ b/sledge/nonstdlib/map.ml
@@ -162,6 +162,9 @@ end) : S with type key = Key.t = struct
   let keys = M.keys
   let values = M.values
   let to_iter = M.to_iter
+  let to_list = M.to_list
+  let of_iter = M.of_iter
+  let of_list = M.of_list
 
   let to_iter2 l r =
     let seq = ref Iter.empty in
diff --git a/sledge/nonstdlib/map_intf.ml b/sledge/nonstdlib/map_intf.ml
index 939096e20..c881fa238 100644
--- a/sledge/nonstdlib/map_intf.ml
+++ b/sledge/nonstdlib/map_intf.ml
@@ -112,6 +112,9 @@ module type S = sig
   val keys : 'a t -> key iter
   val values : 'a t -> 'a iter
   val to_iter : 'a t -> (key * 'a) iter
+  val to_list : 'a t -> (key * 'a) list
+  val of_iter : (key * 'a) iter -> 'a t
+  val of_list : (key * 'a) list -> 'a t
 
   val to_iter2 :
        'a t
diff --git a/sledge/nonstdlib/multiset.ml b/sledge/nonstdlib/multiset.ml
index aae21459a..7dc9e7e40 100644
--- a/sledge/nonstdlib/multiset.ml
+++ b/sledge/nonstdlib/multiset.ml
@@ -33,15 +33,13 @@ struct
   let sexp_of_t s =
     List.sexp_of_t
       (Sexplib.Conv.sexp_of_pair Elt.sexp_of_t Mul.sexp_of_t)
-      (Iter.to_list (M.to_iter s))
+      (M.to_list s)
 
   let t_of_sexp elt_of_sexp sexp =
-    List.fold_left
-      ~f:(fun m (key, data) -> M.add_exn ~key ~data m)
+    M.of_list
       (List.t_of_sexp
          (Sexplib.Conv.pair_of_sexp elt_of_sexp Mul.t_of_sexp)
          sexp)
-      M.empty
 
   let pp sep pp_elt fs s =
     List.pp sep pp_elt fs (Iter.to_list (M.to_iter s))
diff --git a/sledge/nonstdlib/set.ml b/sledge/nonstdlib/set.ml
index 2a2b2d70f..9fc8979cd 100644
--- a/sledge/nonstdlib/set.ml
+++ b/sledge/nonstdlib/set.ml
@@ -80,6 +80,7 @@ end) : S with type elt = Elt.t = struct
   let for_all s ~f = S.for_all f s
   let fold s z ~f = S.fold f s z
   let to_iter = S.to_iter
+  let of_iter = S.of_iter
 
   let pp ?pre ?suf ?(sep = (",@ " : (unit, unit) fmt)) pp_elt fs x =
     List.pp ?pre ?suf sep pp_elt fs (S.elements x)
diff --git a/sledge/nonstdlib/set_intf.ml b/sledge/nonstdlib/set_intf.ml
index 567cde947..3cfd2fc23 100644
--- a/sledge/nonstdlib/set_intf.ml
+++ b/sledge/nonstdlib/set_intf.ml
@@ -60,6 +60,7 @@ module type S = sig
   (** {1 Convert} *)
 
   val to_iter : t -> elt iter
+  val of_iter : elt iter -> t
 
   (** {1 Pretty-print} *)
 
diff --git a/sledge/src/fol.ml b/sledge/src/fol.ml
index a1d4f9648..edb5b003c 100644
--- a/sledge/src/fol.ml
+++ b/sledge/src/fol.ml
@@ -853,9 +853,7 @@ let v_to_ses : var -> Ses.Var.t =
  fun v -> Ses.Var.identified ~id:(Var.id v) ~name:(Var.name v)
 
 let vs_to_ses : Var.Set.t -> Ses.Var.Set.t =
- fun vs ->
-  Var.Set.fold vs Ses.Var.Set.empty ~f:(fun v ->
-      Ses.Var.Set.add (v_to_ses v) )
+ fun vs -> Ses.Var.Set.of_iter (Iter.map ~f:v_to_ses (Var.Set.to_iter vs))
 
 let rec arith_to_ses poly =
   Arith.fold_monomials poly Ses.Term.zero ~f:(fun mono coeff e ->
@@ -936,8 +934,7 @@ let v_of_ses : Ses.Var.t -> var =
  fun v -> Var.identified ~id:(Ses.Var.id v) ~name:(Ses.Var.name v)
 
 let vs_of_ses : Ses.Var.Set.t -> Var.Set.t =
- fun vs ->
-  Ses.Var.Set.fold ~f:(fun v -> Var.Set.add (v_of_ses v)) vs Var.Set.empty
+ fun vs -> Var.Set.of_iter (Iter.map ~f:v_of_ses (Ses.Var.Set.to_iter vs))
 
 let uap1 f = ap1t (fun x -> _Apply f [|x|])
 let uap2 f = ap2t (fun x y -> _Apply f [|x; y|])
diff --git a/sledge/src/llair_to_Fol.ml b/sledge/src/llair_to_Fol.ml
index 2e1f5c654..6518400ac 100644
--- a/sledge/src/llair_to_Fol.ml
+++ b/sledge/src/llair_to_Fol.ml
@@ -16,9 +16,7 @@ let reg r =
   let global = Llair.Reg.is_global r in
   Var.program ~name ~global
 
-let regs rs =
-  Llair.Reg.Set.fold ~f:(fun r -> Var.Set.add (reg r)) rs Var.Set.empty
-
+let regs rs = Var.Set.of_iter (Iter.map ~f:reg (Llair.Reg.Set.to_iter rs))
 let uap0 f = T.apply f [||]
 let uap1 f a = T.apply f [|a|]
 let uap2 f a b = T.apply f [|a; b|]
diff --git a/sledge/src/ses/var0.ml b/sledge/src/ses/var0.ml
index 7964d6a81..1056e3211 100644
--- a/sledge/src/ses/var0.ml
+++ b/sledge/src/ses/var0.ml
@@ -94,12 +94,8 @@ module Make (T : REPR) = struct
       |> check (fun ({sub; _}, _) -> invariant sub)
 
     let fold sub z ~f = Map.fold ~f:(fun ~key ~data -> f key data) sub z
-
-    let domain sub =
-      Map.fold ~f:(fun ~key ~data:_ -> Set.add key) sub Set.empty
-
-    let range sub =
-      Map.fold ~f:(fun ~key:_ ~data -> Set.add data) sub Set.empty
+    let domain sub = Set.of_iter (Map.keys sub)
+    let range sub = Set.of_iter (Map.values sub)
 
     let invert sub =
       Map.fold sub empty ~f:(fun ~key ~data sub' ->