From 034d2e3c81e9c076a23439d44c29a8eff9e6231b Mon Sep 17 00:00:00 2001
From: Sam Blackshear <shb@fb.com>
Date: Tue, 9 Aug 2016 17:27:10 -0700
Subject: [PATCH] adding accessTree.fold

Reviewed By: mbouaziz

Differential Revision: D3692539

fbshipit-source-id: 99fbe61
---
 infer/src/checkers/accessTree.ml  | 20 ++++++++++++++++++++
 infer/src/unit/accessTreeTests.ml | 27 ++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/infer/src/checkers/accessTree.ml b/infer/src/checkers/accessTree.ml
index d09467b30..5c9136015 100644
--- a/infer/src/checkers/accessTree.ml
+++ b/infer/src/checkers/accessTree.ml
@@ -238,6 +238,26 @@ module Make (TraceDomain : AbstractDomain.S) = struct
     then tree1
     else BaseMap.merge (fun _ n1 n2 -> node_merge n1 n2) tree1 tree2
 
+  let rec access_map_fold_ f base accesses m acc =
+    AccessMap.fold (fun access node acc -> node_fold_ f base (accesses @ [access]) node acc) m acc
+  and node_fold_ f base accesses (trace, tree) acc =
+    let cur_ap_raw = base, accesses in
+    match tree with
+    | Subtree access_map ->
+        let acc' = f acc (AccessPath.Exact cur_ap_raw) trace in
+        access_map_fold_ f base accesses access_map acc'
+    | Star ->
+        f acc (AccessPath.Abstracted cur_ap_raw) trace
+
+  let node_fold (f : 'a -> AccessPath.t -> TraceDomain.astate -> 'a) base node acc =
+    node_fold_ f base [] node acc
+
+  let access_map_fold (f : 'a -> AccessPath.t -> TraceDomain.astate -> 'a) base m acc =
+    access_map_fold_ f base [] m acc
+
+  let fold (f : 'a ->  AccessPath.t -> TraceDomain.astate -> 'a) tree acc_ =
+    BaseMap.fold (fun base node acc -> node_fold f base node acc) tree acc_
+
   (* replace the normal leaves of [node] with starred leaves *)
   let rec node_add_stars ((trace, tree) as node) = match tree with
     | Subtree subtree ->
diff --git a/infer/src/unit/accessTreeTests.ml b/infer/src/unit/accessTreeTests.ml
index 7b4746433..928725318 100644
--- a/infer/src/unit/accessTreeTests.ml
+++ b/infer/src/unit/accessTreeTests.ml
@@ -76,6 +76,7 @@ let tests =
   let xF = AccessPath.Exact (make_access_path "x" ["f"]) in
   let xG = AccessPath.Exact (make_access_path "x" ["g"]) in
   let xFG = AccessPath.Exact (make_access_path "x" ["f"; "g"]) in
+  let y = AccessPath.Exact (make_access_path "y" []) in
   let yF = AccessPath.Exact (make_access_path "y" ["f"]) in
   let yG = AccessPath.Exact (make_access_path "y" ["g"]) in
   let yFG = AccessPath.Exact (make_access_path "y" ["f"; "g"]) in
@@ -434,4 +435,28 @@ let tests =
       assert_trees_equal (widen xFG_tree x_tree_y_trace) xFG_tree_widened_trace in
     "widen">::widen_test_ in
 
-  "access_tree_suite">:::[get_trace_test; add_trace_test; lteq_test; join_test; widen_test]
+  let fold_test =
+    let fold_test_ _ =
+      let collect_ap_traces acc ap trace =
+        (ap, trace) :: acc in
+      let ap_traces = Domain.fold collect_ap_traces tree [] in
+      let has_ap_trace_pair ap_in trace_in =
+        IList.exists
+          (fun (ap, trace) -> AccessPath.equal ap ap_in && MockTraceDomain.equal trace trace_in)
+          ap_traces in
+
+      assert_bool "Should have six ap/trace pairs" (IList.length ap_traces = 6);
+      assert_bool "has x pair" (has_ap_trace_pair x x_trace);
+      assert_bool "has xF pair" (has_ap_trace_pair xF xF_trace);
+      assert_bool "has xFG pair" (has_ap_trace_pair xFG xFG_trace);
+      assert_bool "has y pair" (has_ap_trace_pair y y_trace);
+      assert_bool "has yF* pair" (has_ap_trace_pair yF_star yF_trace);
+      assert_bool "has z pair" (has_ap_trace_pair z_star z_trace) in
+    "fold">::fold_test_ in
+
+  "access_tree_suite">:::[get_trace_test;
+                          add_trace_test;
+                          lteq_test;
+                          join_test;
+                          widen_test;
+                          fold_test;]