fix(engine): make dep. analysis sensible to UID assoc. items

engine/bin/lib.ml

@@ -1,6 +1,7 @@
 open Hax_engine
 open Base
 open Stdio
+open Utils
 
 let setup_logs (options : Types.engine_options) =
   let level : Logs.level option =
@@ -47,6 +48,9 @@ let import_thir_items (include_clauses : Types.inclusion_clause list)
     |> List.map ~f:snd
   in
   let items = List.concat_map ~f:fst imported_items in
+  let associated_items (item : Deps.AST.item) =
+    Deps.uid_associated_items items item.attrs
+  in
   (* Build a map from idents to error reports *)
   let ident_to_reports =
     List.concat_map
@@ -56,6 +60,11 @@ let import_thir_items (include_clauses : Types.inclusion_clause list)
     |> Map.of_alist_exn (module Concrete_ident)
   in
   let items = Deps.filter_by_inclusion_clauses include_clauses items in
+  let items =
+    items
+    @ (List.concat_map ~f:associated_items items
+      |> List.filter ~f:(List.mem ~equal:[%eq: Deps.AST.item] items >> not))
+  in
   let items =
     List.filter
       ~f:(fun i ->

engine/lib/dependencies.ml

@@ -26,6 +26,18 @@ module Make (F : Features.T) = struct
       krate :: path
   end
 
+  module Error : Phase_utils.ERROR = Phase_utils.MakeError (struct
+    let ctx = Diagnostics.Context.Dependencies
+  end)
+
+  module Attrs = Attr_payloads.Make (F) (Error)
+
+  let uid_associated_items (items : item list) : attrs -> item list =
+    let open Attrs.WithItems (struct
+      let items = items
+    end) in
+    raw_associated_item >> List.map ~f:(snd >> item_of_uid)
+
   module ItemGraph = struct
     module G = Graph.Persistent.Digraph.Concrete (Concrete_ident)
     module Topological = Graph.Topological.Make_stable (G)
@@ -71,15 +83,22 @@ module Make (F : Features.T) = struct
       in
       set |> Set.to_list
 
-    let vertices_of_items : item list -> G.E.t list =
-      List.concat_map ~f:(fun i ->
-          vertices_of_item i |> List.map ~f:(Fn.const i.ident &&& Fn.id))
-
-    let of_items (items : item list) : G.t =
+    let vertices_of_items ~original_items (items : item list) : G.E.t list =
+      let uid_associated_items = uid_associated_items original_items in
+      List.concat_map
+        ~f:(fun i ->
+          let assoc =
+            uid_associated_items i.attrs |> List.map ~f:(fun i -> i.ident)
+          in
+          vertices_of_item i @ assoc |> List.map ~f:(Fn.const i.ident &&& Fn.id))
+        items
+
+    let of_items ~original_items (items : item list) : G.t =
       let init =
         List.fold ~init:G.empty ~f:(fun g -> ident_of >> G.add_vertex g) items
       in
-      vertices_of_items items |> List.fold ~init ~f:(G.add_edge >> uncurry)
+      vertices_of_items ~original_items items
+      |> List.fold ~init ~f:(G.add_edge >> uncurry)
 
     let transitive_dependencies_of (g : G.t) (selection : Concrete_ident.t list)
         : Concrete_ident.t Hash_set.t =
@@ -92,9 +111,9 @@ module Make (F : Features.T) = struct
       List.filter ~f:(G.mem_vertex g) selection |> List.iter ~f:visit;
       set
 
-    let transitive_dependencies_of_items (items : item list)
-        ?(graph = of_items items) (selection : Concrete_ident.t list) :
-        item list =
+    let transitive_dependencies_of_items ~original_items (items : item list)
+        ?(graph = of_items ~original_items items)
+        (selection : Concrete_ident.t list) : item list =
       let set = transitive_dependencies_of graph selection in
       items |> List.filter ~f:(ident_of >> Hash_set.mem set)
 
@@ -154,14 +173,14 @@ module Make (F : Features.T) = struct
       let edge_attributes _ = []
     end)
 
-    let print oc items = output_graph oc (of_items items)
+    let print oc items = output_graph oc (of_items ~original_items:items items)
   end
 
   module ModGraph = struct
     module G = Graph.Persistent.Digraph.Concrete (Namespace)
 
     let of_items (items : item list) : G.t =
