Merge pull request #1518 from cryspen/rengine-add-resguaring-nodes

Rust Engine: intro. resugared AST fragment

Author: W95Psp

Cargo.lock

@@ -17,3 +17,4 @@ serde_json = { workspace = true, features = ["unbounded_depth"] }
 schemars.workspace = true
 serde-jsonlines = "0.5.0"
 serde_stacker = "0.1.12"
+pretty = "0.12.4"

rust-engine/Cargo.toml

@@ -14,6 +14,7 @@ pub mod diagnostics;
 pub mod fragment;
 pub mod identifiers;
 pub mod literals;
+pub mod resugared;
 pub mod span;
 
 use crate::symbol::Symbol;
@@ -22,6 +23,7 @@ use fragment::Fragment;
 use hax_rust_engine_macros::*;
 use identifiers::*;
 use literals::*;
+use resugared::*;
 use span::Span;
 
 /// Represents a generic value used in type applications (e.g., `T` in `Vec<T>`).
@@ -171,6 +173,11 @@ pub enum TyKind {
     /// ```
     Dyn(Vec<DynTraitGoal>),
 
+    /// A resugared type.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredTyKind),
+
     /// Fallback constructor to carry errors.
     Error(Diagnostic),
 }
@@ -368,6 +375,11 @@ pub enum PatKind {
         fields: Vec<(GlobalId, Pat)>,
     },
 
+    /// A resugared pattern.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredPatKind),
+
     /// Fallback constructor to carry errors.
     Error(Diagnostic),
 }
@@ -582,6 +594,11 @@ pub enum ImplItemKind {
         /// The list of the argument for the associated function (`&self` in the example).
         params: Vec<Param>,
     },
+
+    /// A resugared impl item.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredImplItemKind),
 }
 
 /// Represents a trait item (associated type, fn, or default)
@@ -628,6 +645,11 @@ pub enum TraitItemKind {
         /// The default body of the associated function (`x + 2` in the example).
         body: Expr,
     },
+
+    /// A resugared trait item.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredTraitItemKind),
 }
 
 /// A QuoteContent is a component of a quote: it can be a verbatim string, a Rust expression to embed in the quote, a pattern etc.
@@ -1019,6 +1041,11 @@ pub enum ExprKind {
         contents: Quote,
     },
 
+    /// A resugared expression.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredExprKind),
+
     /// Fallback constructor to carry errors.
     Error(Diagnostic),
 }
@@ -1388,6 +1415,11 @@ pub enum ItemKind {
     /// Fallback constructor to carry errors.
     Error(Diagnostic),
 
+    /// A resugared item.
+    /// This variant is introduced before printing only.
+    /// Phases must not produce this variant.
+    Resugared(ResugaredItemKind),
+
     /// Item that is not implemented yet
     NotImplementedYet,
 }
@@ -1455,8 +1487,26 @@ pub mod traits {
         };
     }
 
