[sledge] Strengthen normalization of division

Summary: `x / (q × y) ==> (1/q × x) / y`

Reviewed By: jvillard

Differential Revision: D21042526

fbshipit-source-id: 232a159c9

sledge/lib/import/map.ml

@@ -51,19 +51,34 @@ end) : S with type key = Key.t = struct
     let f s (k, v) = f ~key:k ~data:v s in
     Container.fold_until ~fold ~init ~f ~finish m
 
-  let choose m =
+  let root_key_exn m =
+    with_return
+    @@ fun {return} ->
+    binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data:_ ->
+        return key )
+    |> ignore ;
+    raise (Not_found_s (Atom __LOC__))
+
+  let choose_exn m =
     with_return
     @@ fun {return} ->
     binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data ->
-        return (Some (key, data)) )
+        return (key, data) )
     |> ignore ;
-    None
+    raise (Not_found_s (Atom __LOC__))
 
+  let choose m = try Some (choose_exn m) with Not_found_s _ -> None
   let pop m = choose m |> Option.map ~f:(fun (k, v) -> (k, v, remove m k))
 
   let pop_min_elt m =
     min_elt m |> Option.map ~f:(fun (k, v) -> (k, v, remove m k))
 
+  let is_singleton m =
+    try
+      let l, _, r = split m (root_key_exn m) in
+      is_empty l && is_empty r
+    with Not_found_s _ -> false
+
   let find_and_remove m k =
     let found = ref None in
     let m =

sledge/lib/import/map_intf.ml

@@ -19,6 +19,7 @@ module type S = sig
   include Core_kernel.Map_intf.Make_S_plain_tree(Key).S
 
   val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int
+  val is_singleton : 'a t -> bool
 
   val map_endo : 'a t -> f:('a -> 'a) -> 'a t
   (** Like map, but specialized to require [f] to be an endofunction, which
@@ -46,6 +47,9 @@ module type S = sig
     -> finish:('a -> 'r)
     -> 'r
 
+  val choose_exn : 'a t -> key * 'a
+  (** Find an unspecified element. [O(1)]. *)
+
   val choose : 'a t -> (key * 'a) option
   (** Find an unspecified element. [O(1)]. *)
 

sledge/lib/import/qset.ml

@@ -74,9 +74,11 @@ struct
 
   let map_counts m ~f = M.mapi ~f:(fun ~key ~data -> f key data) m
   let is_empty = M.is_empty
+  let is_singleton = M.is_singleton
   let length m = M.length m
   let count m x = match M.find m x with Some q -> q | None -> Q.zero
   let choose = M.choose
+  let choose_exn = M.choose_exn
   let pop = M.pop
   let min_elt = M.min_elt
   let pop_min_elt = M.pop_min_elt

sledge/lib/import/qset_intf.ml

@@ -46,6 +46,7 @@ module type S = sig
   (* queries *)
 
   val is_empty : t -> bool
+  val is_singleton : t -> bool
 
   val length : t -> int
   (** Number of elements with non-zero multiplicity. [O(1)]. *)
@@ -53,6 +54,9 @@ module type S = sig
   val count : t -> elt -> Q.t
   (** Multiplicity of an element. [O(log n)]. *)
 
+  val choose_exn : t -> elt * Q.t
+  (** Find an unspecified element. [O(1)]. *)
+
   val choose : t -> (elt * Q.t) option
   (** Find an unspecified element. [O(1)]. *)
 

sledge/lib/term.ml

@@ -592,6 +592,10 @@ let rec simp_div x y =
   | _, Rational {data} -> Sum.to_term (Sum.of_ ~coeff:Q.(inv data) x)
   (* e / (n / d) ==> (e × d) / n *)
   | e, Ap2 (Div, n, d) -> simp_div (simp_mul2 e d) n
+  (* x / (q × y) ==> (1/q × x) / y *)
+  | _, Add sum when Qset.is_singleton sum ->
+      let y, q = Qset.choose_exn sum in
+      simp_div (simp_mul2 (rational (Q.inv q)) x) y
   (* x / y *)
   | _ -> Ap2 (Div, x, y)