feat(engine): introduce a new generic printer

cli/driver/src/exporter.rs

@@ -46,9 +46,24 @@ fn write_files(
         std::fs::write(&path, file.contents).unwrap_or_else(|e| {
             tcx.dcx().fatal(format!(
                 "Unable to write to file {:#?}. Error: {:#?}",
-                path, e
+                &path, e
             ))
-        })
+        });
+        if let Some(mut sourcemap) = file.sourcemap {
+            sourcemap.inline_sources_content();
+            sourcemap.sourceRoot = manifest_dir.display().to_string();
+            let sourcemap = serde_json::to_string_pretty(&sourcemap).unwrap();
+            let path = path.with_extension(format!(
+                "{}.map",
+                path.extension().unwrap().to_str().unwrap()
+            ));
+            std::fs::write(&path, sourcemap).unwrap_or_else(|e| {
+                session.fatal(format!(
+                    "Unable to write to file {:#?}. Error: {:#?}",
+                    path, e
+                ))
+            });
+        }
     }
 }
 
@@ -396,6 +411,7 @@ impl Callbacks for ExtractionCallbacks {
                     let engine_options = hax_cli_options_engine::EngineOptions {
                         backend: backend.clone(),
                         input: converted_items,
+                        manifest_dir: std::env::var("CARGO_MANIFEST_DIR").unwrap().into(),
                         impl_infos,
                     };
                     let mut engine_subprocess = find_hax_engine()

cli/options/engine/src/lib.rs

@@ -8,16 +8,46 @@ type ThirBody = hax_frontend_exporter::ThirBody;
 pub struct EngineOptions {
     pub backend: BackendOptions,
     pub input: Vec<hax_frontend_exporter::Item<ThirBody>>,
+    pub manifest_dir: String,
     pub impl_infos: Vec<(
         hax_frontend_exporter::DefId,
         hax_frontend_exporter::ImplInfos,
     )>,
 }
 
+#[allow(non_snake_case)]
+#[derive(JsonSchema, Debug, Clone, Serialize, Deserialize)]
+pub struct SourceMap {
+    pub mappings: String,
+    pub sourceRoot: String,
+    pub sources: Vec<String>,
+    pub sourcesContent: Vec<Option<String>>,
+    pub names: Vec<String>,
+    pub version: u8,
+    pub file: String,
+}
+
+impl SourceMap {
+    pub fn inline_sources_content(&mut self) {
+        self.sourcesContent = vec![];
+        for source in &self.sources {
+            let path = if self.sourceRoot.is_empty() {
+                source.clone()
+            } else {
+                format!("{}/{}", &self.sourceRoot, source)
+            };
+            eprintln!("path={path}");
+            let contents = Some(std::fs::read_to_string(path).unwrap());
+            self.sourcesContent.push(contents);
+        }
+    }
+}
+
 #[derive(JsonSchema, Debug, Clone, Serialize, Deserialize)]
 pub struct File {
     pub path: String,
     pub contents: String,
+    pub sourcemap: Option<SourceMap>,
 }
 
 #[derive(JsonSchema, Debug, Clone, Serialize, Deserialize)]

engine/backends/coq/coq/coq_backend.ml

@@ -691,6 +691,7 @@ let translate _ (_bo : BackendOptions.t) (items : AST.item list) :
              path = mod_name ^ ".v";
              contents =
                hardcoded_coq_headers ^ "\n" ^ string_of_items items ^ "\n";
+             sourcemap = None;
            })
 
 open Phase_utils

engine/backends/coq/ssprove/ssprove_backend.ml

@@ -2408,7 +2408,8 @@ let translate _ (_bo : BackendOptions.t) (items : AST.item list) :
            ^ "\n"
          in
 
-         Types.{ path = mod_name ^ ".v"; contents = file_content })
+         Types.
+           { path = mod_name ^ ".v"; contents = file_content; sourcemap = None })
 
 let apply_phases (_bo : BackendOptions.t) (i : Ast.Rust.item list) :
     AST.item list =

