Add is_singleton_or_more to Set and Map

Reviewed By: skcho

Differential Revision: D13062696

fbshipit-source-id: 3a7286f55

infer/src/bufferoverrun/absLoc.ml

@@ -122,9 +122,6 @@ module PowLoc = struct
   let append_field ploc ~fn =
     if is_bot ploc then singleton Loc.unknown
     else fold (fun l -> add (Loc.append_field l ~fn)) ploc empty
-
-
-  let is_singleton x = Int.equal (cardinal x) 1
 end
 
 (** unsound but ok for bug catching *)
@@ -133,5 +130,9 @@ let always_strong_update = true
 let can_strong_update : PowLoc.t -> bool =
  fun ploc ->
   if always_strong_update then true
-  else if Int.equal (PowLoc.cardinal ploc) 1 then Loc.is_var (PowLoc.choose ploc)
-  else false
+  else
+    match PowLoc.is_singleton_or_more ploc with
+    | IContainer.Singleton loc ->
+        Loc.is_var loc
+    | _ ->
+        false

infer/src/bufferoverrun/arrayBlk.ml

@@ -154,10 +154,11 @@ let normalize : astate -> astate = fun a -> map ArrInfo.normalize a
 
 let do_prune : (ArrInfo.t -> ArrInfo.t -> ArrInfo.t) -> astate -> astate -> astate =
  fun arr_info_prune a1 a2 ->
-  if Int.equal (cardinal a2) 1 then
-    let k, v2 = choose a2 in
-    if mem k a1 then add k (arr_info_prune (find k a1) v2) a1 else a1
-  else a1
+  match is_singleton_or_more a2 with
+  | IContainer.Singleton (k, v2) ->
+      if mem k a1 then add k (arr_info_prune (find k a1) v2) a1 else a1
+  | _ ->
+      a1
 
 
 let prune_comp : Binop.t -> astate -> astate -> astate =

infer/src/bufferoverrun/bounds.ml

@@ -104,8 +104,8 @@ module SymLinear = struct
      given coefficient. *)
   let one_symbol_of_coeff : NonZeroInt.t -> t -> Symb.Symbol.t option =
    fun coeff x ->
-    match M.is_singleton x with
-    | Some (k, v) when Z.equal (v :> Z.t) (coeff :> Z.t) ->
+    match M.is_singleton_or_more x with
+    | IContainer.Singleton (k, v) when Z.equal (v :> Z.t) (coeff :> Z.t) ->
         Some k
     | _ ->
         None

infer/src/bufferoverrun/bufferOverrunChecker.ml

@@ -52,7 +52,11 @@ module TransferFunctions (CFG : ProcCfg.S) = struct
         Option.find_map (Loc.get_path l) ~f:(fun partial ->
             try
               let locs = eval_locs_sympath_partial partial in
-              if PowLoc.is_singleton locs then Some (PowLoc.choose locs) else None
+              match PowLoc.is_singleton_or_more locs with
+              | IContainer.Singleton loc ->
+                  Some loc
+              | _ ->
+                  None
             with Caml.Not_found -> None )
       in
       let ret_alias =
@@ -182,19 +186,22 @@ module TransferFunctions (CFG : ProcCfg.S) = struct
         in
         let mem = Dom.Mem.update_mem locs v mem in
         let mem =
