diff --git a/sledge/lib/import/map.ml b/sledge/lib/import/map.ml index dc755027c..0764bb7ea 100644 --- a/sledge/lib/import/map.ml +++ b/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 = diff --git a/sledge/lib/import/map_intf.ml b/sledge/lib/import/map_intf.ml index 5bdf92a6f..6bac231cb 100644 --- a/sledge/lib/import/map_intf.ml +++ b/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)]. *) diff --git a/sledge/lib/import/qset.ml b/sledge/lib/import/qset.ml index e2e89423b..dc38469a5 100644 --- a/sledge/lib/import/qset.ml +++ b/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 diff --git a/sledge/lib/import/qset_intf.ml b/sledge/lib/import/qset_intf.ml index cf0067b61..5446290f9 100644 --- a/sledge/lib/import/qset_intf.ml +++ b/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)]. *) diff --git a/sledge/lib/term.ml b/sledge/lib/term.ml index 271395968..de1e2fc57 100644 --- a/sledge/lib/term.ml +++ b/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)