engine/backends/fstar/fstar_backend.ml

@@ -1523,7 +1523,49 @@ let fstar_headers (bo : BackendOptions.t) =
   in
   [ opts; "open Core"; "open FStar.Mul" ] |> String.concat ~sep:"\n"
 
-let translate m (bo : BackendOptions.t) (items : AST.item list) :
+module NewGenericPrinter = New_rust_printer.Make (InputLanguage)
+
+(** Use the generic printer instead of the F* printer. For now, there
+is no generic printer for F*, that's why we currently just use the
+Rust generic printer. Thus currently this exists only for debugging
+purposes. *)
+let translate_as_experimental_rust m (bo : BackendOptions.t)
+    (items : AST.item list) : Types.file list =
+  let show_view Concrete_ident.{ crate; path; definition } =
+    crate :: (path @ [ definition ]) |> String.concat ~sep:"::"
+  in
+  U.group_items_by_namespace items
+  |> Map.to_alist
+  |> List.concat_map ~f:(fun (ns, items) ->
+         let mod_name =
+           String.concat ~sep:"."
+             (List.map
+                ~f:(map_first_letter String.uppercase)
+                (fst ns :: snd ns))
+         in
+         let string_of_items _ _ items =
+           let r = NewGenericPrinter.items () items in
+           let str = New_generic_printer_api.AnnotatedString.to_string r in
+           let sm = New_generic_printer_api.AnnotatedString.to_sourcemap r in
+           let r = (str, sm) in
+           (r, r)
+         in
+         let impl, intf = string_of_items bo m items in
+         let make ~ext (body, sourcemap) =
+           if String.is_empty body then None
+           else
+             Some
+               Types.
+                 {
+                   path = mod_name ^ "." ^ ext;
+                   contents = body;
+                   sourcemap = Some sourcemap;
+                 }
+         in
+         List.filter_map ~f:Fn.id [ make ~ext:"rs" impl ])
+
+(** Translate as F* (the "legacy" printer) *)
+let translate_as_fstar m (bo : BackendOptions.t) (items : AST.item list) :
     Types.file list =
   let show_view Concrete_ident.{ crate; path; definition } =
     crate :: (path @ [ definition ]) |> String.concat ~sep:"::"
@@ -1555,11 +1597,19 @@ let translate m (bo : BackendOptions.t) (items : AST.item list) :
                    contents =
                      "module " ^ mod_name ^ "\n" ^ fstar_headers bo ^ "\n\n"
                      ^ body ^ "\n";
+                   sourcemap = None;
                  }
          in
          List.filter_map ~f:Fn.id
            [ make ~ext:"fst" impl; make ~ext:"fsti" intf ])
 
+let translate =
+  if
+    Sys.getenv "HAX_ENGINE_EXPERIMENTAL_RUST_PRINTER_INSTEAD_OF_FSTAR"
+    |> Option.is_some
+  then translate_as_experimental_rust
+  else translate_as_fstar
+
 open Phase_utils
 module DepGraph = Dependencies.Make (InputLanguage)
 module DepGraphR = Dependencies.Make (Features.Rust)

engine/backends/proverif/proverif_backend.ml

@@ -873,9 +873,12 @@ let translate m (bo : BackendOptions.t) (items : AST.item list) :
     ^ M.Processes.print items
   in
   let analysis_contents = M.Toplevel.print items in
-  let lib_file = Types.{ path = "lib.pvl"; contents = lib_contents } in
+  let lib_file =
+    Types.{ path = "lib.pvl"; contents = lib_contents; sourcemap = None }
+  in
   let analysis_file =
-    Types.{ path = "analysis.pv"; contents = analysis_contents }
+    Types.
+      { path = "analysis.pv"; contents = analysis_contents; sourcemap = None }
   in
   [ lib_file; analysis_file ]
 

engine/default.nix

@@ -10,6 +10,7 @@
   gnused,
   lib,
   removeReferencesTo,