-          if PowLoc.is_singleton locs && not v.represents_multiple_values then
-            let loc_v = PowLoc.min_elt locs in
-            let pname = Procdesc.get_proc_name pdesc in
-            match Typ.Procname.get_method pname with
-            | "__inferbo_empty" when Loc.is_return loc_v -> (
-              match Sem.get_formals pdesc with
-              | [(formal, _)] ->
-                  let formal_v = Dom.Mem.find (Loc.of_pvar formal) mem in
-                  Dom.Mem.store_empty_alias formal_v loc_v exp2 mem
-              | _ ->
-                  assert false )
+          if not v.represents_multiple_values then
+            match PowLoc.is_singleton_or_more locs with
+            | IContainer.Singleton loc_v -> (
+                let pname = Procdesc.get_proc_name pdesc in
+                match Typ.Procname.get_method pname with
+                | "__inferbo_empty" when Loc.is_return loc_v -> (
+                  match Sem.get_formals pdesc with
+                  | [(formal, _)] ->
+                      let formal_v = Dom.Mem.find (Loc.of_pvar formal) mem in
+                      Dom.Mem.store_empty_alias formal_v loc_v exp2 mem
+                  | _ ->
+                      assert false )
+                | _ ->
+                    Dom.Mem.store_simple_alias loc_v exp2 mem )
             | _ ->
-                Dom.Mem.store_simple_alias loc_v exp2 mem
+                mem
           else mem
         in
         let mem = Dom.Mem.update_latest_prune exp1 exp2 mem in

infer/src/bufferoverrun/bufferOverrunDomain.ml

@@ -526,9 +526,11 @@ module Alias = struct
    fun formal loc e a ->
     let a = lift (AliasMap.store loc e) a in
     let locs = Val.get_all_locs formal in
-    if PowLoc.is_singleton locs then
-      (fst a, AliasRet.L (AliasTarget.of_empty (PowLoc.min_elt locs)))
-    else a
+    match PowLoc.is_singleton_or_more locs with
+    | IContainer.Singleton loc ->
+        (fst a, AliasRet.L (AliasTarget.of_empty loc))
+    | _ ->
+        a
 
 
   let remove_temp : Ident.t -> astate -> astate =

infer/src/bufferoverrun/bufferOverrunDomainRelation.ml

@@ -1364,25 +1364,24 @@ module Make (Manager : Manager_S) = struct
         VarSet.fold (fun var acc -> VarMap.add var id acc) vars pack_ids
       in
       let vars_ids = pack_ids_of_vars vars x in
-      let num_vars_ids = PackSet.cardinal vars_ids in
-      if Int.equal num_vars_ids 0 then
-        let id = get_new_id () in
-        {x with pack_ids= set_pack_id_of_vars vars id x.pack_ids}
-      else if Int.equal num_vars_ids 1 then
-        let id = PackSet.choose vars_ids in
-        {x with pack_ids= set_pack_id_of_vars vars id x.pack_ids}
-      else
-        let id = PackSet.min_elt vars_ids in
-        let other_ids = PackSet.remove id vars_ids in
-        let pack_ids =
-          x.pack_ids |> set_pack_id_of_vars vars id
-          |> VarMap.map (PackSet.subst ~from:other_ids ~to_:id)
-        in
-        let packs =
-          let v = val_of_pack_ids vars_ids x in
-          x.packs |> PackMap.remove_packs other_ids |> PackMap.add id v
-        in
-        {pack_ids; packs}
+      match PackSet.is_singleton_or_more vars_ids with
+      | IContainer.Empty ->
+          let id = get_new_id () in
+          {x with pack_ids= set_pack_id_of_vars vars id x.pack_ids}
+      | IContainer.Singleton id ->
+          {x with pack_ids= set_pack_id_of_vars vars id x.pack_ids}
+      | IContainer.More ->
+          let id = PackSet.min_elt vars_ids in
+          let other_ids = PackSet.remove id vars_ids in
+          let pack_ids =
+            x.pack_ids |> set_pack_id_of_vars vars id
+            |> VarMap.map (PackSet.subst ~from:other_ids ~to_:id)
+          in
+          let packs =
+            let v = val_of_pack_ids vars_ids x in
+            x.packs |> PackMap.remove_packs other_ids |> PackMap.add id v
+          in
+          {pack_ids; packs}
 
 
     let subst ~forget_free subst_map x =

infer/src/bufferoverrun/bufferOverrunUtils.ml

