fix(lean): Add trait-level arguments as explicit arguments to trait items

CHANGELOG.md

@@ -37,6 +37,8 @@ Changes to the Lean backend:
  - Add support for binding subpatterns in match constructs (#1790)
  - Add error when using patterns in function parameters (#1792)
  - Add support for constant parameters to functions and traits (#1797)
+ - Make trait-level arguments explicit for all trait functions, adding them as
+   extra parameters (#1803)
 
 Miscellaneous:
 - Reserve extraction folder for auto-generated files in Lean examples (#1754)

hax-lib/proof-libs/lean/Hax/Core/Convert.lean

@@ -12,22 +12,32 @@ open Std.Do
 
 set_option mvcgen.warning false
 
-/- Warning : this function has been specialized, it should be turned into a typeclass -/
-def Core.Convert.TryInto.try_into {α n} (a: Array α) :
-   RustM (Core.Result.Result (Vector α n) Core.Array.TryFromSliceError) :=
+namespace Core.Convert
+
+class TryInto (Self: Type) (T: Type) {E: Type} where
+  try_into (Self T) : Self → (RustM (Core.Result.Result T E))
+
+instance {α : Type} {n : Nat} : TryInto (RustSlice α) (RustArray α n) (E := Core.Array.TryFromSliceError) where
+  try_into a :=
    pure (
      if h: a.size = n then
        Core.Result.Result.Ok (a.toVector.cast h)
      else
        .Err Core.Array.TryFromSliceError.array.TryFromSliceError
      )
 
+#check fun  {α : Type} {n: Nat} (a: RustSlice α) => TryInto.try_into (RustSlice α) (RustArray α n) a
+
+#check fun {α n} => (inferInstance : TryInto (RustSlice α) (RustArray α n))
+
 @[spec]
-theorem Core.Convert.TryInto.try_into.spec {α n} (a: Array α) :
+theorem TryInto.try_into.spec {α : Type} {n: Nat} (a: RustSlice α) :
   (h: a.size = n) →
   ⦃ ⌜ True ⌝ ⦄
-  ( Core.Convert.TryInto.try_into a)
+  (TryInto.try_into (RustSlice α) (RustArray α n) (E := Core.Array.TryFromSliceError) a )
   ⦃ ⇓ r => ⌜ r = .Ok (a.toVector.cast h) ⌝ ⦄ := by
   intro h
-  mvcgen [Core.Result.Impl.unwrap.spec, Core.Convert.TryInto.try_into]
+  mvcgen [TryInto.try_into]
   grind
+
+end Core.Convert

hax-lib/proof-libs/lean/Hax/Core/Ops/Function.lean

@@ -16,7 +16,7 @@ class FnOnce (Self Args : Type) {Output: Type} where
 
 class Fn (Self Args : Type) {Output: Type} where
   [_constr_10219838627318688691 : (FnOnce Self Args (Output := Output))]
-  call : Self -> Args -> RustM Output
+  call (Self Args) : Self -> Args -> RustM Output
 
 instance {α β} : FnOnce (α → RustM β) α (Output := β) where
   Output := β

hax-lib/proof-libs/lean/Hax/Lib.lean

@@ -421,7 +421,7 @@ section Cast
 
 /-- Hax-introduced explicit cast. It is partial (returns a `RustM`) -/
 @[simp, spec]
-def Core.Convert.From.from {α β} [Coe α (RustM β)] (x:α) : (RustM β) := x
+def Core.Convert.From.from (β α) [Coe α (RustM β)] (x:α) : (RustM β) := x
 
 /-- Rust-supported casts on base types -/
 class Cast (α β: Type) where
@@ -883,7 +883,7 @@ end RustVectors
 
 
 -- Miscellaneous
-def Core.Ops.Deref.Deref.deref {α Allocator} (v: Alloc.Vec.Vec α Allocator)
+def Core.Ops.Deref.Deref.deref {α Allocator} (β : Type) (v: Alloc.Vec.Vec α Allocator)
   : RustM (Array α)
   := pure v
 

rust-engine/src/ast.rs

@@ -1131,9 +1131,9 @@ pub struct GenericParam {
 /// Generic parameters and constraints (contained between `<>` in function declarations)
 #[derive_group_for_ast]
 pub struct Generics {
-    /// A vector of genreric parameters.
+    /// A vector of generic parameters.
     pub params: Vec<GenericParam>,
-    /// A vector of genreric constraints.
+    /// A vector of generic constraints.
     pub constraints: Vec<GenericConstraint>,
 }
 

rust-engine/src/backends/lean.rs

@@ -243,7 +243,7 @@ const _: () = {
         };
     }
 
-    // Methods for handling arguments of variants (or struct constructor)
+    // Extra methods, specific to the LeanPrinter
     impl LeanPrinter {
         /// Prints arguments a variant or constructor of struct, using named or unamed arguments based
         /// on the `is_record` flag. Used for both expressions and patterns
@@ -326,14 +326,11 @@ const _: () = {
             docs!["do", line!(), body].group()
         }
 
-        /// Produces a fresh name for a constraint on an associated type. It needs a fresh name to
-        /// be added as an extra field
-        fn fresh_constraint_name(
-            &self,
-            associated_type_name: &String,
-            constraint: &ImplIdent,
-        ) -> String {
-            format!("_constr_{}_{}", associated_type_name, constraint.name)
+        /// Produces a name for a constraint on an trait-level constraint, or an associated
+        /// type. The name is obtained by combining the type it applies to and the name of the
+        /// constraint (and should be unique)
+        fn constraint_name(&self, type_name: &String, constraint: &ImplIdent) -> String {
+            format!("_constr_{}_{}", type_name, constraint.name)
         }
 
         /// Turns an expression of type `RustM T` into one of type `T` (out of the monad), providing
@@ -359,6 +356,90 @@ const _: () = {
                 }
             }
         }
+
+        /// Print trait items, adding trait-level params as extra arguments
+        fn trait_item_with_trait_params<A: 'static + Clone>(
+            &self,
+            trait_generics: &[GenericParam],
+            TraitItem {
+                meta: _,
+                kind,
+                generics: item_generics,
+                ident,
+            }: &TraitItem,
+        ) -> DocBuilder<A> {
+            {
+                let name = self.render_last(ident);
+                let trait_generics = intersperse!(
+                    trait_generics
+                        .iter()
+                        .map(|GenericParam { ident, .. }| ident),
+                    softline!()
+                )
+                .parens()
+                .group()
+                .append(line!());
+                docs![match kind {
+                    TraitItemKind::Fn(ty) => {
+                        docs![
+                            name,
+                            softline!(),
+                            trait_generics,
+                            item_generics,
+                            ":",
+                            line!(),
+                            ty
+                        ]
+                        .group()
+                        .nest(INDENT)
+                    }
+                    TraitItemKind::Type(constraints) => {
+                        docs![
+                            name.clone(),
+                            reflow!(" : Type"),
+                            concat!(constraints.iter().map(|c| docs![
+                            hardline!(),
+                            docs![
+                                self.constraint_name(&name, c),
+                                reflow!(" :"),
+                                line!(),
+                                &c.goal
+                            ]
+                            .group()
+                            .nest(INDENT)
+                            .brackets()
+                        ]))
+                        ]
+                    }
+                    TraitItemKind::Default { params, body } => docs![
+                        docs![
+                            name,
+                            softline!(),
+                            trait_generics,
+                            item_generics,
+                            zip_right!(params, line!()).group(),
+                            docs![": RustM ", body.ty].group(),
+                            line!(),
+                            ":= do",
+                        ]
+                        .group(),
+                        line!(),
+                        body,
+                    ]
+                    .group()
+                    .nest(INDENT),
+                    TraitItemKind::Resugared(_) => {
+                        unreachable!("This backend has no resugaring for trait items")
+                    }
+                }]
+            }
+        }
+    }
+
+    impl<A: 'static + Clone> ToDocument<LeanPrinter, A> for (Vec<GenericParam>, &TraitItem) {
+        fn to_document(&self, printer: &LeanPrinter) -> DocBuilder<A> {
+            printer.trait_item_with_trait_params(&self.0, self.1)
+        }
     }
 
     impl<A: 'static + Clone> PrettyAst<A> for LeanPrinter {
@@ -486,7 +567,7 @@ set_option linter.unusedVariables false
                     args,
                     generic_args,
                     bounds_impls: _,
-                    trait_: _,
+                    trait_,
                 } => {
                     match (&args[..], &generic_args[..], head.kind()) {
                         ([arg], [], ExprKind::GlobalId(LIFT)) => docs![reflow!("← "), arg].parens(),
@@ -497,6 +578,9 @@ set_option linter.unusedVariables false
                             // Fallback for any application
                             docs![
                                 head,
+                                trait_
+                                    .as_ref()
+                                    .map(|(impl_expr, _)| zip_left!(line!(), &impl_expr.goal.args)),
                                 zip_left!(line!(), generic_args).group(),
                                 zip_left!(line!(), args).group(),
                             ]
@@ -998,27 +1082,27 @@ set_option linter.unusedVariables false
                     generics,
                     items,
                 } => {
-                    // Type parameters are also parameters of the class, but constraints are fields of the class
                     docs![
                         docs![
-                            docs![reflow!("class "), name],
-                            (!generics.params.is_empty()).then_some(docs![
-                                line!(),
-                                intersperse!(&generics.params, line!()).group()
-                            ]),
+                            reflow!("class "),
+                            name,
+                            // Trait-level type parameters are also parameters of the class (but
+                            // constraints are fields of the class)
+                            zip_left!(line!(), &generics.params),
                             line!(),
                             "where"
                         ]
                         .group(),
-                        hardline!(),
-                        (!generics.constraints.is_empty()).then_some(docs![zip_right!(
+                        zip_left!(
+                            hardline!(),
+                            // Trait-level constraints are extra fields of the class
                             generics
                                 .constraints
                                 .iter()
                                 .map(|constraint: &GenericConstraint| {
                                     match constraint {
                                         GenericConstraint::Type(tc_constraint) => docs![
-                                            self.fresh_constraint_name(&self.render_last(name), tc_constraint),
+                                            self.constraint_name(&self.render_last(name), tc_constraint),
                                             " :",
                                             line!(),
                                             constraint
@@ -1028,17 +1112,14 @@ set_option linter.unusedVariables false
                                         GenericConstraint::Lifetime(_) => unreachable_by_invariant!(Drop_references),
                                         GenericConstraint::Projection(_) => emit_error!(issue 1710, "Unsupported equality constraints on associated types"),
                                     }
-                                }),
-                            hardline!()
-                        )]),
-                        intersperse!(
+                                })),
+                        zip_left!(
+                            hardline!(),
                             items.iter().filter(|item| {
                                 // TODO: should be treated directly by name rendering, see :
                                 // https://github.com/cryspen/hax/issues/1646
-                                !(item.ident.is_precondition() || item.ident.is_postcondition())
-                            }),
-                            hardline!()
-                        )
+                                 !(item.ident.is_precondition() || item.ident.is_postcondition())
+                            }).map(|item| docs![(generics.params.clone(), item)] )),
                     ]
                     .nest(INDENT)
                 }