+  ocaml-sourcemaps,
 }: let
   non_empty_list = ocamlPackages.buildDunePackage rec {
     pname = "non_empty_list";
@@ -65,6 +66,7 @@
         re
         js_of_ocaml
         ocamlgraph
+        ocaml-sourcemaps
       ]
       ++
       # F* dependencies

engine/lib/dune

@@ -12,6 +12,7 @@
   core
   logs
   re
+  sourcemaps
   ocamlgraph)
  (preprocessor_deps
   ; `ppx_inline` is used on the `Subtype` module, thus we need it at PPX time

engine/lib/new_generic_printer/new_generic_printer_api.ml

@@ -0,0 +1,92 @@
+open! Prelude
+open New_generic_printer_base
+
+module AnnotatedString = struct
+  type t = string * Annotation.t list [@@deriving show, yojson, eq]
+
+  let to_string = fst
+
+  let to_spanned_strings ((s, annots) : t) : (Ast.span * string) list =
+    Annotation.split_with_string s annots
+
+  let to_sourcemap : t -> Types.source_map =
+    snd >> List.filter_map ~f:Annotation.to_mapping >> Sourcemaps.Source_maps.mk
+    >> fun ({
+              mappings;
+              sourceRoot;
+              sources;
+              sourcesContent;
+              names;
+              version;
+              file;
+            } :
+             Sourcemaps.Source_maps.t) ->
+    Types.
+      { mappings; sourceRoot; sources; sourcesContent; names; version; file }
+end
+
+module Make (F : Features.T) = struct
+  module AST = Ast.Make (F)
+  open Ast.Make (F)
+  open SecretTypes
+
+  type print_object =
+    < printer_name : string
+    ; get_span_data : unit -> Annotation.t list
+    ; ty : (par_state -> ty fn) no_override
+    ; pat : (par_state -> pat fn) no_override
+    ; arm : arm fn no_override
+    ; expr : (par_state -> expr fn) no_override
+    ; item : item fn no_override
+    ; items : item list fn >
+  (** In the end, an printer *object* should be of the type {!print_object}. *)
+
+  module type API = sig
+    type aux_info
+
+    val items : aux_info -> item list -> AnnotatedString.t
+    val item : aux_info -> item -> AnnotatedString.t
+    val expr : aux_info -> expr -> AnnotatedString.t
+    val pat : aux_info -> pat -> AnnotatedString.t
+    val ty : aux_info -> ty -> AnnotatedString.t
+  end
+
+  module Api (NewPrint : sig
+    type aux_info
+
+    val new_print : aux_info -> print_object
+  end) =
+  struct
+    open NewPrint
+
+    let mk' (f : print_object -> 'a -> PPrint.document) (aux : aux_info)
+        (x : 'a) : AnnotatedString.t =
+      let printer = new_print aux in
+      let doc = f printer x in
+      let buf = Buffer.create 0 in
+      PPrint.ToBuffer.pretty 1.0 80 buf doc;
+      (Buffer.contents buf, printer#get_span_data ())
+
+    let mk (f : print_object -> 'a fn no_override) =
+      mk' (fun (po : print_object) ->
+          let f : 'a fn no_override = f po in
+          let f = !:f in
+          f)
+
+    type aux_info = NewPrint.aux_info
+
+    let items : aux_info -> item list -> AnnotatedString.t =
+      mk' (fun p -> p#items)
+
+    let item : aux_info -> item -> AnnotatedString.t = mk (fun p -> p#item)
+
+    let expr : aux_info -> expr -> AnnotatedString.t =
+      mk' (fun p -> !:(p#expr) AlreadyPar)
+
+    let pat : aux_info -> pat -> AnnotatedString.t =
+      mk' (fun p -> !:(p#pat) AlreadyPar)
+
+    let ty : aux_info -> ty -> AnnotatedString.t =
+      mk' (fun p -> !:(p#ty) AlreadyPar)
+  end
+end