[pulse] Avoid partitioning abstract values

Summary:
`partition` always constructs two new maps, which is expensive when
there are a lot of entries.  Let's avoid it if possible.

Reviewed By: jvillard

Differential Revision: D21684298

fbshipit-source-id: a8674d358

infer/src/pulse/PulseAbductiveDomain.ml

@@ -26,13 +26,10 @@ module type BaseDomainSig = sig
   val update : ?stack:BaseStack.t -> ?heap:BaseMemory.t -> ?attrs:BaseAddressAttributes.t -> t -> t
 
   val filter_addr : f:(AbstractValue.t -> bool) -> t -> t
-  (**filter both heap and attrs *)
+  (** filter both heap and attrs *)
 
-  val partition_addr :
-       f:(AbstractValue.t -> bool)
-    -> t
-    -> (BaseMemory.t * BaseAddressAttributes.t) * (BaseMemory.t * BaseAddressAttributes.t)
-  (**partition both heap and attrs *)
+  val filter_addr_with_discarded_attrs : f:(AbstractValue.t -> bool) -> t -> t * Attributes.t list
+  (** filter both heap and attrs with returning discarded attrs together *)
 
   val pp : F.formatter -> t -> unit
 end
@@ -63,12 +60,12 @@ module BaseDomainCommon = struct
     update ~heap:heap' ~attrs:attrs' foot
 
 
-  let partition_addr ~f foot =
-    let heap_yes, heap_no = BaseMemory.partition (fun address _ -> f address) foot.heap in
-    let attrs_yes, attrs_no =
-      BaseAddressAttributes.partition (fun address _ -> f address) foot.attrs
+  let filter_addr_with_discarded_attrs ~f foot =
+    let heap' = BaseMemory.filter (fun address _ -> f address) foot.heap in
+    let attrs', discarded_attributes =
+      BaseAddressAttributes.filter_with_discarded_attrs (fun address _ -> f address) foot.attrs
     in
-    ((heap_yes, attrs_yes), (heap_no, attrs_no))
+    (update ~heap:heap' ~attrs:attrs' foot, discarded_attributes)
 end
 
 (** represents the post abstract state at each program point *)
@@ -339,10 +336,11 @@ let discard_unreachable ({pre; post} as astate) =
   in
   let post_addresses = BaseDomain.reachable_addresses (post :> BaseDomain.t) in
   let live_addresses = AbstractValue.Set.union pre_addresses post_addresses in
-  let (heap_new, attrs_new), (_, attrs_unreachable) =
-    PostDomain.partition_addr ~f:(fun address -> AbstractValue.Set.mem address live_addresses) post
+  let post_new, attrs_unreachable =
+    PostDomain.filter_addr_with_discarded_attrs
+      ~f:(fun address -> AbstractValue.Set.mem address live_addresses)
+      post
   in
-  let post_new = PostDomain.update ~heap:heap_new ~attrs:attrs_new post in
   (* note: we don't call {!PulsePathCondition.simplify} *)
   let astate =
     if phys_equal pre_new pre && phys_equal post_new post then astate

infer/src/pulse/PulseAbductiveDomain.mli

@@ -30,13 +30,10 @@ module type BaseDomainSig = sig
   val update : ?stack:BaseStack.t -> ?heap:BaseMemory.t -> ?attrs:BaseAddressAttributes.t -> t -> t
 
   val filter_addr : f:(AbstractValue.t -> bool) -> t -> t
-  (**filter both heap and attrs *)
+  (** filter both heap and attrs *)
 
-  val partition_addr :
-       f:(AbstractValue.t -> bool)
-    -> t
-    -> (BaseMemory.t * BaseAddressAttributes.t) * (BaseMemory.t * BaseAddressAttributes.t)
-  (**partition both heap and attrs *)
+  val filter_addr_with_discarded_attrs : f:(AbstractValue.t -> bool) -> t -> t * Attributes.t list
+  (** filter both heap and attrs with returning discarded attrs together *)
 
   val pp : F.formatter -> t -> unit
 end
@@ -143,7 +140,7 @@ val is_local : Var.t -> t -> bool
 
 val find_post_cell_opt : AbstractValue.t -> t -> BaseDomain.cell option
 
-val discard_unreachable : t -> t * AbstractValue.Set.t * BaseAddressAttributes.t
+val discard_unreachable : t -> t * AbstractValue.Set.t * Attributes.t list
 (** [discard_unreachable astate] garbage collects unreachable addresses in the state to make it
     smaller, and retuns the new state, the live addresses, and the attributes of discarded addresses *)
 

infer/src/pulse/PulseBaseAddressAttributes.ml

@@ -54,7 +54,11 @@ let empty = Graph.empty
 
 let filter = Graph.filter
 
-let partition = Graph.partition
+let filter_with_discarded_attrs f x =
+  fold
+    (fun k v ((x, discarded) as acc) -> if f k v then acc else (Graph.remove k x, v :: discarded))
+    x (x, [])
+
 
 let pp = Graph.pp
 

infer/src/pulse/PulseBaseAddressAttributes.mli

@@ -14,7 +14,8 @@ val empty : t
 
 val filter : (AbstractValue.t -> Attributes.t -> bool) -> t -> t
 
-val partition : (AbstractValue.t -> Attributes.t -> bool) -> t -> t * t
+val filter_with_discarded_attrs :
+  (AbstractValue.t -> Attributes.t -> bool) -> t -> t * Attributes.t list
 
 val find_opt : AbstractValue.t -> t -> Attributes.t option
 

infer/src/pulse/PulseOperations.ml

@@ -316,7 +316,7 @@ let mark_address_of_stack_variable history variable location address astate =
 
 
 let check_memory_leak_unreachable unreachable_attrs location astate =
-  let check_memory_leak _ attributes result =
+  let check_memory_leak result attributes =
     let allocated_not_freed_opt =
       Attributes.fold attributes ~init:(None (* allocation trace *), false (* freed *))
         ~f:(fun acc attr ->
@@ -335,7 +335,7 @@ let check_memory_leak_unreachable unreachable_attrs location astate =
     | _ ->
         result
   in
-  BaseAddressAttributes.fold check_memory_leak unreachable_attrs (Ok ())
+  List.fold unreachable_attrs ~init:(Ok ()) ~f:check_memory_leak
 
 
 let remove_vars vars location astate =