Cache override graph

Summary: We can finally enable caching the override graph. In this diff, we 1. Implement `save_to_cache` and `load_from_cache` for the key-value data structure defined in `saveLoadSharedMemory`. 2. Use the above APIs to enable saving / loading the override graph into / from the shared memory. 3. Attempt to reuse the override graph from the previous run if deciding it is not stale. Otherwise, remove the previous override graph and compute a new one. Reviewed By: arthaud Differential Revision: D48079764 fbshipit-source-id: 580d6755d8ea98898154d44c0453952478065957
facebook · Aug 19, 2023 · 4c9539b · 4c9539b
1 parent 3e7fd86
commit 4c9539b
Show file tree

Hide file tree

Showing 8 changed files with 193 additions and 16 deletions.
diff --git a/source/interprocedural/overrideGraph.ml b/source/interprocedural/overrideGraph.ml
@@ -199,6 +199,8 @@ module SharedMemory = struct
   (** Records a heap override graph in shared memory. *)
   let from_heap overrides = overrides |> Target.Map.Tree.to_alist |> T.of_alist
 
+  let to_heap handle = handle |> T.to_alist |> Target.Map.Tree.of_alist_exn
+
   (** Remove an override graph from shared memory. This must be called before storing another
       override graph. *)
   let cleanup = T.cleanup
@@ -216,6 +218,11 @@ module SharedMemory = struct
           :: List.fold overrides ~f:expand_and_gather ~init:expanded
     in
     List.fold callees ~init:[] ~f:expand_and_gather |> List.dedup_and_sort ~compare:Target.compare
+
+
+  let save_to_cache = T.save_to_cache
+
+  let load_from_cache = T.load_from_cache
 end
 
 let get_source ~environment qualifier =

diff --git a/source/interprocedural/overrideGraph.mli b/source/interprocedural/overrideGraph.mli
@@ -55,9 +55,15 @@ module SharedMemory : sig
   (** Record a heap override graph in shared memory and return the handle to the storage location. *)
   val from_heap : Heap.t -> t
 
+  val to_heap : t -> Heap.t
+
   (** Remove an override graph from shared memory. This must be called before storing another
       override graph. *)
   val cleanup : t -> unit
+
+  val save_to_cache : t -> unit
+
+  val load_from_cache : unit -> (t, SaveLoadSharedMemory.Usage.t) result
 end
 
 type skipped_overrides = Target.t list

diff --git a/source/interprocedural/saveLoadSharedMemory.ml b/source/interprocedural/saveLoadSharedMemory.ml
@@ -15,7 +15,10 @@ let exception_to_error ~error ~message ~f =
 
 
 module Usage = struct
-  type error = LoadError [@@deriving show { with_path = false }]
+  type error =
+    | LoadError
+    | Stale
+  [@@deriving show { with_path = false }]
 
   type t =
     | Used
@@ -124,4 +127,24 @@ struct
     let first_class_handle = FirstClass.create () in
     List.iter list ~f:(save_entry ~first_class_handle);
     { Handle.first_class_handle; keys = list |> List.map ~f:fst |> FirstClass.KeySet.of_list }
+
+
+  let to_alist { Handle.first_class_handle; keys } =
+    FirstClass.get_batch first_class_handle keys
+    |> FirstClass.KeyMap.elements
+    |> List.filter_map ~f:(fun (key, value) ->
+           match value with
+           | Some value -> Some (key, value)
+           | None -> None)
+
+
+  module HandleSharedMemory = MakeSingleValue (struct
+    type t = Handle.t
+
+    let name = Format.asprintf "Handle of %s" Value.description
+  end)
+
+  let save_to_cache = HandleSharedMemory.save
+
+  let load_from_cache = HandleSharedMemory.load
 end
diff --git a/source/interprocedural/saveLoadSharedMemory.mli b/source/interprocedural/saveLoadSharedMemory.mli
@@ -12,7 +12,10 @@ val exception_to_error
   ('b, 'a) result
 
 module Usage : sig
-  type error = LoadError [@@deriving show]
+  type error =
+    | LoadError
+    | Stale
+  [@@deriving show]
 
   type t =
     | Used
