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

engine/lib/new_generic_printer/new_generic_printer_base.ml

@@ -0,0 +1,175 @@
+open! Prelude
+open! Ast
+open PPrint
+
+module SecretTypes = struct
+  type 't no_override = 't
+  type 'location no_direct_call = unit
+end
+
+let ( !: ) (type a) (f : a SecretTypes.no_override) : a = f
+
+include New_generic_printer_base_sig.Types
+
+module Make (F : Features.T) = struct
+  module AST = Ast.Make (F)
+  open Ast.Make (F)
+  open New_generic_printer_base_sig.Make (F) (SecretTypes)
+
+  class virtual base : print_base_type =
+    object (print)
+      val mutable current_span = Span.default
+      val mutable span_data : Annotation.t list = []
+      val mutable current_namespace : (string * string list) option = None
+      method get_span_data () = span_data
+
+      method with_span ~(span : span) (f : unit -> document) : document =
+        let prev_span = current_span in
+        current_span <- span;
+        let doc = f () |> print#spanned_doc |> custom in
+        current_span <- prev_span;
+        doc
+
+      method spanned_doc (doc : document) : custom =
+        let span = current_span in
+        object
+          method requirement : requirement = requirement doc
+
+          method pretty : output -> state -> int -> bool -> unit =
+            fun o s i b ->
+              span_data <-
+                ({ line = s.line; col = s.column }, span) :: span_data;
+              pretty o s i b doc
+
+          method compact : output -> unit = fun o -> compact o doc
+        end
+
+      method expr_at = print#par_state >> print#expr
+      method ty_at = print#par_state >> print#ty
+      method pat_at = print#par_state >> print#pat
+
+      method expr ctx (e : expr) =
+        let span = e.span in
+        print#with_span ~span (fun _ ->
+            try print#__expr ctx e
+            with Diagnostics.SpanFreeError.Exn (Data (_context, _kind)) ->
+              failwith "todo")
+
+      method ty ctx full =
+        match full with
+        | TApp { ident = `Concrete ident; args } ->
+            print#ty_TApp_application ~full ident args |> group
+        | TApp
+            {
+              ident = `Primitive _ | `TupleCons _ | `TupleField _ | `Projector _;
+              _;
+            } ->
+            print#assertion_failure "TApp not concrete"
+        | TApp { ident = `TupleType size; args } ->
+            let args =
+              List.filter_map ~f:(function GType t -> Some t | _ -> None) args
+            in
+            if [%equal: int] (List.length args) size |> not then
+              print#assertion_failure "malformed [ty.TApp] tuple";
+            print#ty_TApp_tuple ~full args
+        | TApp _ -> .
+        | _ -> print#ty_ () ctx full
+
+      method pat ctx (full : pat) =
+        print#with_span ~span:full.span (fun _ -> print#pat_ () ctx full)
+
+      method item i =
+        print#set_current_namespace
+          (print#namespace_of_concrete_ident i.ident |> Option.some);
+        print#with_span ~span:i.span (fun _ ->
+            try print#item_ () i
+            with Diagnostics.SpanFreeError.Exn (Data (context, kind)) ->
+              let error = Diagnostics.pretty_print_context_kind context kind in
+              (* let cast_item : item -> Ast.Full.item = Stdlib.Obj.magic in *)
+              (* let ast = cast_item i |> Print_rust.pitem_str in *)
+              let msg =
+                error ^ "\nLast available AST for this item:\n\n" ^ "ast"
+              in
+              (* TODO: if the printer is extremely broken, this results in a stack overflow *)
+              make_hax_error_item i.span i.ident msg |> print#item)
+
+      method private __expr ctx full =
+        match full.e with
+        | App { f = { e = GlobalVar i; _ } as f; args; generic_args; _ } -> (
+            let expect_one_arg where =
+              match args with
+              | [ arg ] -> arg
+              | _ -> print#assertion_failure @@ "Expected one arg at " ^ where
+            in
+            match i with
+            | `Concrete _ | `Primitive _ -> (
+                match (args, i) with
+                | [], `Concrete i ->
+                    print#expr_App_constant ~full i generic_args
+                | [], _ -> print#assertion_failure "Primitive app of arity 0"
+                | _ -> print#expr_App_application ~full f args generic_args)
+            | `TupleType _ | `TupleCons _ | `TupleField _ ->
+                print#assertion_failure "App: unexpected tuple"
+            | `Projector (`TupleField (nth, size)) ->
+                let arg = expect_one_arg "projector tuple field" in
+                print#expr_App_tuple_projection ~full ~size ~nth arg
+            | `Projector (`Concrete i) ->
+                let arg = expect_one_arg "projector concrete" in
+                print#expr_App_field_projection ~full i arg)
+        | App { f; args; generic_args; _ } ->
+            print#expr_App_application ~full f args generic_args
+        | Construct { constructor; fields; base; is_record; is_struct } -> (
+            match constructor with
+            | `Concrete constructor ->
+                print#expr_Construct_inductive ~full ~is_record ~is_struct
+                  ~constructor ~base fields
+            | `TupleCons _ ->
+                List.map ~f:snd fields |> print#expr_Construct_tuple ~full
+            | `Primitive _ | `TupleType _ | `TupleField _ | `Projector _ ->
+                print#assertion_failure "Construct unexpected constructors")
+        | App _ | Construct _ -> .
+        | _ -> print#expr_ () ctx full
+
+      method arm (full : arm) =
+        print#with_span ~span:full.span (fun _ -> print#arm_ () full)
+
+      method generic_param (full : generic_param) =
+        print#with_span ~span:full.span (fun _ -> print#generic_param_ () full)
+
+      method param_ty (full : param) =
+        match full.typ_span with
+        | Some span -> print#with_span ~span (fun _ -> print#param_ty_ () full)
+        | None -> print#param_ty_ () full
+
+      method impl_item (full : impl_item) =
+        print#with_span ~span:full.ii_span (fun _ -> print#impl_item_ () full)
+
+      method trait_item (full : trait_item) =
+        print#with_span ~span:full.ti_span (fun _ -> print#trait_item_ () full)
+
+      method attr (full : attr) =
+        print#with_span ~span:full.span (fun _ -> print#attr_ () full)
+
+      method concrete_ident id =
+        let current_ns = print#get_current_namespace () in
+        let id_ns = print#namespace_of_concrete_ident id in
+        print#concrete_ident_ ()
+          ~under_current_ns:
+            ([%equal: (string * string list) option] current_ns (Some id_ns))
+          id
+
+      method items = separate_map (twice hardline) print#item
+      method attrs = separate_map hardline print#attr
+
+      method assertion_failure : 'any. string -> 'any =
+        fun details ->
+          let span = Span.to_thir current_span in
+          let kind = Types.AssertionFailure { details } in
+          let ctx = Diagnostics.Context.GenericPrinter print#printer_name in
+          Diagnostics.SpanFreeError.raise ~span ctx kind
+
+      method set_current_namespace ns = current_namespace <- ns
+      method get_current_namespace () = current_namespace
+      method unreachable : 'any. unit -> 'any = failwith "Unreachable!"
+    end
+end

engine/lib/new_generic_printer/new_generic_printer_base.mli

@@ -0,0 +1,23 @@
+open! Prelude
+open! Ast
+
+include module type of struct
+  (** Protects some methods from being called or overrided. *)
+  module SecretTypes = struct
+    type 't no_override = private 't
+    (** Hello *)
+
+    type 'location no_direct_call = private unit
+    (** Hello *)
+  end
+
+  include New_generic_printer_base_sig.Types
+end
+
+val ( !: ) : 'a. 'a SecretTypes.no_override -> 'a
+
+module Make (F : Features.T) : sig
+  open New_generic_printer_base_sig.Make(F)(SecretTypes)
+
+  class virtual base : print_base_type
+end

engine/lib/new_generic_printer/new_generic_printer_base_sig.ml

@@ -0,0 +1,373 @@
+[@@@warning "-37-34-27"]
+
+open! Prelude
+open! Ast
+open PPrint
+
+module Types = struct
+  (** Generic printer for the {!module:Ast} ASTs. It uses the [PPrint]
+library, and additionaly computes {!Annotation.t}.  *)
+
+  (** Identifies a position in the AST. This is useful for figuring out
+wether we should wrap a chunk of AST in parenthesis. or not *)
+  type ast_position =
+    | GenericValue_GType
+    | GenericValue_GConst
+    | Lhs_LhsArbitraryExpr
+    | Lhs_LhsArrayAccessor
+    | Ty_TArrow
+    | Ty_TRef
+    | Ty_Tuple
+    | Ty_TSlice
+    | Ty_TArray_typ
+    | Ty_TArray_length
+    | Expr_If_cond
+    | Expr_If_then
+    | Expr_If_else
+    | Expr_Array
+    | Expr_Assign
+    | Expr_Closure_param
+    | Expr_Closure_body
+    | Expr_Ascription_e
+    | Expr_Ascription_typ
+    | Expr_Let_lhs
+    | Expr_Let_rhs
+    | Expr_Let_body
+    | Expr_Match_scrutinee
+    | Expr_QuestionMark
+    | Expr_Borrow
+    | Expr_TupleProjection
+    | Expr_ConstructTuple
+    | Expr_FieldProjection
+    | Expr_App_f
+    | Expr_App_arg
+    | Expr_ConcreteInductive_base
+    | Expr_ConcreteInductive_field
+    | Pat_PBinding_subpat
+    | Pat_PDeref
+    | Pat_PArray
+    | Pat_ConstructTuple
+    | Pat_ConcreteInductive
+    | Pat_Ascription_pat
+    | Pat_Ascription_typ
+    | Pat_Or
+    | Param_pat
+    | Param_typ
+    | GenericParam_GPType
+    | GenericParam_GPConst
+    | Arm_pat
+    | Arm_body
+    | Item_Fn_body
+  [@@warning "-37"]
+
+  module Annotation = struct
+    type loc = { line : int; col : int } [@@deriving show, yojson, eq]
+    type t = loc * span [@@deriving show, yojson, eq]
+
+    let compare ((a, _) : t) ((b, _) : t) : int =
+      let line = Int.compare a.line b.line in
+      if Int.equal line 0 then Int.compare a.col b.col else line
+
+    (** Converts a list of annotation and a string to a list of annotated string *)
+    let split_with_string (s : string) (annots : t list) =
+      let lines_position =
+        String.to_list s
+        |> List.filter_mapi ~f:(fun i ch ->
+               match ch with '\n' -> Some i | _ -> None)
+        |> List.to_array |> Array.get
+      in
+      let annots = List.sort ~compare annots in
+      let init = ({ line = 0; col = 0 }, None) in
+      let slices =
+        List.folding_map
+          ~f:(fun (start, start_span) (end_, end_span) ->
+            let span = Option.value ~default:end_span start_span in
+            ((end_, Some end_span), (span, start, end_)))
+          ~init annots
+      in
+      List.map slices ~f:(fun (span, start, end_) ->
+          let pos = lines_position start.line + start.col in
+          let len = lines_position end_.line + end_.col - pos in
+          (span, String.sub s ~pos ~len))
+
+    let to_mapping ((loc, span) : t) : Sourcemaps.Source_maps.mapping option =
+      let real_path (x : Types.file_name) =
+        match x with
+        | Real (LocalPath p) | Real (Remapped { local_path = Some p; _ }) ->
+            Some p
+        | _ -> None
+      in
+      let loc_to_loc ({ line; col } : loc) : Sourcemaps.Location.t =
+        { line; col }
+      in
+      let to_loc ({ col; line } : Types.loc) : loc =
+        { col = Int.of_string col; line = Int.of_string line - 1 }
+      in
+      let* span =
+        Span.to_thir span
+        |> List.find ~f:(fun (s : Types.span) ->
+               real_path s.filename |> Option.is_some)
+      in
+      let* src_filename = real_path span.filename in
+      let src_start = to_loc span.lo |> loc_to_loc in
+      let src_end = to_loc span.hi |> loc_to_loc in
+      let dst_start = loc_to_loc loc in
+      Some
+        Sourcemaps.Source_maps.
+          {
+            src = { start = src_start; end_ = Some src_end };
+            gen = { start = dst_start; end_ = None };
+            source = src_filename;
+            name = None;
+          }
+  end
+
+  type annot_str = string * Annotation.t list [@@deriving show, yojson, eq]
+
+  (** When printing a chunk of AST, should we wrap parenthesis
+({!NeedsPar}) or not ({!AlreadyPar})? *)
+  type par_state = NeedsPar | AlreadyPar
+
+  type 't fn = 't -> document
+end
+
+open Types
+
+module Make
+    (F : Features.T) (SecretTypes : sig
+      type 't no_override
+      type 'location no_direct_call
+    end) =
+struct
+  module AST = Ast.Make (F)
+  open Ast.Make (F)
+  open SecretTypes
+
+  (** Raw generic printers base class. Those are useful for building a
+  printer, not for consuming printers. Consumers should use
+  the {!module:Api} functor. *)
+  class type virtual print_base_type =
+    object
+
+      (** {1 Span handling} *)
+
+      (** Every piece of string rendered is contextualized with span information automatically. *)
+
+      method get_span_data : unit -> Annotation.t list
+      (** Retreive the mapping between locations in the rendered
+      string and Rust locations. *)
+
+      method with_span : span:span -> (unit -> document) -> document
+      (** [with_span ~span f] runs `f` in the context of [span]. *)
+
+      method spanned_doc : document -> custom
+      (** [spanned_doc doc] constructs a custom wrapping document for
+      [doc]. Rendering this document in [pretty] mode has a
+      side-effect: we push a [Annotation.t] to internal state. An
+      annotation maps a location within the rendered string to a Rust
+      span (that is, a location in the original Rust source code). *)
+
+      (** {1 [*_at] methods} *)
+
+      (** Always use [<name>_at] methods rather than [<name>]
+      ones. The former takes an [ast_position], that contextualizes
+      from where we are printing something. Printing the body of a
+      [let .. = ..;] expression (position [Expr_Let_body]) and
+      printing a function argument (position [Expr_App_arg]) will
+      probably require different parenthesizing: [ast_position] gives
+      contextual information upon which such parenthesizing decisions
+      can be taken. *)
+
+      method expr_at : ast_position -> expr fn
+      (** Renders an [expr] at some [ast_position]. *)
+
+      method ty_at : ast_position -> ty fn
+      (** Renders a [ty] at some [ast_position]. *)
+
+      method pat_at : ast_position -> pat fn
+      (** Renders a [pat] at some [ast_position]. *)
+
+      (** {1 Driver methods} *)
+
+      (** The methods in this section are defined in two flavors:
+          `<name>` and `<name>_`. `<name>` methods are not
+          overridable. Indeed, they take care of various things for
+          you:
+          
+          {ul {- catch exceptions and translate them as
+          pretty-printed errors with the original Rust AST;}
+              {- set contextual span information in a systematic way;}
+              {- disambiguate certain variant of the AST (see {!section-"specialized-printers"}).}}
+
+         Your can override `<name>_` methods.
+      *)
+
+      (** {2 Expressions} *)
+      method expr : (par_state -> expr fn) no_override
+      (** Prints an expression. Pre-handles the variants [App] and
+      [Construct]: see {!section-"specialize-expr"}. *)
+
+      method virtual expr_ : [ `Expr ] no_direct_call -> par_state -> expr fn
+      (** Overridable printer for expressions. Please mark the cases
+      [App] and [Construct] as unreachable. *)
+
+      (** {2 Types} *)
+      method ty : (par_state -> ty fn) no_override
+      (** Prints a type. Pre-handles [TApp]. *)
+
+      method virtual ty_ : [ `Ty ] no_direct_call -> par_state -> ty fn
+      (** Overridable printer for types. Please mark the case [TApp]
+      as unreachable. *)
+
+      (** {2 Patterns} *)
+      method pat : (par_state -> pat fn) no_override
+      (** Prints a pattern. *)
+
+      method virtual pat_ : [ `Pat ] no_direct_call -> par_state -> pat fn
+      (** Overridable printer for patterns. *)
+
+      (** {2 Items} *)
+      method item : item fn no_override
+      (** Prints a item. *)
+
+      method virtual item_ : [ `Item ] no_direct_call -> item fn
+      (** Overridable printer for items. *)
+
+      (** {2 Arms} *)
+      method arm : arm fn no_override
+      (** Prints an arm (in a match). *)
+
+      method virtual arm_ : [ `Arm ] no_direct_call -> arm fn
+      (** Overridable printer for arms (in matches).*)
+
+      (** {2 Generic parameters} *)
+      method generic_param : generic_param fn no_override
+      (** Prints a generic parameter. *)
+
+      method virtual generic_param_ : [ `GP ] no_direct_call -> generic_param fn
+      (** Overridable printer for generic parameters. *)
+
+      (** {2 Parameters} *)
+      method param_ty : param fn no_override
+      (** Prints the type of a parameter. This is special because of `typ_span`. *)
+
+      method virtual param_ty_ : [ `Param ] no_direct_call -> param fn
+      (** Overridable printer for parameter types. *)
+
+      (** {2 Impl items} *)
+      method impl_item : impl_item fn no_override
+      (** Prints an impl item. *)
+
+      method virtual impl_item_ : [ `II ] no_direct_call -> impl_item fn
+      (** Overridable printer for impl items. *)
+
+      (** {2 Trait items} *)
+      method trait_item : trait_item fn no_override
+      (** Prints an trait item. *)
+
+      method virtual trait_item_ : [ `TI ] no_direct_call -> trait_item fn
+      (** Overridable printer for trait items. *)
+
+      (** {2 Attributes} *)
+
+      method attr : attr fn no_override
+      (** Prints an attribute. *)
+
+      method virtual attr_ : [ `Attr ] no_direct_call -> attr fn
+      (** Overridable printer for attributes. *)
+
+      (** {2 Concrete idents} *)
+
+      method concrete_ident : concrete_ident fn no_override
+      (** Prints a concrete ident. *)
+
+      method virtual concrete_ident_ :
+        [ `CIdent ] no_direct_call -> under_current_ns:bool -> concrete_ident fn
+      (** Overridable printer for concrete idents. *)
+
+      (** {1:specialized-printers Specialized printers} *)
+
+      (** Some nodes in the AST are ambiguous as they encode multiple
+      language constructs: the `App` constructor of `expr` for
+      instance encodes (1) function applications, (2) fields
+      projectors, (3) constants... This is the same for `Construct`,
+      `TApp`, and some other.
+
+      This section defines specialized methods for those language
+      constructs. When the variant `<Variant>` of a type `<type>` in
+      the AST is encoding various language constructs, we defined
+      various methods named `<type>_<Variant>_<name>`. *)
+
+      (** {2:specialize-expr Specialized printers for [expr]} *)
+
+      method virtual expr_App_constant :
+        full:expr -> concrete_ident -> generic_value list fn
+      (** [expr_App_constant ~full e generics] prints the constant
+      [e] with generics [generics]. [full] is the unspecialized [expr]. *)
+
+      method virtual expr_App_application :
+        full:expr -> expr -> expr list -> generic_value list fn
+      (** [expr_App_application ~full e args generics] prints the
+      function application [e<...generics>(...args)]. [full] is the unspecialized [expr]. *)
+
+      method virtual expr_App_tuple_projection :
+        full:expr -> size:int -> nth:int -> expr fn
+      (** [expr_App_tuple_projection ~full ~size ~nth expr] prints
+      the projection of the [nth] component of the tuple [expr] of
+      size [size]. [full] is the unspecialized [expr]. *)
+
+      method virtual expr_App_field_projection :
+        full:expr -> concrete_ident -> expr fn
+      (** [expr_App_field_projection ~full field expr] prints the
+      projection of the field [field] in the expression [expr]. [full]
+      is the unspecialized [expr]. *)
+
+      method virtual expr_Construct_inductive :
+        full:expr ->
+        is_record:bool ->
+        is_struct:bool ->
+        constructor:concrete_ident ->
+        base:(expr * F.construct_base) option ->
+        (global_ident * expr) list fn
+      (** [expr_Construct_inductive ~full ~is_record ~is_struct
+      ~constructor ~base fields] prints the construction of an
+      inductive with base [base] and fields [fields]. [full] is the
+      unspecialized [expr]. TODO doc is_record is_struct *)
+
+      method virtual expr_Construct_tuple : full:expr -> expr list fn
+
+      (** {2:specialize-expr Specialized printers for [ty]} *)
+
+      method virtual ty_TApp_tuple : full:ty -> ty list fn
+      (** [ty_TApp_tuple ~full types] prints a tuple type with
+      compounds types [types]. [full] is the unspecialized [ty]. *)
+
+      method virtual ty_TApp_application :
+        full:ty -> concrete_ident -> generic_value list fn
+      (** [ty_TApp_application ~full typ generic_args] prints the type
+      [typ<...generic_args>]. [full] is the unspecialized [ty]. *)
+
+      method items : item list fn
+
+      (** {1 Misc methods} *)
+
+      (** {1 Convenience methods} *)
+
+      method attrs : attrs fn
+
+      method assertion_failure : 'any. string -> 'any
+      (** Helper that throws and reports an [Types.AssertionFailure] error. *)
+
+      method set_current_namespace : (string * string list) option -> unit
+      method get_current_namespace : unit -> (string * string list) option
+
+      method virtual namespace_of_concrete_ident :
+        concrete_ident -> string * string list
+
+      method virtual printer_name : string
+      method virtual par_state : ast_position -> par_state
+
+      method unreachable : 'any. unit -> 'any
+      (** Mark an unreachable place in the printer. *)
+    end
+end

engine/lib/new_generic_printer/new_generic_printer_template.ml

@@ -0,0 +1,64 @@
+module Make (F : Features.T) = struct
+  module AST = Ast.Make (F)
+  open Ast.Make (F)
+  module P = New_generic_printer_base.Make (F)
+  open PPrint
+
+  let unimplemented s = string ("unimplemented: " ^ s)
+
+  class print =
+    object
+      inherit P.base as _super
+      method ty_TApp_tuple ~full:_ _args = unimplemented "ty_TApp_tuple"
+
+      method ty_TApp_application ~full:_ _f _args =
+        unimplemented "ty_TApp_application"
+
+      method expr_App_constant ~full:_ _ident _generic_values =
+        unimplemented "expr_App_constant"
+
+      method expr_App_application ~full:_ _f _args _generics =
+        unimplemented "expr_App_application"
+
+      method expr_App_tuple_projection ~full:_ ~size:_ ~nth:_ _tuple =
+        unimplemented "expr_App_tuple_projection"
+
+      method expr_App_field_projection ~full:_ _ident _data =
+        unimplemented "expr_App_field_projection"
+
+      method expr_Construct_inductive ~full:_ ~is_record:_ ~is_struct:_
+          ~constructor:_ ~base:_ _fields =
+        unimplemented "expr_Construct_inductive"
+
+      method expr_Construct_tuple ~full:_ _components =
+        unimplemented "expr_Construct_tuple"
+
+      method expr_ _ _ctx _expr = unimplemented "expr_"
+      method ty_ _ _ctx _typ = unimplemented "ty_"
+      method pat_ _ _ctx _pat = unimplemented "pat_"
+      method item_ _ _item = unimplemented "item_"
+      method arm_ _ _arm = unimplemented "arm_"
+      method generic_param_ _ _gp = unimplemented "generic_param_"
+      method param_ty_ _ _param_ty = unimplemented "param_ty_"
+      method impl_item_ _ _ii = unimplemented "impl_item_"
+      method trait_item_ _ _ti = unimplemented "trait_item_"
+      method attr_ _ _attr = unimplemented "attr_"
+
+      method namespace_of_concrete_ident =
+        Concrete_ident.DefaultViewAPI.to_namespace
+
+      method printer_name = "blank-template"
+      method par_state _ = AlreadyPar
+
+      method concrete_ident_ _ ~under_current_ns:_ _ident =
+        unimplemented "concrete_ident_"
+    end
+
+  open New_generic_printer_api.Make (F)
+
+  include Api (struct
+    type aux_info = unit
+
+    let new_print _ = (new print :> print_object)
+  end)
+end

engine/lib/new_generic_printer/new_rust_printer.ml

@@ -0,0 +1,170 @@
+open Prelude
+
+module Make (F : Features.T) = struct
+  module AST = Ast.Make (F)
+  open Ast.Make (F)
+  open New_generic_printer_base
+  module P = New_generic_printer_base.Make (F)
+  open PPrint
+
+  let unimplemented s = string ("unimplemented: " ^ s)
+
+  module View = Concrete_ident.DefaultViewAPI
+
+  let iblock f = group >> jump 2 0 >> terminate (break 0) >> f >> group
+  let call = ( !: )
+
+  class print =
+    object (print)
+      inherit P.base as _super
+      method ty_TApp_tuple ~full:_ _args = unimplemented "ty_TApp_tuple"
+
+      method ty_TApp_application ~full:_ _f _args =
+        unimplemented "ty_TApp_application"
+
+      method expr_App_constant ~full:_ _ident _generic_values =
+        unimplemented "expr_App_constant"
+
+      method expr_App_application ~full:_ f _args _generics =
+        print#expr_at Expr_App_f f
+      (* print#expr_at Expr_App_f f ^/^ separate_map space (print#expr_at Expr_App_arg) args *)
+      (* unimplemented "expr_App_application" *)
+
+      method expr_App_tuple_projection ~full:_ ~size:_ ~nth:_ _tuple =
+        unimplemented "expr_App_tuple_projection"
+
+      method expr_App_field_projection ~full:_ _ident _data =
+        unimplemented "expr_App_field_projection"
+
+      method expr_Construct_inductive ~full:_ ~is_record:_ ~is_struct:_
+          ~constructor:_ ~base:_ _fields =
+        unimplemented "expr_Construct_inductive"
+
+      method expr_Construct_tuple ~full:_ _components =
+        unimplemented "expr_Construct_tuple"
+
+      method expr_ _ ctx expr =
+        let wrap_parens =
+          group
+          >> match ctx with AlreadyPar -> Fn.id | NeedsPar -> iblock braces
+        in
+        match expr.e with
+        | If { cond; then_; else_ } ->
+            let if_then =
+              (string "if" ^//^ nest 2 (print#expr_at Expr_If_cond cond))
+              ^/^ string "then"
+              ^//^ (print#expr_at Expr_If_then then_ |> braces |> nest 1)
+            in
+            (match else_ with
+            | None -> if_then
+            | Some else_ ->
+                if_then ^^ break 1 ^^ string "else" ^^ space
+                ^^ (print#expr_at Expr_If_else else_ |> iblock braces))
+            |> wrap_parens
+        | Match { scrutinee; arms } ->
+            let header =
+              string "match" ^^ space
+              ^^ (print#expr_at Expr_Match_scrutinee scrutinee
+                 |> terminate space |> iblock Fn.id)
+              |> group
+            in
+            let arms =
+              separate_map hardline
+                (call print#arm >> group >> nest 2
+                >> precede (bar ^^ space)
+                >> group)
+                arms
+            in
+            header ^^ iblock braces arms
+        | Let { monadic; lhs; rhs; body } ->
+            (Option.map ~f:(fun monad -> print#expr_monadic_let ~monad) monadic
+            |> Option.value ~default:print#expr_let)
+              ~lhs ~rhs body
+            |> wrap_parens
+        | Literal l -> print#literal l
+        | Block (e, _) -> call print#expr ctx e
+        | _ -> unimplemented "expr_todo"
+
+      method expr_monadic_let
+          : monad:supported_monads * F.monadic_binding ->
+            lhs:pat ->
+            rhs:expr ->
+            expr fn =
+        fun ~monad:_ ~lhs ~rhs body -> print#expr_let ~lhs ~rhs body
+
+      method expr_let : lhs:pat -> rhs:expr -> expr fn =
+        fun ~lhs ~rhs body ->
+          string "let"
+          ^/^ iblock Fn.id (print#pat_at Expr_Let_lhs lhs)
+          ^/^ equals
+          ^/^ iblock Fn.id (print#expr_at Expr_Let_rhs rhs)
+          ^^ semi
+          ^/^ (print#expr_at Expr_Let_body body |> group)
+
+      method literal =
+        function
+        | String s -> utf8string s |> dquotes
+        | Char c -> char c |> bquotes
+        | Int { value; negative; _ } ->
+            string value |> precede (if negative then minus else empty)
+        | Float { value; kind; negative } ->
+            string value
+            |> precede (if negative then minus else empty)
+            |> terminate
+                 (string (match kind with F32 -> "f32" | F64 -> "f64"))
+        | Bool b -> OCaml.bool b
+
+      method ty_ _ _ctx _typ = unimplemented "ty_"
+      method pat_ _ _ctx _pat = unimplemented "pat_"
+
+      method item_ _ item =
+        match item.v with
+        | Fn { name; generics; body; params } ->
+            let params =
+              iblock parens (separate_map (comma ^^ break 1) print#param params)
+            in
+            let generics =
+              separate_map comma (call print#generic_param) generics.params
+            in
+            string "fn" ^^ space
+            ^^ call print#concrete_ident name
+            ^^ generics ^^ params
+            ^^ iblock braces (print#expr_at Item_Fn_body body)
+        | _ -> string "item not implemented"
+
+      method param_ty_ _ _param_ty = unimplemented "param_ty_"
+
+      method param (param : param) =
+        let { pat; typ = _; typ_span = _; attrs } = param in
+        print#attrs attrs ^^ print#pat_at Param_pat pat ^^ space ^^ colon
+        ^^ space ^^ !:(print#param_ty) param
+
+      method arm_ _ _arm = unimplemented "arm_"
+      method generic_param_ _ _gp = unimplemented "generic_param_"
+      method impl_item_ _ _ii = unimplemented "impl_item_"
+      method trait_item_ _ _ti = unimplemented "trait_item_"
+      method attr_ _ _attr = unimplemented "attr_"
+
+      method namespace_of_concrete_ident =
+        Concrete_ident.DefaultViewAPI.to_namespace
+
+      method printer_name = "blank-template"
+      method par_state _ = AlreadyPar
+
+      method concrete_ident_ _ ~under_current_ns id =
+        let id = View.to_view id in
+        let chunks =
+          if under_current_ns then [ id.definition ]
+          else id.crate :: (id.path @ [ id.definition ])
+        in
+        separate_map (colon ^^ colon) utf8string chunks
+    end
+
+  open New_generic_printer_api.Make (F)
+
+  include Api (struct
+    type aux_info = unit
+
+    let new_print _ = (new print :> print_object)
+  end)
+end

engine/lib/new_generic_printer/new_rust_printer.mli

@@ -0,0 +1,4 @@
+module Make (F : Features.T) : sig
+  open New_generic_printer_api.Make(F)
+  include API with type aux_info = unit
+end

engine/lib/print_rust.ml

@@ -600,24 +600,67 @@ let rustfmt_annotated (x : AnnotatedString.t) : AnnotatedString.t =
   if String.equal rf "no" then x
   else try rustfmt_annotated' x with RetokenizationFailure -> x
 
-let pitem : item -> AnnotatedString.Output.t =
-  Raw.pitem >> rustfmt_annotated >> AnnotatedString.Output.convert
-
-let pitems : item list -> AnnotatedString.Output.t =
-  List.concat_map ~f:Raw.pitem
-  >> rustfmt_annotated >> AnnotatedString.Output.convert
-
-let pitem_str : item -> string = pitem >> AnnotatedString.Output.raw_string
-
-let pexpr_str (e : expr) : string =
-  let e = Raw.pexpr e in
-  let ( ! ) = AnnotatedString.pure @@ Span.dummy () in
-  let ( & ) = AnnotatedString.( & ) in
-  let prefix = "fn expr_wrapper() {" in
-  let suffix = "}" in
-  let item = !prefix & e & !suffix in
-  rustfmt_annotated item |> AnnotatedString.Output.convert
-  |> AnnotatedString.Output.raw_string |> Stdlib.String.trim
-  |> String.chop_suffix_if_exists ~suffix
-  |> String.chop_prefix_if_exists ~prefix
-  |> Stdlib.String.trim
+module type T = sig
+  val pitem : item -> AnnotatedString.Output.t
+  val pitems : item list -> AnnotatedString.Output.t
+  val pitem_str : item -> string
+  val pexpr_str : expr -> string
+end
+
+module Traditional : T = struct
+  let pitem : item -> AnnotatedString.Output.t =
+    Raw.pitem >> rustfmt_annotated >> AnnotatedString.Output.convert
+
+  let pitems : item list -> AnnotatedString.Output.t =
+    List.concat_map ~f:Raw.pitem
+    >> rustfmt_annotated >> AnnotatedString.Output.convert
+
+  let pitem_str : item -> string = pitem >> AnnotatedString.Output.raw_string
+
+  let pexpr_str (e : expr) : string =
+    let e = Raw.pexpr e in
+    let ( ! ) = AnnotatedString.pure @@ Span.dummy () in
+    let ( & ) = AnnotatedString.( & ) in
+    let prefix = "fn expr_wrapper() {" in
+    let suffix = "}" in
+    let item = !prefix & e & !suffix in
+    rustfmt_annotated item |> AnnotatedString.Output.convert
+    |> AnnotatedString.Output.raw_string |> Stdlib.String.trim
+    |> String.chop_suffix_if_exists ~suffix
+    |> String.chop_prefix_if_exists ~prefix
+    |> Stdlib.String.trim
+end
+
+module Experimental : T = struct
+  module NewRustGenericPrinter = New_rust_printer.Make (Features.Full)
+
+  let pitem : item -> AnnotatedString.Output.t =
+    NewRustGenericPrinter.item ()
+    >> New_generic_printer_api.AnnotatedString.to_spanned_strings
+    >> AnnotatedString.Output.convert
+
+  let pitems : item list -> AnnotatedString.Output.t =
+    NewRustGenericPrinter.items ()
+    >> New_generic_printer_api.AnnotatedString.to_spanned_strings
+    >> AnnotatedString.Output.convert
+
+  let pexpr : expr -> AnnotatedString.Output.t =
+    NewRustGenericPrinter.expr ()
+    >> New_generic_printer_api.AnnotatedString.to_spanned_strings
+    >> AnnotatedString.Output.convert
+
+  let pitem_str : item -> string =
+    NewRustGenericPrinter.item ()
+    >> New_generic_printer_api.AnnotatedString.to_string
+
+  let pexpr_str : expr -> string =
+    NewRustGenericPrinter.expr ()
+    >> New_generic_printer_api.AnnotatedString.to_string
+end
+
+let experimental =
+  Sys.getenv "HAX_ENGINE_EXPERIMENTAL_RUST_PRINTER" |> Option.is_some
+
+include
+  (val if experimental then (module Experimental : T)
+       else (module Traditional : T))

flake.nix

@@ -23,6 +23,14 @@
       url = "github:hacl-star/hacl-star";
       flake = false;
     };
+    ocaml-sourcemaps = {
+      url = "github:W95Psp/ocaml-sourcemaps";
+      flake = true;
+      inputs = {
+        nixpkgs.follows = "nixpkgs";
+        flake-utils.follows = "flake-utils";
+      };
+    };
   };
 
   outputs = {
@@ -56,6 +64,8 @@
             cat "${hax-env-file}" | xargs -I{} echo "export {}"
           fi
         '';
+        ocaml-sourcemaps = inputs.ocaml-sourcemaps.packages.${system}.default;
+        ocamlPackages = pkgs.ocamlPackages;
       in rec {
         packages = {
           inherit rustc ocamlformat rustfmt fstar hax-env;
@@ -72,7 +82,7 @@
               #!${pkgs.stdenv.shell}
               ${packages.hax-rust-frontend.hax-engine-names-extract}/bin/hax-engine-names-extract | sed 's|/nix/store/\(.\{6\}\)|/nix_store/\1-|g'
             '';
-            inherit rustc;
+            inherit rustc ocaml-sourcemaps ocamlPackages;
           };
           hax-rust-frontend = pkgs.callPackage ./cli {
             inherit rustc craneLib;
@@ -146,11 +156,11 @@
         in let
           packages = [
             ocamlformat
-            pkgs.ocamlPackages.ocaml-lsp
-            pkgs.ocamlPackages.ocamlformat-rpc-lib
-            pkgs.ocamlPackages.ocaml-print-intf
-            pkgs.ocamlPackages.odoc
-            pkgs.ocamlPackages.utop
+            ocamlPackages.ocaml-lsp
+            ocamlPackages.ocamlformat-rpc-lib
+            ocamlPackages.ocaml-print-intf
+            ocamlPackages.odoc
+            ocamlPackages.utop
 
             pkgs.cargo-expand
             pkgs.cargo-insta