[sledge] Optimize map operations over formulas

Summary:
An invariant of `And` and `Or` formulas is that they are flattened,
that is, `And` formulas do not have positive immediate subformulas of
form `And` nor negative immediate subformulas of form `Or`, and
similarly for `Or`. This invariant is ensured by the formula
constructors, which scan the immediate subformulas for cases that need
to be flattened.

When mapping over formulas, this repeated scanning is a performance
bottleneck. Most of the work of flattening is wasted since the input
formulas are necessarily already in flattened form, and the common
case is that the transformation preserves the flattened form. This
diff optimizes this case by detecting violations of flattening during
the transformation, performing the needed flattening, and avoiding the
general constructor that scans for flattening violations.

Reviewed By: jvillard

Differential Revision: D26338014

fbshipit-source-id: 9f15cca58

sledge/src/fol/exp.ml

@@ -394,8 +394,8 @@ module Formula = struct
     | Eq0 x -> lift_map1 f b Fml.eq0 x
     | Pos x -> lift_map1 f b Fml.pos x
     | Not x -> map1 (map_terms ~f) b Fml.not_ x
-    | And {pos; neg} -> Fml.map_pos_neg (map_terms ~f) b Fml.andN ~pos ~neg
-    | Or {pos; neg} -> Fml.map_pos_neg (map_terms ~f) b Fml.orN ~pos ~neg
+    | And {pos; neg} -> Fml.map_and b ~pos ~neg (map_terms ~f)
+    | Or {pos; neg} -> Fml.map_or b ~pos ~neg (map_terms ~f)
     | Iff (x, y) -> map2 (map_terms ~f) b Fml.iff x y
     | Cond {cnd; pos; neg} ->
         map3 (map_terms ~f) b

sledge/src/fol/fml.ml

@@ -147,9 +147,6 @@ let iff = _Iff
 let cond ~cnd ~pos ~neg = _Cond cnd pos neg
 let lit = _Lit
 
-let map_pos_neg f e cons ~pos ~neg =
-  map2 (Set.map ~f) e (fun pos neg -> cons ~pos ~neg) pos neg
-
 let rec map_trms b ~f =
   match b with
   | Tt -> b
@@ -157,8 +154,8 @@ let rec map_trms b ~f =
   | Eq0 x -> map1 f b _Eq0 x
   | Pos x -> map1 f b _Pos x
   | Not x -> map1 (map_trms ~f) b _Not x
-  | And {pos; neg} -> map_pos_neg (map_trms ~f) b _And ~pos ~neg
-  | Or {pos; neg} -> map_pos_neg (map_trms ~f) b _Or ~pos ~neg
+  | And {pos; neg} -> map_and b ~pos ~neg (map_trms ~f)
+  | Or {pos; neg} -> map_or b ~pos ~neg (map_trms ~f)
   | Iff (x, y) -> map2 (map_trms ~f) b _Iff x y
   | Cond {cnd; pos; neg} -> map3 (map_trms ~f) b _Cond cnd pos neg
   | Lit (p, xs) -> mapN f b (_Lit p) xs

sledge/src/fol/fml.mli

@@ -66,9 +66,8 @@ val lit : Predsym.t -> Trm.t array -> t
 
 val map_vars : t -> f:(Var.t -> Var.t) -> t
 val map_trms : t -> f:(Trm.t -> Trm.t) -> t
-
-val map_pos_neg :
-  (t -> t) -> 'a -> (pos:set -> neg:set -> 'a) -> pos:set -> neg:set -> 'a
+val map_and : t -> pos:set -> neg:set -> (t -> t) -> t
+val map_or : t -> pos:set -> neg:set -> (t -> t) -> t
 
 (** Traverse *)
 

sledge/src/fol/propositional.ml

@@ -152,7 +152,9 @@ struct
       |> check invariant
 
     let and_ p q =
-      join _And ff
+      join
+        (_Join (fun ~pos ~neg -> And {pos; neg}) tt ff)
+        ff
         (function
           | And {pos; neg} -> (pos, neg)
           | Tt -> (Fmls.empty, Fmls.empty)
@@ -161,7 +163,9 @@ struct
         p q
 
     let or_ p q =
-      join _Or tt
+      join
+        (_Join (fun ~pos ~neg -> Or {pos; neg}) ff tt)
+        tt
         (function
           | Or {pos; neg} -> (pos, neg)
           | Not Tt -> (Fmls.empty, Fmls.empty)
@@ -263,6 +267,57 @@ struct
       | Lit (_, xs) -> Array.iter ~f xs
 
     let trms p = Iter.from_labelled_iter (iter_trms p)
+
+    type polarity = Con | Dis
+
+    let map_join polarity b ~pos ~neg f =
+      let pos_to_flatten = ref [] in
+      let neg_to_flatten = ref [] in
+      let pos0, neg0 = (pos, neg) in
+      let pos =
+        Fmls.map pos ~f:(fun p ->
+            let p' = f p in
+            if p' == p then p
+            else
+              match (polarity, p') with
+              | Con, And {pos; neg} | Dis, Or {pos; neg} ->
+                  pos_to_flatten := (p, pos, neg) :: !pos_to_flatten ;
+                  p
+              | _ -> p' )
+      in
+      let neg =
+        Fmls.map neg ~f:(fun n ->
+            let n' = f n in
+            if n' == n then n
+            else
+              match (polarity, n') with
+              | Con, Or {pos; neg} | Dis, And {pos; neg} ->
+                  neg_to_flatten := (n, pos, neg) :: !neg_to_flatten ;
+                  n
+              | _ -> n' )
+      in
+      let pos, neg =
+        if List.is_empty !pos_to_flatten then (pos, neg)
+        else
+          List.fold !pos_to_flatten (pos, neg)
+            ~f:(fun (p, p', n') (pos, neg) ->
+              (Fmls.union p' (Fmls.remove p pos), Fmls.union n' neg) )
+      in
+      let pos, neg =
+        if List.is_empty !neg_to_flatten then (pos, neg)
+        else
+          List.fold !neg_to_flatten (pos, neg)
+            ~f:(fun (n, p', n') (pos, neg) ->
+              (Fmls.union n' pos, Fmls.union p' (Fmls.remove n neg)) )
+      in
+      if pos0 == pos && neg0 == neg then b
+      else
+        match polarity with
+        | Con -> _Join (fun ~pos ~neg -> And {pos; neg}) tt ff ~pos ~neg
+        | Dis -> _Join (fun ~pos ~neg -> Or {pos; neg}) ff tt ~pos ~neg
+
+    let map_and = map_join Con
+    let map_or = map_join Dis
   end
 end
 [@@inline]

sledge/src/fol/propositional_intf.ml

@@ -46,6 +46,8 @@ module type FORMULA = sig
   val or_ : t -> t -> t
   val is_negative : t -> bool
   val trms : t -> trm iter
+  val map_and : t -> pos:set -> neg:set -> (t -> t) -> t
+  val map_or : t -> pos:set -> neg:set -> (t -> t) -> t
 end
 
 (** Sets of formulas *)