@@ -56,5 +59,11 @@ module MakeKeyValue (Key : Hack_parallel.Std.SharedMemory.KeyType) (Value : KeyV
 
   val of_alist : (Key.t * Value.t) list -> t
 
+  val to_alist : t -> (Key.t * Value.t) list
+
   val cleanup : t -> unit
+
+  val save_to_cache : t -> unit
+
+  val load_from_cache : unit -> (t, Usage.t) result
 end
diff --git a/source/interprocedural_analyses/taint/cache.ml b/source/interprocedural_analyses/taint/cache.ml
@@ -27,6 +27,7 @@ module Entry = struct
     | ClassHierarchyGraph
     | ClassIntervalGraph
     | PreviousAnalysisSetup
+    | OverrideGraph
   [@@deriving compare, show { with_path = false }]
 
   let show_pretty = function
@@ -35,6 +36,7 @@ module Entry = struct
     | ClassHierarchyGraph -> "class hierarchy graph"
     | ClassIntervalGraph -> "class interval graph"
     | PreviousAnalysisSetup -> "previous analysis setup"
+    | OverrideGraph -> "override graph"
 end
 
 module EntryStatus = struct
@@ -294,8 +296,17 @@ let save_shared_memory ~configuration =
       Ok ())
 
 
-let save ~maximum_overrides ~initial_models ~skipped_overrides { save_cache; configuration; _ } =
+let save
+    ~maximum_overrides
+    ~initial_models
+    ~skipped_overrides
+    ~override_graph_shared_memory
+    { save_cache; configuration; _ }
+  =
   if save_cache then
+    let () =
+      Interprocedural.OverrideGraph.SharedMemory.save_to_cache override_graph_shared_memory
+    in
     let () =
       PreviousAnalysisSetupSharedMemory.save_to_cache
         { AnalysisSetup.maximum_overrides; initial_models; skipped_overrides }
@@ -361,3 +372,94 @@ let initial_callables = InitialCallablesSharedMemory.load_or_compute
 let class_hierarchy_graph = ClassHierarchyGraphSharedMemory.load_or_compute
 
 let class_interval_graph = ClassIntervalGraphSharedMemory.load_or_compute
+
+module OverrideGraphSharedMemory = struct
+  let is_reusable
+      ~initial_models
+      ~maximum_overrides
+      {
+        AnalysisSetup.maximum_overrides = previous_maximum_overrides;
+        initial_models = previous_initial_models;
+        _;
+      }
+    =
+    let no_change_in_skip_overrides =
+      Ast.Reference.Set.equal
+        (Registry.skip_overrides previous_initial_models)
+        (Registry.skip_overrides initial_models)
+    in
+    let no_change_in_maximum_overrides =
+      Option.equal Int.equal maximum_overrides previous_maximum_overrides
+    in
+    no_change_in_skip_overrides && no_change_in_maximum_overrides
+
+
+  let load_or_compute_if_unloadable
+      ~initial_models
+      ~previous_analysis_setup:{ AnalysisSetup.skipped_overrides; _ }
+      ~maximum_overrides
+      entry_status
+      compute_value
+    =
+    match Interprocedural.OverrideGraph.SharedMemory.load_from_cache () with
+    | Ok override_graph_shared_memory ->
+        let override_graph_heap =
+          Interprocedural.OverrideGraph.SharedMemory.to_heap override_graph_shared_memory
+        in
+        ( {
+            Interprocedural.OverrideGraph.override_graph_heap;
+            override_graph_shared_memory;
+            skipped_overrides;
+          },
+          EntryStatus.add
+            ~name:Entry.OverrideGraph
+            ~usage:SaveLoadSharedMemory.Usage.Used
+            entry_status )
+    | Error error ->
+        ( compute_value ~initial_models ~maximum_overrides (),
+          EntryStatus.add ~name:Entry.OverrideGraph ~usage:error entry_status )
+
+
+  let remove_previous () =
+    match Interprocedural.OverrideGraph.SharedMemory.load_from_cache () with
+    | Ok override_graph_shared_memory ->
+        Log.info "Removing the previous override graph.";
+        Interprocedural.OverrideGraph.SharedMemory.cleanup override_graph_shared_memory
+    | Error _ -> Log.warning "Failed to remove the previous override graph."
+
+
+  let load_or_compute_if_stale_or_unloadable
+      ~initial_models
+      ~maximum_overrides
+      ({ status; _ } as cache)
+      compute_value
+    =
+    match status with
+    | Loaded ({ previous_analysis_setup; entry_status } as loaded) ->
+        let reusable = is_reusable ~initial_models ~maximum_overrides previous_analysis_setup in
+        if reusable then
+          let () = Log.info "Try to reuse the override graph from the previous run." in
+          let value, entry_status =
+            load_or_compute_if_unloadable
+              ~initial_models
+              ~previous_analysis_setup
+              ~maximum_overrides
+              entry_status
+              compute_value
+          in
+          let status = SharedMemoryStatus.Loaded { loaded with entry_status } in
+          value, { cache with status }
+        else
+          let () = Log.info "Override graph from the previous run is stale." in
+          let () = remove_previous () in
+          let cache =
+            set_entry_usage
+              ~entry:Entry.OverrideGraph
+              ~usage:(SaveLoadSharedMemory.Usage.Unused Stale)
+              cache
+          in
+          compute_value ~initial_models ~maximum_overrides (), cache
+    | _ -> compute_value ~initial_models ~maximum_overrides (), cache
+end
+
+let override_graph = OverrideGraphSharedMemory.load_or_compute_if_stale_or_unloadable
diff --git a/source/interprocedural_analyses/taint/cache.mli b/source/interprocedural_analyses/taint/cache.mli
@@ -21,6 +21,7 @@ val save
   :  maximum_overrides:int option ->
   initial_models:Registry.t ->
   skipped_overrides:Interprocedural.OverrideGraph.skipped_overrides ->