-      let ig = ItemGraph.of_items items in
+      let ig = ItemGraph.of_items ~original_items:items items in
       assert (ItemGraph.MutRec.all_homogeneous_namespace ig);
       List.map ~f:(ident_of >> (Namespace.of_concrete_ident &&& Fn.id)) items
       |> Map.of_alist_multi (module Namespace)
@@ -215,7 +234,9 @@ module Make (F : Features.T) = struct
     List.map ~f:Concrete_ident.DefaultViewAPI.show >> String.concat ~sep:", "
 
   let sort (items : item list) : item list =
-    let g = ItemGraph.of_items items |> ItemGraph.Oper.mirror in
+    let g =
+      ItemGraph.of_items ~original_items:items items |> ItemGraph.Oper.mirror
+    in
     let lookup (name : concrete_ident) =
       List.find ~f:(ident_of >> Concrete_ident.equal name) items
     in
@@ -232,9 +253,10 @@ module Make (F : Features.T) = struct
       Set.equal items items');
     items'
 
-  let filter_by_inclusion_clauses' (clauses : Types.inclusion_clause list)
-      (items : item list) : item list * Concrete_ident.t Hash_set.t =
-    let graph = ItemGraph.of_items items in
+  let filter_by_inclusion_clauses' ~original_items
+      (clauses : Types.inclusion_clause list) (items : item list) :
+      item list * Concrete_ident.t Hash_set.t =
+    let graph = ItemGraph.of_items ~original_items items in
     let of_list = Set.of_list (module Concrete_ident) in
     let selection = List.map ~f:ident_of items |> of_list in
     let deps_of =
@@ -300,7 +322,7 @@ module Make (F : Features.T) = struct
 
   let filter_by_inclusion_clauses (clauses : Types.inclusion_clause list)
       (items : item list) : item list =
-    let f = filter_by_inclusion_clauses' clauses in
+    let f = filter_by_inclusion_clauses' ~original_items:items clauses in
     let selection =
       let items', items_drop_body = f items in
       let items', _ =
