Implement return position impl trait / opaque type support

This is working, but I'm not that happy with how the lowering works. We might
need an additional representation between `TypeRef` and `Ty` where names are
resolved and `impl Trait` bounds are separated out, but things like inference
variables don't exist and `impl Trait` is always represented the same
way.

Also note that this doesn't implement correct handling of RPIT *inside* the
function (which involves turning the `impl Trait`s into variables and creating
obligations for them). That intermediate representation might help there as
well.

1 files changed, 122 insertions, 21 deletions

diff --git a/crates/ra_hir_ty/src/lower.rs b/crates/ra_hir_ty/src/lower.rs
index 35ac86a46..dfc018b0b 100644
--- a/crates/ra_hir_ty/src/lower.rs
+++ b/crates/ra_hir_ty/src/lower.rs
@@ -31,8 +31,9 @@ use crate::{
         all_super_trait_refs, associated_type_by_name_including_super_traits, generics,
         make_mut_slice, variant_data,
     },
-    Binders, BoundVar, DebruijnIndex, FnSig, GenericPredicate, PolyFnSig, ProjectionPredicate,
-    ProjectionTy, Substs, TraitEnvironment, TraitRef, Ty, TypeCtor, TypeWalk,
+    Binders, BoundVar, DebruijnIndex, FnSig, GenericPredicate, OpaqueTy, OpaqueTyId, PolyFnSig,
+    ProjectionPredicate, ProjectionTy, ReturnTypeImplTrait, ReturnTypeImplTraits, Substs,
+    TraitEnvironment, TraitRef, Ty, TypeCtor, TypeWalk,
 };
 use hir_expand::name::Name;
 
@@ -47,7 +48,16 @@ pub struct TyLoweringContext<'a> {
     /// possible currently, so this should be fine for now.
     pub type_param_mode: TypeParamLoweringMode,
     pub impl_trait_mode: ImplTraitLoweringMode,
