Filter by attach to leaves before collapsing

Summary: The function that computes attach to features is pretty slow since it collapses all the taint before filtering by attach to kinds. Let's filter before collapsing.

Reviewed By: dkgi

Differential Revision: D30758356

fbshipit-source-id: c72fadb219b8c61076900d92f8bd9f9651a34187

Author: arthaud

source/interprocedural_analyses/taint/domains.ml

@@ -7,7 +7,6 @@
 
 open Core
 open Ast
-open Pyre
 
 module type SET_ARG = sig
   include Abstract.SetDomain.ELEMENT
@@ -709,12 +708,9 @@ module MakeTaintEnvironment (Taint : TAINT_DOMAIN) () = struct
   let extract_features_to_attach ~root ~attach_to_leaf taint =
     let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
     read ~root ~path:[] taint
+    |> Tree.transform Taint.leaf Filter ~f:(Taint.equal_leaf attach_to_leaf)
     |> Tree.collapse ~transform:Fn.id
-    |> Taint.partition Taint.leaf ByFilter ~f:(fun leaf ->
-           if Taint.equal_leaf attach_to_leaf leaf then Some true else None)
-    |> (fun map -> Map.Poly.find map true)
-    >>| Taint.fold Taint.simple_feature_self ~f:gather_features ~init:Features.SimpleSet.bottom
-    |> Option.value ~default:Features.SimpleSet.bottom
+    |> Taint.fold Taint.simple_feature_self ~f:gather_features ~init:Features.SimpleSet.bottom
 end
 
 module ForwardState = MakeTaintEnvironment (ForwardTaint) ()