@@ -26,8 +26,12 @@ module Exec = struct
     let locs = val_ |> Dom.Val.get_all_locs in
     let v = Dom.Mem.find_set locs mem in
     let mem = Dom.Mem.add_stack (Loc.of_id id) v mem in
-    if PowLoc.is_singleton locs && not v.represents_multiple_values then
-      Dom.Mem.load_simple_alias id (PowLoc.min_elt locs) mem
+    if not v.represents_multiple_values then
+      match PowLoc.is_singleton_or_more locs with
+      | IContainer.Singleton loc ->
+          Dom.Mem.load_simple_alias id loc mem
+      | _ ->
+          mem
     else mem
 
 

infer/src/bufferoverrun/symb.ml

@@ -123,13 +123,4 @@ module SymbolMap = struct
         true
     | exception Exit ->
         false
-
-
-  let is_singleton : 'a t -> (key * 'a) option =
-   fun m ->
-    if is_empty m then None
-    else
-      let ((kmin, _) as binding) = min_binding m in
-      let kmax, _ = max_binding m in
-      if Symbol.equal kmin kmax then Some binding else None
 end

infer/src/bufferoverrun/symb.mli

@@ -63,8 +63,6 @@ module SymbolMap : sig
   include PrettyPrintable.PPMap with type key = Symbol.t
 
   val for_all2 : f:(key -> 'a option -> 'b option -> bool) -> 'a t -> 'b t -> bool
-
-  val is_singleton : 'a t -> (key * 'a) option
 end
 
 module SymbolTable : sig

infer/src/checkers/functionPointers.ml

@@ -46,10 +46,15 @@ module Analyzer = AbstractInterpreter.MakeRPO (TransferFunctions (CFG))
 
 let find_procname var astate =
   match Domain.find_opt (Ident.to_string var) astate with
-  | Some procnames ->
-      if ProcnameSet.is_empty procnames then None
-      else Some (ProcnameSet.min_elt procnames)
-        (* TODO: handle multiple procnames, e.g. with non-determinism branching *)
+  | Some procnames -> (
+    match ProcnameSet.is_singleton_or_more procnames with
+    | IContainer.Empty ->
+        None
+    | IContainer.Singleton procname ->
+        Some procname
+    | IContainer.More ->
+        Some (ProcnameSet.min_elt procnames)
+        (* TODO: handle multiple procnames, e.g. with non-determinism branching *) )
   | None ->
       None
 

infer/src/checkers/loopInvariant.ml

@@ -44,27 +44,23 @@ let is_non_primitive typ = Typ.is_pointer typ || Typ.is_struct typ
 let is_def_unique_and_satisfy tenv var (loop_nodes : LoopNodes.t) ~is_inv_by_default
     is_exp_invariant =
   let equals_var id = Var.equal var (Var.of_id id) in