@@ -325,7 +347,7 @@ module Make (F : Features.T) = struct
     ({ item with v = Alias { name = item.ident; item = item'.ident } }, item')
 
   let name_me' (items : item list) : item list =
-    let g = ItemGraph.of_items items in
+    let g = ItemGraph.of_items ~original_items:items items in
     let from_ident ident : item option =
       List.find ~f:(fun i -> [%equal: Concrete_ident.t] i.ident ident) items
     in

engine/lib/dependencies.mli

@@ -1,6 +1,7 @@
 module Make (F : Features.T) : sig
   module AST : module type of Ast.Make (F)
 
+  val uid_associated_items : AST.item list -> Ast.attrs -> AST.item list
   val name_me : AST.item list -> AST.item list
   val sort : AST.item list -> AST.item list
 

engine/lib/diagnostics.ml

@@ -69,6 +69,7 @@ module Context = struct
     | Phase of Phase.t
     | Backend of Backend.t
     | ThirImport
+    | Dependencies
     | DebugPrintRust
     | GenericPrinter of string
     | Other of string
@@ -79,6 +80,7 @@ module Context = struct
     | Backend backend -> [%show: Backend.t] backend ^ " backend"
     | ThirImport -> "AST import"
     | DebugPrintRust -> "Rust debug printer"
+    | Dependencies -> "Dependenciy analysis"
     | GenericPrinter kind -> kind ^ " generic printer"
     | Other s -> "Other (" ^ s ^ ")"
 end

engine/lib/phases/phase_newtype_as_refinement.ml

@@ -14,7 +14,7 @@ module Make (F : Features.T) =
         let ctx = Diagnostics.Context.Phase phase_id
       end)
 
-      module Attrs = Attr_payloads.MakeBase (Error)
+      module Attrs = Attr_payloads.Make (F) (Error)
 
       let visitor =
         object

test-harness/src/snapshots/toolchain__attributes into-fstar.snap

@@ -32,21 +32,21 @@ module Attributes.Ensures_on_arity_zero_fns
 open Core
 open FStar.Mul
 
-let basically_a_constant (_: Prims.unit)
-    : Prims.Pure u8
-      (requires true)
-      (ensures
-        fun x ->
-          let x:u8 = x in
-          x >. 100uy) = 127uy
-
 let doing_nothing (_: Prims.unit)
     : Prims.Pure Prims.unit
       (requires true)
       (ensures
         fun v__x ->
           let v__x:Prims.unit = v__x in
           true) = ()
+
+let basically_a_constant (_: Prims.unit)
+    : Prims.Pure u8
+      (requires true)
+      (ensures
+        fun x ->
+          let x:u8 = x in
+          x >. 100uy) = 127uy
 '''
 "Attributes.Inlined_code_ensures_requires.fst" = '''
 module Attributes.Inlined_code_ensures_requires
@@ -92,15 +92,15 @@ open FStar.Mul
 
 unfold type t_Int = int
 
-unfold let add x y = x + y
-
 unfold instance impl: Core.Ops.Arith.t_Sub t_Int t_Int =
   {
     f_Output = t_Int;
     f_sub_pre = (fun (self: t_Int) (other: t_Int) -> true);
     f_sub_post = (fun (self: t_Int) (other: t_Int) (out: t_Int) -> true);
     f_sub = fun (self: t_Int) (other: t_Int) -> self + other
   }
+
+unfold let add x y = x + y
 '''
 "Attributes.Nested_refinement_elim.fst" = '''
 module Attributes.Nested_refinement_elim
@@ -454,12 +454,6 @@ module Attributes
 open Core
 open FStar.Mul
 
-let add3_lemma (x: u32)
-    : Lemma Prims.l_True
-      (ensures
-        x <=. 10ul || x >=. (u32_max /! 3ul <: u32) || (add3 x x x <: u32) =. (x *! 3ul <: u32)) =
-  ()
-
 let inlined_code__V: u8 = 12uy
 
 let issue_844_ (v__x: u8)
@@ -472,7 +466,18 @@ let issue_844_ (v__x: u8)
 
 let u32_max: u32 = 90000ul
 
-unfold let some_function _ = "hello from F*"
+let swap_and_mut_req_ens (x y: u32)
+    : Prims.Pure (u32 & u32 & u32)
+      (requires x <. 40ul && y <. 300ul)
+      (ensures
+        fun temp_0_ ->
+          let x_future, y_future, result:(u32 & u32 & u32) = temp_0_ in
+          x_future =. y && y_future =. x && result =. (x +! y <: u32)) =
+  let x0:u32 = x in
+  let x:u32 = y in
+  let y:u32 = x0 in
+  let hax_temp_output:u32 = x +! y in
+  x, y, hax_temp_output <: (u32 & u32 & u32)
 
 /// A doc comment on `add3`
 ///another doc comment on add3
@@ -487,25 +492,20 @@ let add3 (x y z: u32)
                 let _:Prims.unit = temp_0_ in
                 result >. 32ul <: bool)) = (x +! y <: u32) +! z
 
-let swap_and_mut_req_ens (x y: u32)
-    : Prims.Pure (u32 & u32 & u32)
-      (requires x <. 40ul && y <. 300ul)
+let add3_lemma (x: u32)
+    : Lemma Prims.l_True
       (ensures
-        fun temp_0_ ->
-          let x_future, y_future, result:(u32 & u32 & u32) = temp_0_ in
-          x_future =. y && y_future =. x && result =. (x +! y <: u32)) =
-  let x0:u32 = x in
-  let x:u32 = y in
-  let y:u32 = x0 in
-  let hax_temp_output:u32 = x +! y in
-  x, y, hax_temp_output <: (u32 & u32 & u32)
+        x <=. 10ul || x >=. (u32_max /! 3ul <: u32) || (add3 x x x <: u32) =. (x *! 3ul <: u32)) =
+  ()
 
 type t_Foo = {
   f_x:u32;
   f_y:f_y: u32{f_y >. 3ul};
   f_z:f_z: u32{((f_y +! f_x <: u32) +! f_z <: u32) >. 3ul}
 }
 
+unfold let some_function _ = "hello from F*"
+
 let before_inlined_code = "example before"
 
 let inlined_code (foo: t_Foo) : Prims.unit =