Try to keep spliced-in functions in place

... rather than moving them to the end (or, after top-sorting,
immediately before the first use).

Motivated by trying to preserve the file structure of the Arm 9.4
specification when generating Lem definitions, where without this patch
some spliced-in functions ended up in the wrong files, causing
gratuitous dependencies between files.

This does require a topological sorting after splicing, but the prover
backends have to perform a topological sorting anyway, and patching in
more prover-friendly versions of functions currently seems to be the
main use case of the splicing feature.

This also requires a few tweaks to the topological sorting code to make
it usable early on in the translation pipeline, before other rewrites.

src/bin/sail.ml

@@ -437,7 +437,7 @@ let run_sail (config : Yojson.Basic.t option) tgt =
           )
   in
   let ast, env = Frontend.initial_rewrite effect_info env ast in
-  let ast, env = List.fold_right (fun file (ast, _) -> Splice.splice ast file) !opt_splice (ast, env) in
+  let ast, env = match !opt_splice with [] -> (ast, env) | files -> Splice.splice_files ast (List.rev files) in
   let effect_info = Effects.infer_side_effects (Target.asserts_termination tgt) ast in
 
   (* Don't show warnings during re-writing for now *)

src/lib/callgraph.ml

@@ -222,6 +222,24 @@ let add_def_to_graph graph (DEF_aux (def, _)) =
           | Typ_aux ((Typ_id id | Typ_app (id, _)), _) -> graph := G.add_edge self (Type id) !graph
           | _ -> Reporting.unreachable (fst annot) __POS__ "Struct without struct type"
         end
+      | E_lit (L_aux (L_undef, l)) -> begin
+          (* Make undefined literals depend on the undefined functions generated for the type (if any)
+             to ensure that `rewrite_undefined` works *)
+          try
+            let typ = Env.expand_synonyms env (typ_of_annot annot) in
+            let funcalls_of_exp =
+              let e_app (id, args) =
+                let arg_funcalls = List.fold_left IdSet.union IdSet.empty args in
+                if Bindings.mem id (Env.get_val_specs env) then IdSet.add id arg_funcalls else arg_funcalls
+              in
+              fold_exp { (pure_exp_alg IdSet.empty IdSet.union) with e_app }
+            in
+            (* The `mwords` parameter of `undefined_of_type` shouldn't change the set of functions called,
+               so just pass `true`. *)
+            funcalls_of_exp (undefined_of_typ true l (fun _ -> empty_uannot) typ)
+            |> IdSet.iter (fun f -> graph := G.add_edge self (Function f) !graph)
+          with _ -> ()
+        end
       | _ -> ()
     end;
     E_aux (e_aux, annot)
@@ -337,10 +355,11 @@ let add_def_to_graph graph (DEF_aux (def, _)) =
         IdSet.iter (fun id -> graph := G.add_edges (Letbind id) [] !graph) ids;
         IdSet.iter (fun id -> ignore (rewrite_let (rewriters (Letbind id)) lb)) ids
     | DEF_type tdef -> add_type_def_to_graph tdef
-    | DEF_register (DEC_aux (DEC_reg (typ, id, opt_exp), _)) ->
-        begin
-          match opt_exp with Some exp -> ignore (fold_exp (rw_exp (Register id)) exp) | None -> ()
-        end;
+    | DEF_register (DEC_aux (DEC_reg (typ, id, opt_exp), annot)) ->
+        (* Determine dependencies of initial expressions (or `undefined` if missing, which will add
+           dependencies to `undefined_*` functions) *)
+        let exp = match opt_exp with Some exp -> exp | None -> E_aux (E_lit (mk_lit L_undef), annot) in
+        ignore (fold_exp (rw_exp (Register id)) exp);
         IdSet.iter (fun typ_id -> graph := G.add_edge (Register id) (Type typ_id) !graph) (typ_ids typ)
     | DEF_measure (id, pat, exp) ->
         graph := G.add_edges (FunctionMeasure id) [Function id] !graph;
@@ -353,7 +372,11 @@ let add_def_to_graph graph (DEF_aux (def, _)) =
         graph := G.add_edges (Outcome id) [] !graph;
         scan_typquant (Outcome id) typq;
         IdSet.iter (fun typ_id -> graph := G.add_edge (Outcome id) (Type typ_id) !graph) (typ_ids typ);
-        List.iter (scan_outcome_def l (Outcome id)) outcome_defs
+        List.iter (scan_outcome_def l (Outcome id)) outcome_defs;
+        (* Remove dependencies on functions declared within the outcome;  these are parameters of the outcome,
+           and instantiations of these functions (possibly with a more constrained, architecture-specific type)
+           and an `instantiation` declaration normally come later *)
+        IdSet.iter (fun f -> graph := G.delete_edge (Outcome id) (Function f) !graph) (val_spec_ids outcome_defs)
     | DEF_instantiation (IN_aux (IN_id id, _), substs) ->
         graph := G.add_edges (Function id) [Outcome id] !graph;
         List.iter

src/lib/effects.ml

@@ -363,13 +363,21 @@ let infer_side_effects asserts_termination ast =
   let mapping_effects = ref Bindings.empty in
 
   let all_nodes = Callgraph.G.nodes cg in
+  let register_nodes =
+    List.filter (function Callgraph.Register _ -> true | _ -> false) all_nodes |> NodeSet.of_list
+  in
   let total = List.length all_nodes in
   List.iteri
     (fun i node ->
       Util.progress "Effects (transitive) " (string_of_int (i + 1) ^ "/" ^ string_of_int total) (i + 1) total;
       match node with
       | Callgraph.Function id | Callgraph.Letbind id | Callgraph.Mapping id ->
-          let reachable = Callgraph.G.reachable (NodeSet.singleton node) NodeSet.empty cg in
+          (* Compute dependencies of this node, i.e., the reachable nodes in the callgraph, but cut at registers;
+             we don't want to transitively include the dependencies of registers here, otherwise a let-binding
+             using a reference to a register, for example, would inherit the effects of all functions used for
+             initialising the register.  The register nodes themselves will still be included in `reachable`,
+             however, so the effects of directly accessing registers will be included in `side_effects` below. *)
+          let reachable = Callgraph.G.reachable (NodeSet.singleton node) register_nodes cg in
           (* First, a function has any side effects it directly causes *)
           let side_effects =
             match Bindings.find_opt id !direct_effects with Some effs -> effs | None -> EffectSet.empty

src/lib/graph.ml

@@ -85,6 +85,8 @@ module type S = sig
 
   val has_edge : node -> node -> graph -> bool
 
+  val delete_edge : node -> node -> graph -> graph
+
   val add_edges : node -> node list -> graph -> graph
 
   val children : graph -> node -> node list

src/lib/graph.mli

@@ -87,6 +87,8 @@ module type S = sig
 
   val has_edge : node -> node -> graph -> bool
 
+  val delete_edge : node -> node -> graph -> graph
+
   val add_edges : node -> node list -> graph -> graph
 
   val children : graph -> node -> node list

src/lib/splice.ml

@@ -85,20 +85,22 @@ let scan_ast { defs; _ } =
   List.fold_left scan (IdSet.empty, Bindings.empty) defs
 
 let filter_old_ast repl_ids repl_specs { defs; _ } =
-  let check (rdefs, spec_found) (DEF_aux (aux, def_annot) as def) =
+  let check (rdefs, specs_found) (DEF_aux (aux, def_annot) as def) =
     match aux with
     | DEF_fundef fd ->
         let id = id_of_fundef fd in
-        if IdSet.mem id repl_ids then (rdefs, spec_found) else (def :: rdefs, spec_found)
+        if IdSet.mem id repl_ids then
+          (DEF_aux (DEF_pragma ("spliced_function#", string_of_id id, def_annot.loc), def_annot) :: rdefs, specs_found)
+        else (def :: rdefs, specs_found)
     | DEF_val (VS_aux (VS_val_spec (_, id, _), _)) -> (
         match Bindings.find_opt id repl_specs with
-        | Some vs -> (DEF_aux (DEF_val vs, def_annot) :: rdefs, IdSet.add id spec_found)
-        | None -> (def :: rdefs, spec_found)
+        | Some vs -> (DEF_aux (DEF_val vs, def_annot) :: rdefs, IdSet.add id specs_found)
+        | None -> (def :: rdefs, specs_found)
       )
-    | _ -> (def :: rdefs, spec_found)
+    | _ -> (def :: rdefs, specs_found)
   in
-  let rdefs, spec_found = List.fold_left check ([], IdSet.empty) defs in
-  (List.rev rdefs, spec_found)
+  let rdefs, specs_found = List.fold_left check ([], IdSet.empty) defs in
+  (List.rev rdefs, specs_found)
 
 let filter_replacements spec_found { defs; _ } =
   let not_found = function
@@ -107,6 +109,20 @@ let filter_replacements spec_found { defs; _ } =
   in
   List.filter not_found defs
 
+let move_replacement_fundefs ast =
+  let rec aux acc = function
+    | DEF_aux (DEF_pragma ("spliced_function#", id, _), _) :: defs ->
+        let is_replacement = function
+          | DEF_aux (DEF_fundef fd, _) -> string_of_id (id_of_fundef fd) = id
+          | _ -> false
+        in
+        let replacement, other_defs = List.partition is_replacement defs in
+        aux (replacement @ acc) other_defs
+    | def :: defs -> aux (def :: acc) defs
+    | [] -> List.rev acc
+  in
+  { ast with defs = aux [] ast.defs }
+
 let annotate_ast ast =
   let annotate_def (DEF_aux (d, a)) = DEF_aux (d, add_def_attribute a.loc "spliced" "" a) in
   { ast with defs = List.map annotate_def ast.defs }
@@ -120,4 +136,11 @@ let splice ast file =
   let repl_ids, repl_specs = scan_ast repl_ast in
   let defs1, specs_found = filter_old_ast repl_ids repl_specs ast in
   let defs2 = filter_replacements specs_found repl_ast in
-  Type_error.check Type_check.initial_env (Type_check.strip_ast { ast with defs = defs1 @ defs2 })
+  { ast with defs = defs1 @ defs2 }
+
+let splice_files ast files =
+  let spliced_ast = List.fold_left (fun ast file -> splice ast file) ast files in
+  let checked_ast, env = Type_error.check Type_check.initial_env (Type_check.strip_ast spliced_ast) in
+  (* Move replacement functions into place and top-sort, in case dependencies have changed *)
+  let sorted_ast = Callgraph.top_sort_defs (move_replacement_fundefs checked_ast) in
+  (sorted_ast, env)