Use full annotation store in resolution fixpoint

Summary:
Store and fetch the full set of local annotations, as we used to.

iirc, the resolution fixpoint is used so we can rebuild a resolution and have typing context at each step in the fixpoint / before and after each statement.

Before this diff to make the temporary annotations change a little easier to tackle at once, this continued to only store non-temporary annotations.

Temporary annotations *are* part of the type context though, so we should both store the full annotation store (both permanent and temporary local annotations) when we save a pre/post condition, as well as read in both when constructing a resolution from a saved pre/post condition set.

Reviewed By: grievejia

Differential Revision: D30413804

fbshipit-source-id: 59139cfc1fa855a70a6654da42d119666af9095d

Author: shannonzhu

source/analysis/localAnnotationMap.ml

@@ -10,18 +10,33 @@ open Ast
 open Pyre
 
 module Annotations = struct
+  type annotation_store = {
+    annotations: RefinementUnit.t Reference.Map.Tree.t;
+    temporary_annotations: RefinementUnit.t Reference.Map.Tree.t;
+  }
+  [@@deriving eq]
+
   type t = {
-    precondition: RefinementUnit.t Reference.Map.Tree.t;
-    postcondition: RefinementUnit.t Reference.Map.Tree.t;
+    precondition: annotation_store;
+    postcondition: annotation_store;
   }
   [@@deriving eq]
 
-  let pp_state formatter map =
-    let pp_element ~key ~data =
-      Format.fprintf formatter "\"%a\": \"%a\", " Reference.pp key RefinementUnit.pp data
+  let pp_state formatter { annotations; temporary_annotations } =
+    let pp_element ~temporary ~key ~data =
+      let temporary_suffix = if temporary then "(temp)" else "" in
+      Format.fprintf
+        formatter
+        "\"%a\": \"%a\"%s, "
+        Reference.pp
+        key
+        RefinementUnit.pp
+        data
+        temporary_suffix
     in
     Format.fprintf formatter "{";
-    Reference.Map.Tree.iteri map ~f:pp_element;
+    Reference.Map.Tree.iteri annotations ~f:(pp_element ~temporary:false);
+    Reference.Map.Tree.iteri temporary_annotations ~f:(pp_element ~temporary:true);
     Format.fprintf formatter "}"
 
 
@@ -58,32 +73,47 @@ let pp formatter statements =
 let show map = Format.asprintf "%a" pp map
 
 let set
-    ?(precondition = Reference.Map.empty)
-    ?(postcondition = Reference.Map.empty)
+    ?(precondition =
+      { Resolution.annotations = Reference.Map.empty; temporary_annotations = Reference.Map.empty })
+    ?(postcondition =
+      { Resolution.annotations = Reference.Map.empty; temporary_annotations = Reference.Map.empty })
     ~key
     local_annotations
   =
+  let convert_to_tree { Resolution.annotations; temporary_annotations } =
+    {
+      Annotations.annotations = Reference.Map.to_tree annotations;
+      temporary_annotations = Reference.Map.to_tree temporary_annotations;
+    }
+  in
   Hashtbl.set
     local_annotations
     ~key
     ~data:
       {
-        Annotations.precondition = Reference.Map.to_tree precondition;
-        postcondition = Reference.Map.to_tree postcondition;
+        Annotations.precondition = convert_to_tree precondition;
+        postcondition = convert_to_tree postcondition;
       }
 
 
 module ReadOnly = struct
   type t = Annotations.t Int.Map.Tree.t
 
+  let convert_to_map { Annotations.annotations; temporary_annotations } =
+    {
+      Resolution.annotations = Reference.Map.of_tree annotations;
+      temporary_annotations = Reference.Map.of_tree temporary_annotations;
+    }
+
+
   let get_precondition local_annotations key =
     Int.Map.Tree.find local_annotations key
-    >>| fun { Annotations.precondition; _ } -> Reference.Map.of_tree precondition
+    >>| fun { Annotations.precondition; _ } -> convert_to_map precondition
 
 
   let get_postcondition local_annotations key =
     Int.Map.Tree.find local_annotations key
-    >>| fun { Annotations.postcondition; _ } -> Reference.Map.of_tree postcondition
+    >>| fun { Annotations.postcondition; _ } -> convert_to_map postcondition
 end
 
 let read_only local_annotations = Hashtbl.to_alist local_annotations |> Int.Map.Tree.of_alist_exn

source/analysis/localAnnotationMap.mli

@@ -5,8 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  *)
 
-open Ast
-
 type t
 
 val empty : unit -> t
@@ -18,18 +16,18 @@ val pp : Format.formatter -> t -> unit
 val show : t -> string
 
 val set
-  :  ?precondition:RefinementUnit.t Reference.Map.t ->
-  ?postcondition:RefinementUnit.t Reference.Map.t ->
+  :  ?precondition:Resolution.annotation_store ->
+  ?postcondition:Resolution.annotation_store ->
   key:int ->
   t ->
   unit
 
 module ReadOnly : sig
   type t
 
-  val get_precondition : t -> int -> RefinementUnit.t Reference.Map.t option
+  val get_precondition : t -> int -> Resolution.annotation_store option
 
-  val get_postcondition : t -> int -> RefinementUnit.t Reference.Map.t option
+  val get_postcondition : t -> int -> Resolution.annotation_store option
 end
 
 val read_only : t -> ReadOnly.t