+  override_graph_shared_memory:Interprocedural.OverrideGraph.SharedMemory.t ->
   t ->
   unit
 
@@ -39,3 +40,13 @@ val class_interval_graph
   ClassIntervalSetGraph.Heap.t * t
 
 val metadata_to_json : t -> Yojson.Safe.t
+
+val override_graph
+  :  initial_models:Registry.t ->
+  maximum_overrides:int option ->
+  t ->
+  (initial_models:Registry.t ->
+  maximum_overrides:int option ->
+  unit ->
+  Interprocedural.OverrideGraph.whole_program_overrides) ->
+  Interprocedural.OverrideGraph.whole_program_overrides * t
diff --git a/source/interprocedural_analyses/taint/taintAnalysis.ml b/source/interprocedural_analyses/taint/taintAnalysis.ml
@@ -462,20 +462,26 @@ let run_taint_analysis
   Log.info "Computing overrides...";
   let timer = Timer.start () in
   let maximum_overrides = TaintConfiguration.maximum_overrides_to_analyze taint_configuration in
-  let {
-    Interprocedural.OverrideGraph.override_graph_heap;
-    override_graph_shared_memory;
-    skipped_overrides;
-  }
+  let ( {
+          Interprocedural.OverrideGraph.override_graph_heap;
+          override_graph_shared_memory;
+          skipped_overrides;
+        },
+        cache )
     =
-    Interprocedural.OverrideGraph.build_whole_program_overrides
-      ~static_analysis_configuration
-      ~scheduler
-      ~environment:(Analysis.TypeEnvironment.read_only environment)
-      ~include_unit_tests:false
-      ~skip_overrides:(Registry.skip_overrides initial_models)
+    Cache.override_graph
+      ~initial_models
       ~maximum_overrides
-      ~qualifiers
+      cache
+      (fun ~initial_models ~maximum_overrides () ->
+        Interprocedural.OverrideGraph.build_whole_program_overrides
+          ~static_analysis_configuration
+          ~scheduler
+          ~environment:(Analysis.TypeEnvironment.read_only environment)
+          ~include_unit_tests:false
+          ~skip_overrides:(Registry.skip_overrides initial_models)
+          ~maximum_overrides
+          ~qualifiers)
   in
   Statistics.performance ~name:"Overrides computed" ~phase_name:"Computing overrides" ~timer ();
 
@@ -543,7 +549,14 @@ let run_taint_analysis
     ~timer
     ();
 
-  let () = Cache.save ~maximum_overrides ~initial_models ~skipped_overrides cache in
+  let () =
+    Cache.save
+      ~maximum_overrides
+      ~initial_models
+      ~skipped_overrides
+      ~override_graph_shared_memory
+      cache
+  in
 
   Log.info "Purging shared memory...";
   let timer = Timer.start () in

diff --git a/stubs/integration_test/run_cache_test.py b/stubs/integration_test/run_cache_test.py
@@ -269,6 +269,7 @@ def run_test_cache_second_run(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "Used",
                     "TypeEnvironment": "Used",
                 }
             },
@@ -335,6 +336,7 @@ def run_test_changed_pysa_file(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "Used",
                     "TypeEnvironment": "Used",
                 }
             },
@@ -393,6 +395,7 @@ def run_test_changed_taint_config_file(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "Used",
                     "TypeEnvironment": "Used",
                 }
             },
@@ -457,6 +460,7 @@ def run_test_changed_models(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "Used",
                     "TypeEnvironment": "Used",
                 }
             },
@@ -617,6 +621,7 @@ def run_test_changed_overrides(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "(Unused Stale)",
                     "TypeEnvironment": "Used",
                 }
             },
@@ -657,6 +662,7 @@ def run_test_changed_overrides_cap(self) -> None:
                     "ClassIntervalGraph": "Used",
                     "InitialCallables": "Used",
                     "PreviousAnalysisSetup": "Used",
+                    "OverrideGraph": "(Unused Stale)",
                     "TypeEnvironment": "Used",
                 }
             },