Make Self clause explicit in trait item declarations

charon-ml/src/generated/Generated_Types.ml

@@ -104,7 +104,6 @@ and 'a0 de_bruijn_var =
           is not used in charon internals, only as a micro-pass before exporting the crate data.
        *)
 
-and region_id = (RegionId.id[@visitors.opaque])
 and type_var_id = (TypeVarId.id[@visitors.opaque])
 and const_generic_var_id = (ConstGenericVarId.id[@visitors.opaque])
 and trait_clause_id = (TraitClauseId.id[@visitors.opaque])
@@ -201,6 +200,7 @@ and global_decl_ref = {
 }
 
 and trait_item_name = string
+and region_id = (RegionId.id[@visitors.opaque])
 
 (** A region variable in a signature or binder. *)
 and region_var = (region_id, string option) indexed_var

charon/src/ast/krate.rs

@@ -221,6 +221,17 @@ impl<'ctx> AnyTransItem<'ctx> {
         }
     }
 
+    /// Get information about the parent of this item, if any.
+    pub fn parent_info(&self) -> &'ctx ItemKind {
+        match self {
+            AnyTransItem::Fun(d) => &d.kind,
+            AnyTransItem::Global(d) => &d.kind,
+            AnyTransItem::Type(_) | AnyTransItem::TraitDecl(_) | AnyTransItem::TraitImpl(_) => {
+                &ItemKind::Regular
+            }
+        }
+    }
+
     /// See [`GenericParams::identity_args`].
     pub fn identity_args(&self) -> GenericArgs {
         self.generic_params()
@@ -238,6 +249,17 @@ impl<'ctx> AnyTransItem<'ctx> {
             AnyTransItem::TraitImpl(d) => visitor.visit(d),
         }
     }
+
+    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
+    pub fn dyn_visit<T: AstVisitable>(&self, f: impl FnMut(&T)) {
+        match *self {
+            AnyTransItem::Type(d) => d.dyn_visit(f),
+            AnyTransItem::Fun(d) => d.dyn_visit(f),
+            AnyTransItem::Global(d) => d.dyn_visit(f),
+            AnyTransItem::TraitDecl(d) => d.dyn_visit(f),
+            AnyTransItem::TraitImpl(d) => d.dyn_visit(f),
+        }
+    }
 }
 
 impl<'ctx> AnyTransItemMut<'ctx> {
@@ -273,6 +295,17 @@ impl<'ctx> AnyTransItemMut<'ctx> {
             AnyTransItemMut::TraitImpl(d) => visitor.visit(*d),
         }
     }