source/analysis/lookup.ml

@@ -265,30 +265,22 @@ let create_of_module type_environment qualifier =
       let pre_annotations, post_annotations =
         let key = [%hash: int * int] (node_id, statement_index) in
         ( LocalAnnotationMap.ReadOnly.get_precondition annotation_lookup key
-          |> Option.value ~default:Reference.Map.empty,
+          |> Option.value ~default:Resolution.empty_annotation_store,
           LocalAnnotationMap.ReadOnly.get_postcondition annotation_lookup key
-          |> Option.value ~default:Reference.Map.empty )
+          |> Option.value ~default:Resolution.empty_annotation_store )
       in
       let pre_resolution =
         (* TODO(T65923817): Eliminate the need of creating a dummy context here *)
         TypeCheck.resolution
           global_resolution
-          ~annotation_store:
-            {
-              Resolution.annotations = pre_annotations;
-              temporary_annotations = Reference.Map.empty;
-            }
+          ~annotation_store:pre_annotations
           (module TypeCheck.DummyContext)
       in
       let post_resolution =
         (* TODO(T65923817): Eliminate the need of creating a dummy context here *)
         TypeCheck.resolution
           global_resolution
-          ~annotation_store:
-            {
-              Resolution.annotations = post_annotations;
-              temporary_annotations = Reference.Map.empty;
-            }
+          ~annotation_store:post_annotations
           (module TypeCheck.DummyContext)
       in
       Visit.visit

source/analysis/resolution.ml

@@ -61,6 +61,10 @@ let pp format { annotation_store = { annotations; temporary_annotations }; type_
 
 let show resolution = Format.asprintf "%a" pp resolution
 
+let empty_annotation_store =
+  { annotations = Reference.Map.empty; temporary_annotations = Reference.Map.empty }
+
+
 let is_global { global_resolution; _ } ~reference =
   Reference.delocalize reference |> GlobalResolution.global global_resolution |> Option.is_some
 

source/analysis/resolution.mli

@@ -30,6 +30,8 @@ val create
   unit ->
   t
 
+val empty_annotation_store : annotation_store
+
 val resolve_expression : t -> Expression.t -> t * Type.t
 
 val resolve_expression_to_type : t -> Expression.t -> Type.t

source/analysis/typeCheck.ml

@@ -1874,8 +1874,7 @@ module State (Context : Context) = struct
     let state =
       let resolution_fixpoint = LocalAnnotationMap.empty () in
       let error_map = LocalErrorMap.empty () in
-      (* TODO(T70545381): Use full annotation store in resolution fixpoint. *)
-      let postcondition = Resolution.annotations resolution in
+      let postcondition = Resolution.annotation_store resolution in
       let key = [%hash: int * int] (Cfg.entry_index, 0) in
       LocalAnnotationMap.set resolution_fixpoint ~key ~postcondition;
       LocalErrorMap.set error_map ~key ~errors;
@@ -5915,9 +5914,8 @@ module State (Context : Context) = struct
           match post_resolution with
           | None -> ()
           | Some post_resolution ->
-              (* TODO(T70545381): Use full annotation store in resolution fixpoint. *)
-              let precondition = Resolution.annotations resolution in
-              let postcondition = Resolution.annotations post_resolution in
+              let precondition = Resolution.annotation_store resolution in
+              let postcondition = Resolution.annotation_store post_resolution in
               LocalAnnotationMap.set state.resolution_fixpoint ~key ~precondition ~postcondition
         in
         { state with resolution = post_resolution }
@@ -5952,8 +5950,7 @@ end
 
 let resolution
     global_resolution
-    ?(annotation_store =
-      { Resolution.annotations = Reference.Map.empty; temporary_annotations = Reference.Map.empty })
+    ?(annotation_store = Resolution.empty_annotation_store)
     (module Context : Context)
   =
   let module State = State (Context) in
@@ -5973,15 +5970,14 @@ let resolution
 
 
 let resolution_with_key ~global_resolution ~local_annotations ~parent ~key context =
-  let annotations =
+  let annotation_store =
     Option.value_map
       local_annotations
       ~f:(fun map ->
         LocalAnnotationMap.ReadOnly.get_precondition map key
-        |> Option.value ~default:Reference.Map.empty)
-      ~default:Reference.Map.empty
+        |> Option.value ~default:Resolution.empty_annotation_store)
+      ~default:Resolution.empty_annotation_store
   in
-  let annotation_store = { Resolution.annotations; temporary_annotations = Reference.Map.empty } in
   resolution global_resolution ~annotation_store context |> Resolution.with_parent ~parent
 
 

source/interprocedural/test/dependencyGraphTest.ml

@@ -354,12 +354,8 @@ let test_type_collection context =
     let test_expect (node_id, statement_index, test_expression, expected_type) =
       let key = [%hash: int * int] (node_id, statement_index) in
       let annotation_store =
-        {
-          Resolution.annotations =
-            (LocalAnnotationMap.ReadOnly.get_precondition lookup key
-            |> fun value -> Option.value_exn value);
-          temporary_annotations = Reference.Map.empty;
-        }
+        LocalAnnotationMap.ReadOnly.get_precondition lookup key
+        |> fun value -> Option.value_exn value
       in
       let global_resolution = TypeEnvironment.ReadOnly.global_resolution environment in
       let resolution =