generalizing get_trace to get_node

Reviewed By: mbouaziz

Differential Revision: D3614509

fbshipit-source-id: b1dc5dd

infer/src/checkers/accessTree.ml

@@ -76,9 +76,9 @@ module Make (TraceDomain : AbstractDomain.S) = struct
     | Star ->
         orig_trace
 
-  (** retrieve the trace associated with [ap] from [tree] *)
+  (** retrieve the trace and subtree associated with [ap] from [tree] *)
-  let get_trace ap tree =
+  let get_node ap tree =
-    let rec accesses_get_trace access_list trace tree =
+    let rec accesses_get_node access_list trace tree =
       match access_list, tree with
       | _, Star ->
           trace, Star
@@ -86,22 +86,26 @@ module Make (TraceDomain : AbstractDomain.S) = struct
           trace, tree
       | access :: accesses, Subtree subtree ->
           let access_trace, access_subtree = AccessMap.find access subtree in
-          accesses_get_trace accesses access_trace access_subtree in
+          accesses_get_node accesses access_trace access_subtree in
-    let get_trace_ base accesses tree =
+    let get_node_ base accesses tree =
       let base_trace, base_tree = BaseMap.find base tree in
-      accesses_get_trace accesses base_trace base_tree in
+      accesses_get_node accesses base_trace base_tree in
     let base, accesses = AccessPath.extract ap in
-    match get_trace_ base accesses tree with
+    match get_node_ base accesses tree with
     | trace, subtree ->
         if AccessPath.is_exact ap
-        then Some trace
+        then Some (trace, subtree)
         else
           (* input query was [ap]*, and [trace] is the trace associated with [ap]. get the traces
              associated with the children of [ap] in [tree] and join them with [trace] *)
-          Some (join_all_traces trace subtree)
+          Some (join_all_traces trace subtree, subtree)
     | exception Not_found ->
         None
 
+  (** retrieve the trace associated with [ap] from [tree] *)
+  let get_trace ap tree =
+    Option.map fst (get_node ap tree)
+
   (* helper for [add_access]. [last_trace] is the trace associated with [tree] in the parent. *)
   let rec access_tree_add_trace ~last_trace ~seen_array_access ~is_exact trace access_list tree =
     match access_list, tree with
@@ -272,11 +276,12 @@ module Make (TraceDomain : AbstractDomain.S) = struct
             prev_node_opt in
       BaseMap.merge (fun _ prev_node next_node -> node_widen prev_node next_node) prev next
 
-  let pp fmt base_tree =
   let rec pp_node fmt (trace, subtree) =
     let pp_subtree fmt = function
       | Subtree access_map -> AccessMap.pp ~pp_value:pp_node fmt access_map
       | Star -> F.fprintf fmt "*" in
-      F.fprintf fmt "(%a, %a)" TraceDomain.pp trace pp_subtree subtree in
+    F.fprintf fmt "(%a, %a)" TraceDomain.pp trace pp_subtree subtree
+
+  let pp fmt base_tree =
     BaseMap.pp ~pp_value:pp_node fmt base_tree
 end

infer/src/unit/accessTreeTests.ml

@@ -104,12 +104,12 @@ let tests =
   let array_f_trace = MockTraceDomain.singleton "arrayF" in
   let x_star_trace = MockTraceDomain.of_list ["x"; "xF"; "xFG"] in
 
-  let x_subtree =
   let g_subtree = Domain.make_access_node xF_trace g xFG_trace in
+  let x_subtree =
     Domain.AccessMap.singleton f g_subtree
     |> Domain.make_node x_trace in
-  let y_subtree =
   let yF_subtree = Domain.make_starred_leaf yF_trace in
+  let y_subtree =
     Domain.AccessMap.singleton f yF_subtree
     |> Domain.make_node y_trace in
   let z_subtree = Domain.make_starred_leaf z_trace in
@@ -146,6 +146,18 @@ let tests =
   let assert_trace_not_found access_path tree =
     assert_traces_eq access_path tree no_trace in
 
+  let assert_node_equal access_path tree expected_node =
+    match Domain.get_node access_path tree with
+    | Some actual_node ->
+        let pp_diff fmt (actual, expected) =
+          F.fprintf
+            fmt
+            "Expected to retrieve node %a but got %a"
+            Domain.pp_node expected
+            Domain.pp_node actual in
+        assert_equal ~pp_diff expected_node actual_node
+    | None -> assert false in
+
   let get_trace_test =
     let get_trace_test_ _ =
       (* exact access path tests *)
@@ -168,7 +180,23 @@ let tests =
       assert_traces_eq y_star tree "{ y, yF }";
       assert_traces_eq yF_star tree "{ yF }";
       assert_traces_eq yFG tree "{ yF }";
-      assert_trace_not_found yG tree in
+      assert_trace_not_found yG tree;
+
+      (* get_trace is just (fst get_node), so light tests here *)
+      (* exact access path tests *)
+      assert_node_equal z tree z_subtree;
+      assert_node_equal xF tree g_subtree;
+      assert_node_equal xFG tree (Domain.make_normal_leaf xFG_trace);
+
+      (* starred tree tests *)
+      assert_node_equal yFG tree yF_subtree;
+
+      (* starred access path tests *)
+      let joined_y_subtree =
+        Domain.AccessMap.singleton f yF_subtree
+        |> Domain.make_node (MockTraceDomain.join y_trace yF_trace) in
+      assert_node_equal y_star tree joined_y_subtree in
+
     "get_trace">::get_trace_test_ in
 
   let add_trace_test =