-    derive_has_metadata!(Item, Expr, Pat);
-    derive_has_kind!(Item => ItemKind, Expr => ExprKind, Pat => PatKind);
+    derive_has_metadata!(
+        Item,
+        Expr,
+        Pat,
+        Guard,
+        Arm,
+        ImplItem,
+        TraitItem,
+        GenericParam
+    );
+    derive_has_kind!(
+        Item => ItemKind, Expr => ExprKind, Pat => PatKind, Guard => GuardKind,
+        GenericParam => GenericParamKind, ImplItem => ImplItemKind, TraitItem => TraitItemKind
+    );
+
+    impl HasSpan for Attribute {
+        fn span(&self) -> Span {
+            self.span.clone()
+        }
+    }
 
     impl Typed for Expr {
         fn ty(&self) -> &Ty {

rust-engine/src/ast.rs

@@ -0,0 +1,38 @@
+//! This module defines *resugared fragments* for the Hax Rust engine's AST.
+//!
+//! A resugared fragment is an additional AST node used solely for pretty-printing purposes.
+//! These nodes carry no semantic meaning in hax core logic but enable more accurate
+//! or backend-specific surface syntax reconstruction.
+//!
+//! For example, the engine represents the `unit` type as a zero-sized tuple `()`,
+//! mirroring Rust's internal representation. However, this may not suit all backends:
+//! in F*, `unit` is explicitly written as `unit`, not `()`.
+//!
+//! To accommodate such differences, we introduce resugared fragments (e.g. `UnitType`) that
+//! allow the printer to emit the expected syntax while maintaining the same internal semantics.
+
+use hax_rust_engine_macros::*;
+
+/// Resugared variants for items. This represent extra printing-only items, see [`super::ItemKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredItemKind {}
+
+/// Resugared variants for expressions. This represent extra printing-only expressions, see [`super::ExprKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredExprKind {}
+
+/// Resugared variants for patterns. This represent extra printing-only patterns, see [`super::PatKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredPatKind {}
+
+/// Resugared variants for types. This represent extra printing-only types, see [`super::TyKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredTyKind {}
+
+/// Resugared variants for impl. items. This represent extra printing-only impl. items, see [`super::ImplItemKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredImplItemKind {}
+
+/// Resugared variants for trait items. This represent extra printing-only trait items, see [`super::TraitItemKind::Resugared`].
+#[derive_group_for_ast]
+pub enum ResugaredTraitItemKind {}

rust-engine/src/ast/resugared.rs

@@ -11,4 +11,5 @@
 pub mod ast;
 pub mod hax_io;
 pub mod ocaml_engine;
+pub mod printer;
 pub mod symbol;

rust-engine/src/lib.rs

@@ -0,0 +1,4 @@
+//! This modules provides types and helper for the printers of hax.
+
+mod allocator;
+pub use allocator::Allocator;

rust-engine/src/printer.rs

@@ -0,0 +1,87 @@
+//! This module provides a custom [`pretty`] allocator, indexed by a printer,
+//! enabling multiple printers to cohexist and to implement the type `Pretty`.
+
+use crate::ast::span::Span;
+use pretty::*;
+
+/// A printer-specific [`pretty`] allocator.
+pub struct Allocator<Printer> {
+    /// The pretty allocator
+    allocator: BoxAllocator,
+    /// Extra printer-specific context
+    pub printer: Printer,
+}
+
+impl<Printer> Allocator<Printer> {
+    /// Creates a new allocator from a printer.
+    pub fn new(printer: Printer) -> Self {
+        Self {
+            allocator: BoxAllocator,
+            printer,
+        }
+    }
+}
+
+impl<'a, P, A: 'a> DocAllocator<'a, A> for Allocator<P> {
+    type Doc = <BoxAllocator as DocAllocator<'a, A>>::Doc;
+
+    fn alloc(&'a self, doc: Doc<'a, Self::Doc, A>) -> Self::Doc {
+        self.allocator.alloc(doc)
+    }
+
+    fn alloc_column_fn(
+        &'a self,
+        f: impl Fn(usize) -> Self::Doc + 'a,
+    ) -> <Self::Doc as DocPtr<'a, A>>::ColumnFn {
+        self.allocator.alloc_column_fn(f)
+    }
+
+    fn alloc_width_fn(
+        &'a self,
+        f: impl Fn(isize) -> Self::Doc + 'a,
+    ) -> <Self::Doc as DocPtr<'a, A>>::WidthFn {
+        self.allocator.alloc_width_fn(f)
+    }
+}
+
+/// A helper type used to manually implement `Pretty` for types that carry spans.
+///
+/// By default, we implement the `Pretty` trait for all span-carrying
+/// types. These implementations annotate spans in the generated document, allowing
+/// source spans to be produced during pretty-printing. However, this default behavior
+/// does not provide access to the underlying data, which is sometimes necessary
+/// for custom printing logic.
+///
+/// For example, when printing an item, it's often useful to access its attributes.
+/// To support this, the default `Pretty` implementations delegate to `Manual<Item>`,
+/// which allows printers to access the inner value directly.
+///
+/// In practice, calling `expr.pretty(..)` will internally use
+/// `Manual(expr).pretty(..)`, enabling more flexible control over printing behavior.
+struct Manual<T>(T);
+
+use crate::ast::*;
+macro_rules! impl_pretty_kind_meta {
+        ($($type:ty),*) => {
+            $(impl<'a, 'b, P> Pretty<'a, Allocator<P>, Span> for &'b $type
+            where
+                Manual<&'b $type>: Pretty<'a, Allocator<P>, Span>,
+            {
+                fn pretty(self, allocator: &'a Allocator<P>) -> DocBuilder<'a, Allocator<P>, Span> {
+                    let doc = Manual(self).pretty(allocator);
+                    doc.annotate(self.span())
+                }
+            })*
+        };
+    }
+impl_pretty_kind_meta!(
+    Item,
+    Expr,
+    Pat,
+    Guard,
+    Arm,
+    ImplItem,
+    TraitItem,
+    GenericParam,
+    Attribute
+);