+
+    /// Visit all occurrences of that type inside `self`, in pre-order traversal.
+    pub fn dyn_visit_mut<T: AstVisitable>(&mut self, f: impl FnMut(&mut T)) {
+        match self {
+            AnyTransItemMut::Type(d) => d.dyn_visit_mut(f),
+            AnyTransItemMut::Fun(d) => d.dyn_visit_mut(f),
+            AnyTransItemMut::Global(d) => d.dyn_visit_mut(f),
+            AnyTransItemMut::TraitDecl(d) => d.dyn_visit_mut(f),
+            AnyTransItemMut::TraitImpl(d) => d.dyn_visit_mut(f),
+        }
+    }
 }
 
 impl fmt::Display for TranslatedCrate {

charon/src/ast/visitor.rs

@@ -47,7 +47,7 @@ use indexmap::IndexMap;
         FnOperand, FunId, FunIdOrTraitMethodRef, FunSig, ImplElem, IntegerTy, Literal, LiteralTy,
         llbc_ast::Block, llbc_ast::ExprBody, llbc_ast::RawStatement, llbc_ast::Switch,
         Locals, Name, NullOp, Opaque, Operand, PathElem, Place, PlaceKind, ProjectionElem, RawConstantExpr,
-        RefKind, RegionId, RegionVar, Rvalue, ScalarValue, TraitClauseId, TraitItemName,
+        RefKind, RegionId, RegionVar, Rvalue, ScalarValue, TraitItemName,
         TranslatedCrate, TypeDeclKind, TypeId, TypeVar, TypeVarId,
         ullbc_ast::BlockData, ullbc_ast::BlockId, ullbc_ast::ExprBody, ullbc_ast::RawStatement,
         ullbc_ast::RawTerminator, ullbc_ast::SwitchTargets, ullbc_ast::Terminator,
@@ -65,7 +65,7 @@ use indexmap::IndexMap;
     override(
         DeBruijnId, Ty, TyKind, Region, ConstGeneric, TraitRef, TraitRefKind,
         FunDeclRef, GlobalDeclRef, TraitDeclRef, TraitImplRef,
-        GenericArgs, GenericParams, TraitClause, TraitTypeConstraint,
+        GenericArgs, GenericParams, TraitClause, TraitClauseId, TraitTypeConstraint,
         for<T: AstVisitable + Idx> DeBruijnVar<T>,
         for<T: AstVisitable> RegionBinder<T>,
         for<T: AstVisitable> Binder<T>,

charon/src/bin/charon-driver/translate/translate_ctx.rs

@@ -248,6 +248,9 @@ pub(crate) struct BodyTransCtx<'tcx, 'ctx> {
     pub parent_trait_clauses: Vector<TraitClauseId, TraitClause>,
     /// (For traits only) accumulated trait clauses on associated types.
     pub item_trait_clauses: HashMap<TraitItemName, Vector<TraitClauseId, TraitClause>>,
+    /// (For method declarations only) the clause id corresponding to the explicit `Self` clause.
+    /// If `None`, use `TraitRefKind::Self` instead.
+    pub self_clause_id: Option<TraitClauseId>,
 
     /// The (regular) variables in the current function body.
     pub locals: Locals,
@@ -908,6 +911,7 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
             binding_levels: Default::default(),
             parent_trait_clauses: Default::default(),
             item_trait_clauses: Default::default(),
+            self_clause_id: Default::default(),
             locals: Default::default(),
             vars_map: Default::default(),
             blocks: Default::default(),
@@ -1099,6 +1103,12 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
         span: Span,
         mut id: usize,
     ) -> Result<ClauseDbVar, Error> {
+        if self.self_clause_id.is_some() {
+            // We added an extra first clause which hax doesn't know about, so we adapt the index
+            // accordingly.
+            // TODO: more robust tracking of clause ids between hax and charon.
+            id += 1;
+        }
         // The clause indices returned by hax count clauses in order, starting from the parentmost.
         // While adding clauses to a binding level we already need to translate types and clauses,
         // so the innermost item binder may not have all the clauses yet. Hence for that binder we

charon/src/bin/charon-driver/translate/translate_predicates.rs

@@ -245,10 +245,17 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
                     trait_ref,
                     path,
                 );
-                // If we are refering to a trait clause, we need to find the
-                // relevant one.
+                // If we are refering to a trait clause, we need to find the relevant one.
                 let mut trait_id = match &impl_source.r#impl {
-                    ImplExprAtom::SelfImpl { .. } => TraitRefKind::SelfId,
+                    ImplExprAtom::SelfImpl { .. } => match self.self_clause_id {
+                        None => TraitRefKind::SelfId,
+                        // An explicit `Self` clause is bound at the top-level; we use it instead
+                        // of the implicit `TraitRefKind::SelfId` one.
+                        Some(id) => TraitRefKind::Clause(DeBruijnVar::bound(
+                            DeBruijnId::new(self.binding_levels.len() - 1),
+                            id,
+                        )),
+                    },
                     ImplExprAtom::LocalBound { index, .. } => {
                         let var = self.lookup_clause_var(span, *index)?;
                         TraitRefKind::Clause(var)

charon/src/bin/charon-driver/translate/translate_traits.rs

@@ -54,7 +54,12 @@ impl BodyTransCtx<'_, '_> {
             raise_error!(self, span, "Trait aliases are not supported");
         }
 
-        let hax::FullDefKind::Trait { items, .. } = &def.kind else {
+        let hax::FullDefKind::Trait {
+            items,
+            self_predicate,
+            ..
+        } = &def.kind
+        else {
             raise_error!(self, span, "Unexpected definition: {def:?}");
         };
         let items: Vec<(TraitItemName, &hax::AssocItem, Arc<hax::FullDef>)> = items
@@ -71,6 +76,12 @@ impl BodyTransCtx<'_, '_> {
         // `self.parent_trait_clauses`.
         // TODO: Distinguish between required and implied trait clauses?
         self.translate_def_generics(span, def)?;
+        let self_trait_ref = TraitRef {
+            kind: TraitRefKind::SelfId,
+            trait_decl_ref: RegionBinder::empty(
+                self.translate_trait_predicate(span, self_predicate)?,
+            ),
+        };
 
         // Translate the associated items
         // We do something subtle here: TODO: explain
@@ -108,22 +119,23 @@ impl BodyTransCtx<'_, '_> {
                                 bt_ctx.register_fun_decl_id_no_enqueue(item_span, item_def_id)
                             };
 
-                            // TODO: there's probably a cleaner way to write this
                             assert_eq!(bt_ctx.binding_levels.len(), 2);
-                            let fun_generics = bt_ctx
-                                .outermost_binder()
-                                .params
-                                .identity_args_at_depth(
+                            let mut fun_generics =
+                                bt_ctx.outermost_binder().params.identity_args_at_depth(
                                     GenericsSource::item(def_id),
                                     DeBruijnId::one(),
-                                )
-                                .concat(
-                                    GenericsSource::item(fun_id),
-                                    &bt_ctx.innermost_binder().params.identity_args_at_depth(
-                                        GenericsSource::Method(def_id.into(), item_name.clone()),
-                                        DeBruijnId::zero(),
-                                    ),
                                 );
+                            // Add the necessary explicit `Self` clause at the start.
+                            fun_generics
+                                .trait_refs
+                                .insert(0.into(), self_trait_ref.clone().move_under_binder());
+                            fun_generics = fun_generics.concat(
+                                GenericsSource::item(fun_id),
+                                &bt_ctx.innermost_binder().params.identity_args_at_depth(
+                                    GenericsSource::Method(def_id.into(), item_name.clone()),
+                                    DeBruijnId::zero(),
+                                ),
+                            );
                             Ok(FunDeclRef {
                                 id: fun_id,
                                 generics: fun_generics,
@@ -139,10 +151,10 @@ impl BodyTransCtx<'_, '_> {
                         // declares them.
                         let id = self.register_global_decl_id(item_span, item_def_id);
                         let generics_target = GenericsSource::item(id);
-                        let gref = GlobalDeclRef {
-                            id,
-                            generics: self.the_only_binder().params.identity_args(generics_target),
-                        };
+                        let mut generics =
+                            self.the_only_binder().params.identity_args(generics_target);
+                        generics.trait_refs.push(self_trait_ref.clone());
+                        let gref = GlobalDeclRef { id, generics };
                         const_defaults.insert(item_name.clone(), gref);
                     };
                     let ty = self.translate_ty(item_span, ty)?;
@@ -226,6 +238,16 @@ impl BodyTransCtx<'_, '_> {
             )?;
             TraitDeclRef { trait_id, generics }
         };
+        // A `TraitRef` that points to this impl with the correct generics.
+        let self_predicate = TraitRef {
+            kind: TraitRefKind::TraitImpl(
+                def_id,
+                self.the_only_binder()
+                    .params
+                    .identity_args(GenericsSource::item(def_id)),
+            ),
+            trait_decl_ref: RegionBinder::empty(implemented_trait.clone()),
+        };
 
         // The trait refs which implement the parent clauses of the implemented trait decl.
         let parent_trait_refs = self.translate_trait_impl_exprs(span, &required_impl_exprs)?;
@@ -325,10 +347,14 @@ impl BodyTransCtx<'_, '_> {
                         Provided { .. } => {
                             self.the_only_binder().params.identity_args(generics_target)
                         }
-                        _ => implemented_trait
-                            .generics
-                            .clone()
-                            .with_target(generics_target),
+                        _ => {
+                            let mut generics = implemented_trait
+                                .generics
+                                .clone()
+                                .with_target(generics_target);
+                            generics.trait_refs.push(self_predicate.clone());
+                            generics
+                        }
                     };
                     let gref = GlobalDeclRef { id, generics };
                     consts.push((name, gref));

charon/src/bin/charon-driver/translate/translate_types.rs

@@ -549,7 +549,7 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
     ) -> Result<(), Error> {
         assert!(self.binding_levels.len() == 0);
         self.binding_levels.push(BindingLevel::new(true));
-        self.push_generics_for_def(span, def, false)?;
+        self.push_generics_for_def(span, def, false, false)?;
         self.innermost_binder_mut().params.check_consistency();
         Ok(())
     }
@@ -561,7 +561,7 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
         def: &hax::FullDef,
     ) -> Result<(), Error> {
         self.binding_levels.push(BindingLevel::new(true));
-        self.push_generics_for_def_without_parents(span, def, true, true)?;
+        self.push_generics_for_def_without_parents(span, def, true, true, false)?;
         self.innermost_binder().params.check_consistency();
         Ok(())
     }
@@ -577,24 +577,33 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
         span: Span,
         def: &hax::FullDef,
         is_parent: bool,
+        explicit_self_clause: bool,
     ) -> Result<(), Error> {
-        use hax::FullDefKind;
+        use hax::FullDefKind::*;
         // Add generics from the parent item, recursively (recursivity is important for closures,
-        // as they could be nested).
-        match &def.kind {
-            FullDefKind::AssocTy { .. }
-            | FullDefKind::AssocFn { .. }
-            | FullDefKind::AssocConst { .. }
-            | FullDefKind::Closure { .. }
-            | FullDefKind::Ctor { .. }
-            | FullDefKind::Variant { .. } => {
+        // as they can be nested).
+        match def.kind() {
+            AssocTy { .. }
+            | AssocFn { .. }
+            | AssocConst { .. }
+            | Closure { .. }
+            | Ctor { .. }
+            | Variant { .. } => {
                 let parent_def_id = def.parent.as_ref().unwrap();
                 let parent_def = self.t_ctx.hax_def(parent_def_id)?;
-                self.push_generics_for_def(span, &parent_def, true)?;
+                // Add an explicit `Self` clause to trait item declarations.
+                let explicit_self_clause = matches!(parent_def.kind(), Trait { .. });
+                self.push_generics_for_def(span, &parent_def, true, explicit_self_clause)?;
             }
             _ => {}
         }
-        self.push_generics_for_def_without_parents(span, def, !is_parent, !is_parent)?;
+        self.push_generics_for_def_without_parents(
+            span,
+            def,
+            !is_parent,
+            !is_parent,
+            explicit_self_clause,
+        )?;
         Ok(())
     }
 
@@ -606,22 +615,33 @@ impl<'tcx, 'ctx> BodyTransCtx<'tcx, 'ctx> {
         def: &hax::FullDef,
         include_late_bound: bool,
         include_assoc_ty_clauses: bool,
+        explicit_self_clause: bool,
     ) -> Result<(), Error> {
         use hax::FullDefKind;
         if let Some((generics, predicates)) = def.generics() {
             // Add the generic params.
             self.push_generic_params(generics)?;
-            // Add the self trait clause.
-            match &def.kind {
-                FullDefKind::TraitImpl {
-                    trait_pred: self_predicate,
-                    ..
-                }
-                | FullDefKind::Trait { self_predicate, .. } => {
-                    self.register_trait_decl_id(span, &self_predicate.trait_ref.def_id);
-                    let _ = self.translate_trait_predicate(span, self_predicate)?;
-                }
-                _ => {}
+            // Add the explicit self trait clause if required.
+            if let FullDefKind::Trait { self_predicate, .. } = &def.kind
+                && explicit_self_clause
+            {
+                // We add an explicit `Self` clause to trait method declarations. Trait method
+                // implementations already don't use the implicit `Self` clause. This way, methods
+                // don't need an implicit `Self` clause: they're normal functions, and the trait
+                // decl/impl takes care to pass the right arguments.
+                let self_predicate =
+                    RegionBinder::empty(self.translate_trait_predicate(span, self_predicate)?);
+                let clause_id =
+                    self.innermost_generics_mut()
+                        .trait_clauses
+                        .push_with(|clause_id| TraitClause {
+                            clause_id,
+                            origin: PredicateOrigin::TraitSelf,
+                            span: Some(span),
+                            trait_: self_predicate,
+                        });
+                // Record the id so we can resolve `ImplExpr::Self`s to it.
+                self.self_clause_id = Some(clause_id);
             }
             // Add the predicates.
             let (origin, location) = match &def.kind {

charon/src/bin/generate-ml/main.rs

@@ -1183,7 +1183,6 @@ fn generate_ml(
                     reduce: true,
                     extra_types: &[],
                 })), &[
-                    "RegionId",
                     "TypeVarId",
                     "ConstGeneric",
                     "TraitClauseId",