@@ -1103,62 +1184,6 @@ set_option linter.unusedVariables false
             docs![meta, body]
         }
 
-        fn trait_item(
-            &self,
-            TraitItem {
-                meta: _,
-                kind,
-                generics,
-                ident,
-            }: &TraitItem,
-        ) -> DocBuilder<A> {
-            let name = self.render_last(ident);
-            docs![match kind {
-                TraitItemKind::Fn(ty) => {
-                    docs![name, softline!(), generics, ":", line!(), ty]
-                        .group()
-                        .nest(INDENT)
-                }
-                TraitItemKind::Type(constraints) => {
-                    docs![
-                        name.clone(),
-                        reflow!(" : Type"),
-                        concat!(constraints.iter().map(|c| docs![
-                            hardline!(),
-                            docs![
-                                self.fresh_constraint_name(&name, c),
-                                reflow!(" :"),
-                                line!(),
-                                &c.goal
-                            ]
-                            .group()
-                            .nest(INDENT)
-                            .brackets()
-                        ]))
-                    ]
-                }
-                TraitItemKind::Default { params, body } => docs![
-                    docs![
-                        name,
-                        softline!(),
-                        generics,
-                        zip_right!(params, line!()).group(),
-                        docs![": RustM ", body.ty].group(),
-                        line!(),
-                        ":= do",
-                    ]
-                    .group(),
-                    line!(),
-                    body,
-                ]
-                .group()
-                .nest(INDENT),
-                TraitItemKind::Resugared(_) => {
-                    unreachable!("This backend has no resugaring for trait items")
-                }
-            }]
-        }
-
         fn impl_item(
             &self,
             ImplItem {