-    pub impl_trait_counter: std::cell::Cell<u16>,
+    impl_trait_counter: std::cell::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: std::cell::RefCell<Vec<ReturnTypeImplTrait>>,
 }
 
 impl<'a> TyLoweringContext<'a> {
@@ -56,26 +66,42 @@ impl<'a> TyLoweringContext<'a> {
         let impl_trait_mode = ImplTraitLoweringMode::Disallowed;
         let type_param_mode = TypeParamLoweringMode::Placeholder;
         let in_binders = DebruijnIndex::INNERMOST;
-        Self { db, resolver, in_binders, impl_trait_mode, impl_trait_counter, type_param_mode }
+        let opaque_type_data = std::cell::RefCell::new(Vec::new());
+        Self {
+            db,
+            resolver,
+            in_binders,
+            impl_trait_mode,
+            impl_trait_counter,
+            type_param_mode,
+            opaque_type_data,
+        }
     }
 
-    pub fn with_shifted_in<T>(
+    pub fn with_debruijn<T>(
         &self,
         debruijn: DebruijnIndex,
         f: impl FnOnce(&TyLoweringContext) -> T,
     ) -> T {
+        let opaque_ty_data_vec = self.opaque_type_data.replace(Vec::new());
         let new_ctx = Self {
-            in_binders: self.in_binders.shifted_in_from(debruijn),
+            in_binders: debruijn,
             impl_trait_counter: std::cell::Cell::new(self.impl_trait_counter.get()),
+            opaque_type_data: std::cell::RefCell::new(opaque_ty_data_vec),
             ..*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());
         result
     }
 
-    pub fn shifted_in(self, debruijn: DebruijnIndex) -> Self {
-        Self { in_binders: self.in_binders.shifted_in_from(debruijn), ..self }
+    pub fn with_shifted_in<T>(
+        &self,
+        debruijn: DebruijnIndex,
+        f: impl FnOnce(&TyLoweringContext) -> T,
+    ) -> T {
+        self.with_debruijn(self.in_binders.shifted_in_from(debruijn), f)
     }
 
     pub fn with_impl_trait_mode(self, impl_trait_mode: ImplTraitLoweringMode) -> Self {
@@ -167,20 +193,47 @@ impl Ty {
             TypeRef::ImplTrait(bounds) => {
                 match ctx.impl_trait_mode {
                     ImplTraitLoweringMode::Opaque => {
-                        let self_ty = Ty::Bound(BoundVar::new(DebruijnIndex::INNERMOST, 0));
-                        let predicates = ctx.with_shifted_in(DebruijnIndex::ONE, |ctx| {
-                            bounds
-                                .iter()
-                                .flat_map(|b| {
-                                    GenericPredicate::from_type_bound(ctx, b, self_ty.clone())
-                                })
-                                .collect()
-                        });
-                        Ty::Opaque(predicates)
+                        let idx = ctx.impl_trait_counter.get();
+                        ctx.impl_trait_counter.set(idx + 1);
+
+                        assert!(idx as usize == ctx.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
+                        ctx.opaque_type_data
+                            .borrow_mut()
+                            .push(ReturnTypeImplTrait { bounds: Binders::new(1, Vec::new()) });
+                        // We don't want to lower the bounds inside the binders
+                        // we're currently in, because they don't end up inside
+                        // those binders. E.g. when we have `impl Trait<impl
+                        // OtherTrait<T>>`, the `impl OtherTrait<T>` can't refer
+                        // to the self parameter from `impl Trait`, and the
+                        // bounds aren't actually stored nested within each
+                        // other, but separately. So if the `T` refers to a type
+                        // parameter of the outer function, it's just one binder
+                        // away instead of two.
+                        let actual_opaque_type_data = ctx
+                            .with_debruijn(DebruijnIndex::INNERMOST, |ctx| {
+                                ReturnTypeImplTrait::from_hir(ctx, &bounds)
+                            });
+                        ctx.opaque_type_data.borrow_mut()[idx as usize] = actual_opaque_type_data;
+
+                        let func = match ctx.resolver.generic_def() {
+                            Some(GenericDefId::FunctionId(f)) => f,
+                            _ => {
+                                // this shouldn't happen
+                                return (Ty::Unknown, None);
+                            }
+                        };
+                        let impl_trait_id = OpaqueTyId::ReturnTypeImplTrait(func, idx);
+                        let generics = generics(ctx.db.upcast(), func.into());
+                        let parameters = Substs::bound_vars(&generics, ctx.in_binders);
+                        Ty::Opaque(OpaqueTy { opaque_ty_id: impl_trait_id, parameters })
                     }
                     ImplTraitLoweringMode::Param => {
                         let idx = ctx.impl_trait_counter.get();
-                        ctx.impl_trait_counter.set(idx + 1);
+                        // FIXME we're probably doing something wrong here
+                        ctx.impl_trait_counter.set(idx + count_impl_traits(type_ref) as u16);
                         if let Some(def) = ctx.resolver.generic_def() {
                             let generics = generics(ctx.db.upcast(), def);
                             let param = generics
@@ -197,7 +250,8 @@ impl Ty {
                     }
                     ImplTraitLoweringMode::Variable => {
                         let idx = ctx.impl_trait_counter.get();
-                        ctx.impl_trait_counter.set(idx + 1);
+                        // FIXME we're probably doing something wrong here
+                        ctx.impl_trait_counter.set(idx + count_impl_traits(type_ref) as u16);
                         let (parent_params, self_params, list_params, _impl_trait_params) =
                             if let Some(def) = ctx.resolver.generic_def() {
                                 let generics = generics(ctx.db.upcast(), def);
@@ -663,6 +717,29 @@ fn assoc_type_bindings_from_type_bound<'a>(
         })
 }
 
+impl ReturnTypeImplTrait {
+    fn from_hir(ctx: &TyLoweringContext, bounds: &[TypeBound]) -> Self {
+        let self_ty = Ty::Bound(BoundVar::new(DebruijnIndex::INNERMOST, 0));
+        let predicates = ctx.with_shifted_in(DebruijnIndex::ONE, |ctx| {
+            bounds
+                .iter()
+                .flat_map(|b| GenericPredicate::from_type_bound(ctx, b, self_ty.clone()))
+                .collect()
+        });
+        ReturnTypeImplTrait { bounds: Binders::new(1, predicates) }
+    }
+}
+
+fn count_impl_traits(type_ref: &TypeRef) -> usize {
+    let mut count = 0;
+    type_ref.walk(&mut |type_ref| {
+        if matches!(type_ref, TypeRef::ImplTrait(_)) {
+            count += 1;
+        }
+    });
+    count
+}
+
 /// Build the signature of a callable item (function, struct or enum variant).
 pub fn callable_item_sig(db: &dyn HirDatabase, def: CallableDef) -> PolyFnSig {
     match def {
@@ -864,7 +941,9 @@ fn fn_sig_for_fn(db: &dyn HirDatabase, def: FunctionId) -> PolyFnSig {
         .with_impl_trait_mode(ImplTraitLoweringMode::Variable)
         .with_type_param_mode(TypeParamLoweringMode::Variable);
     let params = data.params.iter().map(|tr| Ty::from_hir(&ctx_params, tr)).collect::<Vec<_>>();
-    let ctx_ret = ctx_params.with_impl_trait_mode(ImplTraitLoweringMode::Opaque);
+    let ctx_ret = TyLoweringContext::new(db, &resolver)
+        .with_impl_trait_mode(ImplTraitLoweringMode::Opaque)
+        .with_type_param_mode(TypeParamLoweringMode::Variable);
     let ret = Ty::from_hir(&ctx_ret, &data.ret_type);
     let generics = generics(db.upcast(), def.into());
     let num_binders = generics.len();
@@ -1084,3 +1163,25 @@ pub(crate) fn impl_trait_query(db: &dyn HirDatabase, impl_id: ImplId) -> Option<
         TraitRef::from_hir(&ctx, target_trait, Some(self_ty.value))?,
     ))
 }
+
+pub(crate) fn return_type_impl_traits(
+    db: &impl HirDatabase,
+    def: hir_def::FunctionId,
+) -> Option<Arc<Binders<ReturnTypeImplTraits>>> {
+    // FIXME unify with fn_sig_for_fn instead of doing lowering twice, maybe
+    let data = db.function_data(def);
+    let resolver = def.resolver(db.upcast());
+    let ctx_ret = TyLoweringContext::new(db, &resolver)
+        .with_impl_trait_mode(ImplTraitLoweringMode::Opaque)
+        .with_type_param_mode(TypeParamLoweringMode::Variable);
+    let _ret = Ty::from_hir(&ctx_ret, &data.ret_type);
+    let generics = generics(db.upcast(), def.into());
+    let num_binders = generics.len();
+    let return_type_impl_traits =
+        ReturnTypeImplTraits { impl_traits: ctx_ret.opaque_type_data.into_inner() };
+    if return_type_impl_traits.impl_traits.is_empty() {
+        None
+    } else {
+        Some(Arc::new(Binders::new(num_binders, return_type_impl_traits)))
+    }
+}