Auto merge of #14087 - HKalbasi:layout, r=Veykril

Support computing layout of RPIT

And some refactoring to make code more type safe.

8 files changed, 214 insertions, 65 deletions

diff --git a/crates/hir-ty/src/chalk_db.rs b/crates/hir-ty/src/chalk_db.rs
index bbb6625855d..6989e9fb9be 100644
--- a/crates/hir-ty/src/chalk_db.rs
+++ b/crates/hir-ty/src/chalk_db.rs
@@ -205,7 +205,7 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
                     .return_type_impl_traits(func)
                     .expect("impl trait id without impl traits");
                 let (datas, binders) = (*datas).as_ref().into_value_and_skipped_binders();
-                let data = &datas.impl_traits[idx as usize];
+                let data = &datas.impl_traits[idx];
                 let bound = OpaqueTyDatumBound {
                     bounds: make_single_type_binders(data.bounds.skip_binders().to_vec()),
                     where_clauses: chalk_ir::Binders::empty(Interner, vec![]),
diff --git a/crates/hir-ty/src/chalk_ext.rs b/crates/hir-ty/src/chalk_ext.rs
index 329c87c74e9..45c975dfcdc 100644
--- a/crates/hir-ty/src/chalk_ext.rs
+++ b/crates/hir-ty/src/chalk_ext.rs
@@ -234,9 +234,8 @@ impl TyExt for Ty {
                     }
                     ImplTraitId::ReturnTypeImplTrait(func, idx) => {
                         db.return_type_impl_traits(func).map(|it| {
-                            let data = (*it)
-                                .as_ref()
-                                .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                            let data =
+                                (*it).as_ref().map(|rpit| rpit.impl_traits[idx].bounds.clone());
                             data.substitute(Interner, &subst).into_value_and_skipped_binders().0
                         })
                     }
@@ -247,9 +246,8 @@ impl TyExt for Ty {
                 {
                     ImplTraitId::ReturnTypeImplTrait(func, idx) => {
                         db.return_type_impl_traits(func).map(|it| {
-                            let data = (*it)
-                                .as_ref()
-                                .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                            let data =
+                                (*it).as_ref().map(|rpit| rpit.impl_traits[idx].bounds.clone());
                             data.substitute(Interner, &opaque_ty.substitution)
                         })
                     }
diff --git a/crates/hir-ty/src/display.rs b/crates/hir-ty/src/display.rs
index ae2162dd7cd..464860dea3c 100644
--- a/crates/hir-ty/src/display.rs
+++ b/crates/hir-ty/src/display.rs
@@ -458,9 +458,8 @@ impl HirDisplay for Ty {
                             let datas = db
                                 .return_type_impl_traits(func)
                                 .expect("impl trait id without data");
-                            let data = (*datas)
-                                .as_ref()
-                                .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                            let data =
+                                (*datas).as_ref().map(|rpit| rpit.impl_traits[idx].bounds.clone());
                             let bounds = data.substitute(Interner, parameters);
                             let mut len = bounds.skip_binders().len();
 
@@ -718,9 +717,8 @@ impl HirDisplay for Ty {
                     ImplTraitId::ReturnTypeImplTrait(func, idx) => {
                         let datas =
                             db.return_type_impl_traits(func).expect("impl trait id without data");
-                        let data = (*datas)
-                            .as_ref()
-                            .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                        let data =
+                            (*datas).as_ref().map(|rpit| rpit.impl_traits[idx].bounds.clone());
                         let bounds = data.substitute(Interner, &parameters);
                         let krate = func.lookup(db.upcast()).module(db.upcast()).krate();
                         write_bounds_like_dyn_trait_with_prefix(
@@ -828,9 +826,8 @@ impl HirDisplay for Ty {
                     ImplTraitId::ReturnTypeImplTrait(func, idx) => {
                         let datas =
                             db.return_type_impl_traits(func).expect("impl trait id without data");
-                        let data = (*datas)
-                            .as_ref()
-                            .map(|rpit| rpit.impl_traits[idx as usize].bounds.clone());
+                        let data =
+                            (*datas).as_ref().map(|rpit| rpit.impl_traits[idx].bounds.clone());
                         let bounds = data.substitute(Interner, &opaque_ty.substitution);
                         let krate = func.lookup(db.upcast()).module(db.upcast()).krate();
                         write_bounds_like_dyn_trait_with_prefix(
diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 571b3e96863..a76d33c0cde 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -39,7 +39,7 @@ use stdx::always;
 use crate::{
     db::HirDatabase, fold_tys, fold_tys_and_consts, infer::coerce::CoerceMany,
     lower::ImplTraitLoweringMode, to_assoc_type_id, AliasEq, AliasTy, Const, DomainGoal,
-    GenericArg, Goal, ImplTraitId, InEnvironment, Interner, ProjectionTy, Substitution,
+    GenericArg, Goal, ImplTraitId, InEnvironment, Interner, ProjectionTy, RpitId, Substitution,
     TraitEnvironment, TraitRef, Ty, TyBuilder, TyExt, TyKind,
 };
 
@@ -352,6 +352,7 @@ pub struct InferenceResult {
     /// **Note**: When a pattern type is resolved it may still contain
     /// unresolved or missing subpatterns or subpatterns of mismatched types.
     pub type_of_pat: ArenaMap<PatId, Ty>,
+    pub type_of_rpit: ArenaMap<RpitId, Ty>,
     type_mismatches: FxHashMap<ExprOrPatId, TypeMismatch>,
     /// Interned common types to return references to.
     standard_types: InternedStandardTypes,
@@ -525,6 +526,9 @@ impl<'a> InferenceContext<'a> {
         for ty in result.type_of_pat.values_mut() {
             *ty = table.resolve_completely(ty.clone());
         }
+        for ty in result.type_of_rpit.iter_mut().map(|x| x.1) {
+            *ty = table.resolve_completely(ty.clone());
+        }
         for mismatch in result.type_mismatches.values_mut() {
             mismatch.expected = table.resolve_completely(mismatch.expected.clone());
             mismatch.actual = table.resolve_completely(mismatch.actual.clone());
@@ -603,7 +607,7 @@ impl<'a> InferenceContext<'a> {
                         _ => unreachable!(),
                     };
                     let bounds = (*rpits).map_ref(|rpits| {
-                        rpits.impl_traits[idx as usize].bounds.map_ref(|it| it.into_iter())
+                        rpits.impl_traits[idx].bounds.map_ref(|it| it.into_iter())
                     });
                     let var = self.table.new_type_var();
                     let var_subst = Substitution::from1(Interner, var.clone());
@@ -616,6 +620,7 @@ impl<'a> InferenceContext<'a> {
                         always!(binders.is_empty(Interner)); // quantified where clauses not yet handled
                         self.push_obligation(var_predicate.cast(Interner));
                     }
+                    self.result.type_of_rpit.insert(idx, var.clone());
                     var
                 },
                 DebruijnIndex::INNERMOST,
diff --git a/crates/hir-ty/src/layout.rs b/crates/hir-ty/src/layout.rs
index cd7f1b805c5..c82c274524a 100644
--- a/crates/hir-ty/src/layout.rs
+++ b/crates/hir-ty/src/layout.rs
@@ -225,10 +225,21 @@ pub fn layout_of_ty(db: &dyn HirDatabase, ty: &Ty, krate: CrateId) -> Result<Lay
             ptr.valid_range_mut().start = 1;
             Layout::scalar(dl, ptr)
         }
-        TyKind::Closure(_, _)
-        | TyKind::OpaqueType(_, _)
-        | TyKind::Generator(_, _)
-        | TyKind::GeneratorWitness(_, _) => return Err(LayoutError::NotImplemented),
+        TyKind::OpaqueType(opaque_ty_id, _) => {
+            let impl_trait_id = db.lookup_intern_impl_trait_id((*opaque_ty_id).into());
+            match impl_trait_id {
+                crate::ImplTraitId::ReturnTypeImplTrait(func, idx) => {
+                    let infer = db.infer(func.into());
+                    layout_of_ty(db, &infer.type_of_rpit[idx], krate)?
+                }
+                crate::ImplTraitId::AsyncBlockTypeImplTrait(_, _) => {
+                    return Err(LayoutError::NotImplemented)
+                }
+            }
+        }
+        TyKind::Closure(_, _) | TyKind::Generator(_, _) | TyKind::GeneratorWitness(_, _) => {
+            return Err(LayoutError::NotImplemented)
+        }
         TyKind::AssociatedType(_, _)
         | TyKind::Error
         | TyKind::Alias(_)
diff --git a/crates/hir-ty/src/layout/tests.rs b/crates/hir-ty/src/layout/tests.rs
index 53838cf41d2..ba7db9b0e20 100644
--- a/crates/hir-ty/src/layout/tests.rs
+++ b/crates/hir-ty/src/layout/tests.rs
@@ -5,7 +5,7 @@ use hir_def::{
     layout::{Layout, LayoutError},
 };
 
-use crate::{test_db::TestDB, Interner, Substitution};
+use crate::{db::HirDatabase, test_db::TestDB, Interner, Substitution};
 
 use super::layout_of_ty;
 
@@ -45,6 +45,50 @@ fn eval_goal(ra_fixture: &str, minicore: &str) -> Result<Layout, LayoutError> {
     layout_of_ty(&db, &goal_ty, module_id.krate())
 }
 
+/// A version of `eval_goal` for types that can not be expressed in ADTs, like closures and `impl Trait`
+fn eval_expr(ra_fixture: &str, minicore: &str) -> Result<Layout, LayoutError> {
+    // using unstable cargo features failed, fall back to using plain rustc
+    let mut cmd = std::process::Command::new("rustc");
+    cmd.args(["-Z", "unstable-options", "--print", "target-spec-json"]).env("RUSTC_BOOTSTRAP", "1");
+    let output = cmd.output().unwrap();
+    assert!(output.status.success(), "{}", output.status);
+    let stdout = String::from_utf8(output.stdout).unwrap();
+    let target_data_layout =
+        stdout.split_once(r#""data-layout": ""#).unwrap().1.split_once('"').unwrap().0.to_owned();
+
+    let ra_fixture = format!(
+        "{minicore}//- /main.rs crate:test target_data_layout:{target_data_layout}\nfn main(){{let goal = {{{ra_fixture}}};}}",
+    );
+
+    let (db, file_id) = TestDB::with_single_file(&ra_fixture);
+    let module_id = db.module_for_file(file_id);
+    let def_map = module_id.def_map(&db);
+    let scope = &def_map[module_id.local_id].scope;
+    let adt_id = scope
+        .declarations()
+        .find_map(|x| match x {
+            hir_def::ModuleDefId::FunctionId(x) => {
+                let name = db.function_data(x).name.to_smol_str();
+                (name == "main").then_some(x)
+            }
+            _ => None,
+        })
+        .unwrap();
+    let hir_body = db.body(adt_id.into());
+    let pat = hir_body
+        .pats
+        .iter()
+        .find(|x| match x.1 {
+            hir_def::expr::Pat::Bind { name, .. } => name.to_smol_str() == "goal",
+            _ => false,
+        })
+        .unwrap()
+        .0;
+    let infer = db.infer(adt_id.into());
+    let goal_ty = infer.type_of_pat[pat].clone();
+    layout_of_ty(&db, &goal_ty, module_id.krate())
+}
+
 #[track_caller]
 fn check_size_and_align(ra_fixture: &str, minicore: &str, size: u64, align: u64) {
     let l = eval_goal(ra_fixture, minicore).unwrap();
@@ -53,6 +97,13 @@ fn check_size_and_align(ra_fixture: &str, minicore: &str, size: u64, align: u64)
 }
 
 #[track_caller]
+fn check_size_and_align_expr(ra_fixture: &str, minicore: &str, size: u64, align: u64) {
+    let l = eval_expr(ra_fixture, minicore).unwrap();
+    assert_eq!(l.size.bytes(), size);
+    assert_eq!(l.align.abi.bytes(), align);
+}
+
+#[track_caller]
 fn check_fail(ra_fixture: &str, e: LayoutError) {
     let r = eval_goal(ra_fixture, "");
     assert_eq!(r, Err(e));
@@ -85,11 +136,31 @@ macro_rules! size_and_align {
     };
 }
 
+macro_rules! size_and_align_expr {
+    ($($t:tt)*) => {
+        {
+            #[allow(dead_code)]
+            {
+                let val = { $($t)* };
+                check_size_and_align_expr(
+                    stringify!($($t)*),
+                    "",
+                    ::std::mem::size_of_val(&val) as u64,
+                    ::std::mem::align_of_val(&val) as u64,
+                );
+            }
+        }
+    };
+}
+
 #[test]
 fn hello_world() {
     size_and_align! {
         struct Goal(i32);
     }
+    size_and_align_expr! {
+        2i32
+    }
 }
 
 #[test]
@@ -144,6 +215,40 @@ fn generic() {
 }
 
 #[test]
+fn return_position_impl_trait() {
+    size_and_align_expr! {
+        trait T {}
+        impl T for i32 {}
+        impl T for i64 {}
+        fn foo() -> impl T { 2i64 }
+        foo()
+    }
+    size_and_align_expr! {
+        trait T {}
+        impl T for i32 {}
+        impl T for i64 {}
+        fn foo() -> (impl T, impl T, impl T) { (2i64, 5i32, 7i32) }
+        foo()
+    }
+    size_and_align_expr! {
+        struct Foo<T>(T, T, (T, T));
+        trait T {}
+        impl T for Foo<i32> {}
+        impl T for Foo<i64> {}
+
+        fn foo() -> Foo<impl T> { Foo(
+            Foo(1i64, 2, (3, 4)),
+            Foo(5, 6, (7, 8)),
+            (
+                Foo(1i64, 2, (3, 4)),
+                Foo(5, 6, (7, 8)),
+            ),
+        ) }
+        foo()
+    }
+}
+
+#[test]
 fn enums() {
     size_and_align! {
         enum Goal {
diff --git a/crates/hir-ty/src/lib.rs b/crates/hir-ty/src/lib.rs
index cbeb61067df..59a5ef8c14d 100644
--- a/crates/hir-ty/src/lib.rs
+++ b/crates/hir-ty/src/lib.rs
@@ -45,6 +45,7 @@ use chalk_ir::{
 use hir_def::{expr::ExprId, type_ref::Rawness, TypeOrConstParamId};
 use hir_expand::name;
 use itertools::Either;
+use la_arena::{Arena, Idx};
 use rustc_hash::FxHashSet;
 use traits::FnTrait;
 use utils::Generics;
@@ -290,22 +291,24 @@ impl TypeFoldable<Interner> for CallableSig {
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
 pub enum ImplTraitId {
-    ReturnTypeImplTrait(hir_def::FunctionId, u16),
+    ReturnTypeImplTrait(hir_def::FunctionId, RpitId),
     AsyncBlockTypeImplTrait(hir_def::DefWithBodyId, ExprId),
 }
 
 #[derive(Clone, PartialEq, Eq, Debug, Hash)]
 pub struct ReturnTypeImplTraits {
-    pub(crate) impl_traits: Vec<ReturnTypeImplTrait>,
+    pub(crate) impl_traits: Arena<ReturnTypeImplTrait>,
 }
 
 has_interner!(ReturnTypeImplTraits);
 
 #[derive(Clone, PartialEq, Eq, Debug, Hash)]
-pub(crate) struct ReturnTypeImplTrait {
+pub struct ReturnTypeImplTrait {
     pub(crate) bounds: Binders<Vec<QuantifiedWhereClause>>,
 }
 
+pub type RpitId = Idx<ReturnTypeImplTrait>;
+
 pub fn static_lifetime() -> Lifetime {
     LifetimeData::Static.intern(Interner)
 }
diff --git a/crates/hir-ty/src/lower.rs b/crates/hir-ty/src/lower.rs
index b1a7ad3e940..86abe1af68a 100644
--- a/crates/hir-ty/src/lower.rs
+++ b/crates/hir-ty/src/lower.rs
@@ -36,7 +36,7 @@ use hir_def::{
 use hir_expand::{name::Name, ExpandResult};
 use intern::Interned;
 use itertools::Either;
-use la_arena::ArenaMap;
+use la_arena::{Arena, ArenaMap};
 use rustc_hash::FxHashSet;
 use smallvec::SmallVec;
 use stdx::{impl_from, never};
@@ -58,6 +58,51 @@ use crate::{
 };
 
 #[derive(Debug)]
+enum ImplTraitLoweringState {
+    /// When turning `impl Trait` into opaque types, we have to collect the
+    /// bounds at the same time to get the IDs correct (without becoming too
+    /// complicated). I don't like using interior mutability (as for the
+    /// counter), but I've tried and failed to make the lifetimes work for
+    /// passing around a `&mut TyLoweringContext`. The core problem is that
+    /// we're grouping the mutable data (the counter and this field) together
+    /// with the immutable context (the references to the DB and resolver).
+    /// Splitting this up would be a possible fix.
+    Opaque(RefCell<Arena<ReturnTypeImplTrait>>),
+    Param(Cell<u16>),
+    Variable(Cell<u16>),
+    Disallowed,
+}
+impl ImplTraitLoweringState {
+    fn new(impl_trait_mode: ImplTraitLoweringMode) -> ImplTraitLoweringState {
+        match impl_trait_mode {
+            ImplTraitLoweringMode::Opaque => Self::Opaque(RefCell::new(Arena::new())),
+            ImplTraitLoweringMode::Param => Self::Param(Cell::new(0)),
+            ImplTraitLoweringMode::Variable => Self::Variable(Cell::new(0)),
+            ImplTraitLoweringMode::Disallowed => Self::Disallowed,
+        }
+    }
+
+    fn take(&self) -> Self {
+        match self {
+            Self::Opaque(x) => Self::Opaque(RefCell::new(x.take())),
+            Self::Param(x) => Self::Param(Cell::new(x.get())),
+            Self::Variable(x) => Self::Variable(Cell::new(x.get())),
+            Self::Disallowed => Self::Disallowed,
+        }
+    }
+
+    fn swap(&self, impl_trait_mode: &Self) {
+        match (self, impl_trait_mode) {
+            (Self::Opaque(x), Self::Opaque(y)) => x.swap(y),
+            (Self::Param(x), Self::Param(y)) => x.swap(y),
+            (Self::Variable(x), Self::Variable(y)) => x.swap(y),
+            (Self::Disallowed, Self::Disallowed) => (),
+            _ => panic!("mismatched lowering mode"),
+        }
+    }
+}
+
+#[derive(Debug)]
 pub struct TyLoweringContext<'a> {
     pub db: &'a dyn HirDatabase,
     pub resolver: &'a Resolver,
@@ -67,17 +112,7 @@ pub struct TyLoweringContext<'a> {
     /// should be converted to variables. I think in practice, this isn't
     /// possible currently, so this should be fine for now.
     pub type_param_mode: ParamLoweringMode,
-    pub impl_trait_mode: ImplTraitLoweringMode,
-    impl_trait_counter: Cell<u16>,
-    /// When turning `impl Trait` into opaque types, we have to collect the
-    /// bounds at the same time to get the IDs correct (without becoming too
-    /// complicated). I don't like using interior mutability (as for the
-    /// counter), but I've tried and failed to make the lifetimes work for
-    /// passing around a `&mut TyLoweringContext`. The core problem is that
-    /// we're grouping the mutable data (the counter and this field) together
-    /// with the immutable context (the references to the DB and resolver).
-    /// Splitting this up would be a possible fix.
-    opaque_type_data: RefCell<Vec<ReturnTypeImplTrait>>,
+    impl_trait_mode: ImplTraitLoweringState,
     expander: RefCell<Option<Expander>>,
     /// Tracks types with explicit `?Sized` bounds.
     pub(crate) unsized_types: RefCell<FxHashSet<Ty>>,
@@ -85,19 +120,15 @@ pub struct TyLoweringContext<'a> {
 
 impl<'a> TyLoweringContext<'a> {
     pub fn new(db: &'a dyn HirDatabase, resolver: &'a Resolver) -> Self {
-        let impl_trait_counter = Cell::new(0);
-        let impl_trait_mode = ImplTraitLoweringMode::Disallowed;
+        let impl_trait_mode = ImplTraitLoweringState::Disallowed;
         let type_param_mode = ParamLoweringMode::Placeholder;
         let in_binders = DebruijnIndex::INNERMOST;
-        let opaque_type_data = RefCell::new(Vec::new());
         Self {
             db,
             resolver,
             in_binders,
             impl_trait_mode,
-            impl_trait_counter,
             type_param_mode,
-            opaque_type_data,
             expander: RefCell::new(None),
             unsized_types: RefCell::default(),
         }
@@ -108,20 +139,18 @@ impl<'a> TyLoweringContext<'a> {
         debruijn: DebruijnIndex,
         f: impl FnOnce(&TyLoweringContext<'_>) -> T,
     ) -> T {
-        let opaque_ty_data_vec = self.opaque_type_data.take();
+        let impl_trait_mode = self.impl_trait_mode.take();
         let expander = self.expander.take();
         let unsized_types = self.unsized_types.take();
         let new_ctx = Self {
             in_binders: debruijn,
-            impl_trait_counter: Cell::new(self.impl_trait_counter.get()),
-            opaque_type_data: RefCell::new(opaque_ty_data_vec),
+            impl_trait_mode,
             expander: RefCell::new(expander),
             unsized_types: RefCell::new(unsized_types),
             ..*self
         };
         let result = f(&new_ctx);
-        self.impl_trait_counter.set(new_ctx.impl_trait_counter.get());
-        self.opaque_type_data.replace(new_ctx.opaque_type_data.into_inner());
+        self.impl_trait_mode.swap(&new_ctx.impl_trait_mode);
         self.expander.replace(new_ctx.expander.into_inner());
         self.unsized_types.replace(new_ctx.unsized_types.into_inner());
         result
@@ -136,7 +165,7 @@ impl<'a> TyLoweringContext<'a> {
     }
 
     pub fn with_impl_trait_mode(self, impl_trait_mode: ImplTraitLoweringMode) -> Self {
-        Self { impl_trait_mode, ..self }
+        Self { impl_trait_mode: ImplTraitLoweringState::new(impl_trait_mode), ..self }
     }
 
     pub fn with_type_param_mode(self, type_param_mode: ParamLoweringMode) -> Self {
@@ -244,20 +273,17 @@ impl<'a> TyLoweringContext<'a> {
             }
             TypeRef::DynTrait(bounds) => self.lower_dyn_trait(bounds),
             TypeRef::ImplTrait(bounds) => {
-                match self.impl_trait_mode {
-                    ImplTraitLoweringMode::Opaque => {
-                        let idx = self.impl_trait_counter.get();
-                        self.impl_trait_counter.set(idx + 1);
+                match &self.impl_trait_mode {
+                    ImplTraitLoweringState::Opaque(opaque_type_data) => {
                         let func = match self.resolver.generic_def() {
                             Some(GenericDefId::FunctionId(f)) => f,
                             _ => panic!("opaque impl trait lowering in non-function"),
                         };
 
-                        assert!(idx as usize == self.opaque_type_data.borrow().len());
                         // this dance is to make sure the data is in the right
                         // place even if we encounter more opaque types while
                         // lowering the bounds
-                        self.opaque_type_data.borrow_mut().push(ReturnTypeImplTrait {
+                        let idx = opaque_type_data.borrow_mut().alloc(ReturnTypeImplTrait {
                             bounds: crate::make_single_type_binders(Vec::new()),
                         });
                         // We don't want to lower the bounds inside the binders
@@ -273,7 +299,7 @@ impl<'a> TyLoweringContext<'a> {
                             .with_debruijn(DebruijnIndex::INNERMOST, |ctx| {
                                 ctx.lower_impl_trait(bounds, func)
                             });
-                        self.opaque_type_data.borrow_mut()[idx as usize] = actual_opaque_type_data;
+                        opaque_type_data.borrow_mut()[idx] = actual_opaque_type_data;
 
                         let impl_trait_id = ImplTraitId::ReturnTypeImplTrait(func, idx);
                         let opaque_ty_id = self.db.intern_impl_trait_id(impl_trait_id).into();
@@ -281,10 +307,10 @@ impl<'a> TyLoweringContext<'a> {
                         let parameters = generics.bound_vars_subst(self.db, self.in_binders);
                         TyKind::OpaqueType(opaque_ty_id, parameters).intern(Interner)
                     }
-                    ImplTraitLoweringMode::Param => {
-                        let idx = self.impl_trait_counter.get();
+                    ImplTraitLoweringState::Param(counter) => {
+                        let idx = counter.get();
                         // FIXME we're probably doing something wrong here
-                        self.impl_trait_counter.set(idx + count_impl_traits(type_ref) as u16);
+                        counter.set(idx + count_impl_traits(type_ref) as u16);
                         if let Some(def) = self.resolver.generic_def() {
                             let generics = generics(self.db.upcast(), def);
                             let param = generics
@@ -305,10 +331,10 @@ impl<'a> TyLoweringContext<'a> {
                             TyKind::Error.intern(Interner)
                         }
                     }
-                    ImplTraitLoweringMode::Variable => {
-                        let idx = self.impl_trait_counter.get();
+                    ImplTraitLoweringState::Variable(counter) => {
+                        let idx = counter.get();
                         // FIXME we're probably doing something wrong here
-                        self.impl_trait_counter.set(idx + count_impl_traits(type_ref) as u16);
+                        counter.set(idx + count_impl_traits(type_ref) as u16);
                         let (
                             _parent_params,
                             self_params,
@@ -327,7 +353,7 @@ impl<'a> TyLoweringContext<'a> {
                         ))
                         .intern(Interner)
                     }
-                    ImplTraitLoweringMode::Disallowed => {
+                    ImplTraitLoweringState::Disallowed => {
                         // FIXME: report error
                         TyKind::Error.intern(Interner)
                     }
@@ -1863,8 +1889,12 @@ pub(crate) fn return_type_impl_traits(
         .with_type_param_mode(ParamLoweringMode::Variable);
     let _ret = ctx_ret.lower_ty(&data.ret_type);
     let generics = generics(db.upcast(), def.into());
-    let return_type_impl_traits =
-        ReturnTypeImplTraits { impl_traits: ctx_ret.opaque_type_data.into_inner() };
+    let return_type_impl_traits = ReturnTypeImplTraits {
+        impl_traits: match ctx_ret.impl_trait_mode {
+            ImplTraitLoweringState::Opaque(x) => x.into_inner(),
+            _ => unreachable!(),
+        },
+    };
     if return_type_impl_traits.impl_traits.is_empty() {
         None
     } else {