Factor out the function that computes features to attach

Summary: We currently have two functions that extract features to attach, let's factor them out.

Reviewed By: dkgi

Differential Revision: D30758332

fbshipit-source-id: ef81a82e75ba4111573a59460d32f31d1ae18338

Author: arthaud

source/interprocedural_analyses/taint/backwardAnalysis.ml

@@ -1429,8 +1429,9 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
     in
     let taint_in_taint_out =
       let features_to_attach =
-        BackwardState.compute_features_to_attach
+        BackwardState.extract_features_to_attach
           ~root:parameter
+          ~attach_to_leaf:Sinks.Attach
           existing_backward.TaintResult.Backward.taint_in_taint_out
       in
       let candidate_tree =
@@ -1490,8 +1491,9 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
     in
     let sink_taint =
       let features_to_attach =
-        BackwardState.compute_features_to_attach
+        BackwardState.extract_features_to_attach
           ~root:parameter
+          ~attach_to_leaf:Sinks.Attach
           existing_backward.TaintResult.Backward.sink_taint
       in
       if not (Features.SimpleSet.is_bottom features_to_attach) then

source/interprocedural_analyses/taint/domains.ml

@@ -705,30 +705,25 @@ module MakeTaintEnvironment (Taint : TAINT_DOMAIN) () = struct
   let is_empty = is_bottom
 
   let roots environment = fold Key ~f:List.cons ~init:[] environment
-end
-
-module ForwardState = MakeTaintEnvironment (ForwardTaint) ()
-(** Used to infer which sources reach the exit points of a function. *)
-
-(** Used to infer which sinks are reached from parameters, as well as the taint-in-taint-out (TITO)
-    using the special LocalReturn sink. *)
-module BackwardState = struct
-  include MakeTaintEnvironment (BackwardTaint) ()
 
-  let compute_features_to_attach ~root taint =
+  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.collapse ~transform:Fn.id
-    |> BackwardTaint.partition BackwardTaint.leaf ByFilter ~f:(fun sink ->
-           if Sinks.equal Sinks.Attach sink then Some true else None)
+    |> 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)
-    >>| BackwardTaint.fold
-          BackwardTaint.simple_feature_self
-          ~f:gather_features
-          ~init:Features.SimpleSet.bottom
+    >>| Taint.fold Taint.simple_feature_self ~f:gather_features ~init:Features.SimpleSet.bottom
     |> Option.value ~default:Features.SimpleSet.bottom
 end
 
+module ForwardState = MakeTaintEnvironment (ForwardTaint) ()
+(** Used to infer which sources reach the exit points of a function. *)
+
+module BackwardState = MakeTaintEnvironment (BackwardTaint) ()
+(** Used to infer which sinks are reached from parameters, as well as the taint-in-taint-out (TITO)
+    using the special LocalReturn sink. *)
+
 (* Special sink as it needs the return access path *)
 let local_return_taint =
   BackwardTaint.create

source/interprocedural_analyses/taint/forwardAnalysis.ml

@@ -1543,23 +1543,6 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
     Fixpoint.Make (FixpointState)
 end
 
-let extract_features_to_attach existing_taint =
-  ForwardState.read ~root:AccessPath.Root.LocalResult ~path:[] existing_taint
-  |> ForwardState.Tree.collapse ~transform:Fn.id
-  |> ForwardTaint.partition ForwardTaint.leaf ByFilter ~f:(fun source ->
-         if Sources.equal Sources.Attach source then Some true else None)
-  |> (fun map -> Map.Poly.find map true)
-  |> function
-  | Some taint ->
-      let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
-      ForwardTaint.fold
-        ForwardTaint.simple_feature_self
-        ~f:gather_features
-        ~init:Features.SimpleSet.bottom
-        taint
-  | None -> Features.SimpleSet.bottom
-
-
 let extract_source_model ~define ~resolution ~features_to_attach exit_taint =
   let {
     Statement.Define.signature =
@@ -1758,8 +1741,9 @@ let run ~environment ~qualifier ~define ~call_graph_of_define ~existing_model =
              ~port:AccessPath.Root.LocalResult
       in
       let features_to_attach =
-        BackwardState.compute_features_to_attach
+        BackwardState.extract_features_to_attach
           ~root:AccessPath.Root.LocalResult
+          ~attach_to_leaf:Sinks.Attach
           existing_model.TaintResult.backward.sink_taint
       in
       if not (Features.SimpleSet.is_bottom features_to_attach) then
@@ -1805,7 +1789,12 @@ let run ~environment ~qualifier ~define ~call_graph_of_define ~existing_model =
   in
   let resolution = TypeEnvironment.ReadOnly.global_resolution environment in
   let extract_model { FixpointState.taint; _ } =
-    let features_to_attach = extract_features_to_attach existing_model.forward.source_taint in
+    let features_to_attach =
+      ForwardState.extract_features_to_attach
+        ~root:AccessPath.Root.LocalResult
+        ~attach_to_leaf:Sources.Attach
+        existing_model.forward.source_taint
+    in
     let source_taint =
       extract_source_model ~define:define.value ~resolution ~features_to_attach taint
     in