-  (* Use O(1) is_singleton check *)
-  (* tedious parameter wrangling to make IContainer's fold interface happy *)
-  IContainer.is_singleton
-    ~fold:(fun s ~init ~f -> LoopNodes.fold (fun node acc -> f acc node) s init)
-    loop_nodes
-  && LoopNodes.for_all
-       (fun node ->
-         Procdesc.Node.get_instrs node
-         |> Instrs.exists ~f:(function
-              | Sil.Load (id, exp_rhs, _, _) when equals_var id && is_exp_invariant exp_rhs ->
-                  true
-              | Sil.Store (exp_lhs, _, exp_rhs, _)
-                when Exp.equal exp_lhs (Var.to_exp var) && is_exp_invariant exp_rhs ->
-                  true
-              | Sil.Call ((id, _), Const (Cfun callee_pname), args, _, _) when equals_var id ->
-                  PurityDomain.is_pure (get_purity tenv ~is_inv_by_default callee_pname args)
-                  && (* check if all params are invariant *)
-                     List.for_all ~f:(fun (exp, _) -> is_exp_invariant exp) args
-              | _ ->
-                  false ) )
-       loop_nodes
+  match LoopNodes.is_singleton_or_more loop_nodes with
+  | IContainer.Singleton node ->
+      Procdesc.Node.get_instrs node
+      |> Instrs.exists ~f:(function
+           | Sil.Load (id, exp_rhs, _, _) when equals_var id && is_exp_invariant exp_rhs ->
+               true
+           | Sil.Store (exp_lhs, _, exp_rhs, _)
+             when Exp.equal exp_lhs (Var.to_exp var) && is_exp_invariant exp_rhs ->
+               true
+           | Sil.Call ((id, _), Const (Cfun callee_pname), args, _, _) when equals_var id ->
+               PurityDomain.is_pure (get_purity tenv ~is_inv_by_default callee_pname args)
+               && (* check if all params are invariant *)
+                  List.for_all ~f:(fun (exp, _) -> is_exp_invariant exp) args
+           | _ ->
+               false )
+  | _ ->
+      false
 
 
 let is_exp_invariant inv_vars invalidated_vars loop_nodes reaching_defs exp =

infer/src/istd/IContainer.mli

@@ -16,6 +16,7 @@ val singleton_or_more :
 (* O(1) *)
 
 val is_singleton : fold:('t, 'a, 'a singleton_or_more) Container.fold -> 't -> bool
+  [@@warning "-32"]
 
 val mem_nth : fold:('t, _, int) Container.fold -> 't -> int -> bool
 

infer/src/istd/PrettyPrintable.ml

@@ -19,6 +19,8 @@ end
 module type PPSet = sig
   include Caml.Set.S
 
+  val is_singleton_or_more : t -> elt IContainer.singleton_or_more
+
   val pp_element : F.formatter -> elt -> unit
 
   val pp : F.formatter -> t -> unit
@@ -27,6 +29,8 @@ end
 module type PPMap = sig
   include Caml.Map.S
 
+  val is_singleton_or_more : 'a t -> (key * 'a) IContainer.singleton_or_more
+
   val pp_key : F.formatter -> key -> unit
 
   val pp : pp_value:(F.formatter -> 'a -> unit) -> F.formatter -> 'a t -> unit
@@ -37,6 +41,14 @@ let pp_collection ~pp_item fmt c = IContainer.pp_collection ~fold:List.fold ~pp_
 module MakePPSet (Ord : PrintableOrderedType) = struct
   include Caml.Set.Make (Ord)
 
+  let is_singleton_or_more s =
+    if is_empty s then IContainer.Empty
+    else
+      let mi = min_elt s in
+      let ma = max_elt s in
+      if phys_equal mi ma then IContainer.Singleton mi else IContainer.More
+
+
   let pp_element = Ord.pp
 
   let pp fmt s = pp_collection ~pp_item:pp_element fmt (elements s)
@@ -45,6 +57,14 @@ end
 module MakePPMap (Ord : PrintableOrderedType) = struct
   include Caml.Map.Make (Ord)
 
+  let is_singleton_or_more m =
+    if is_empty m then IContainer.Empty
+    else
+      let ((kmi, _) as binding) = min_binding m in
+      let kma, _ = max_binding m in
+      if phys_equal kmi kma then IContainer.Singleton binding else IContainer.More
+
+
   let pp_key = Ord.pp
 
   let pp ~pp_value fmt m =

infer/src/istd/PrettyPrintable.mli

@@ -21,6 +21,8 @@ end
 module type PPSet = sig
   include Caml.Set.S
 
+  val is_singleton_or_more : t -> elt IContainer.singleton_or_more
+
   val pp_element : F.formatter -> elt -> unit
 
   val pp : F.formatter -> t -> unit
@@ -29,6 +31,8 @@ end
 module type PPMap = sig
   include Caml.Map.S
 
+  val is_singleton_or_more : 'a t -> (key * 'a) IContainer.singleton_or_more
+
   val pp_key : F.formatter -> key -> unit
 
   val pp : pp_value:(F.formatter -> 'a -> unit) -> F.formatter -> 'a t -> unit