Make closures impl closure traits

3 files changed, 185 insertions, 38 deletions

diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index eb8770ec4..13090f89d 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -3990,6 +3990,7 @@ fn test<F: FnOnce(u32, u64) -> u128>(f: F) {
 fn closure_1() {
     assert_snapshot!(
         infer(r#"
+#[lang = "fn_once"]
 trait FnOnce<Args> {
     type Output;
 }
@@ -4000,39 +4001,39 @@ impl<T> Option<T> {
 }
 
 fn test() {
-    let x = Option::Some(1i32);
+    let x = Option::Some(1u32);
     x.map(|v| v + 1);
     x.map(|_v| 1u64);
     let y: Option<i64> = x.map(|_v| 1);
 }
 "#),
         @r###"
-    [128; 132) 'self': Option<T>
-    [134; 135) 'f': F
-    [145; 147) '{}': ()
-    [161; 280) '{     ... 1); }': ()
-    [171; 172) 'x': Option<i32>
-    [175; 187) 'Option::Some': Some<i32>(T) -> Option<T>
-    [175; 193) 'Option...(1i32)': Option<i32>
-    [188; 192) '1i32': i32
-    [199; 200) 'x': Option<i32>
-    [199; 215) 'x.map(...v + 1)': {unknown}
-    [205; 214) '|v| v + 1': |{unknown}| -> i32
-    [206; 207) 'v': {unknown}
-    [209; 210) 'v': {unknown}
-    [209; 214) 'v + 1': i32
-    [213; 214) '1': i32
-    [221; 222) 'x': Option<i32>
-    [221; 237) 'x.map(... 1u64)': {unknown}
-    [227; 236) '|_v| 1u64': |{unknown}| -> u64
-    [228; 230) '_v': {unknown}
-    [232; 236) '1u64': u64
-    [247; 248) 'y': Option<i64>
-    [264; 265) 'x': Option<i32>
-    [264; 277) 'x.map(|_v| 1)': Option<i64>
-    [270; 276) '|_v| 1': |{unknown}| -> i32
-    [271; 273) '_v': {unknown}
-    [275; 276) '1': i32
+    [148; 152) 'self': Option<T>
+    [154; 155) 'f': F
+    [165; 167) '{}': ()
+    [181; 300) '{     ... 1); }': ()
+    [191; 192) 'x': Option<u32>
+    [195; 207) 'Option::Some': Some<u32>(T) -> Option<T>
+    [195; 213) 'Option...(1u32)': Option<u32>
+    [208; 212) '1u32': u32
+    [219; 220) 'x': Option<u32>
+    [219; 235) 'x.map(...v + 1)': {unknown}
+    [225; 234) '|v| v + 1': |u32| -> i32
+    [226; 227) 'v': u32
+    [229; 230) 'v': u32
+    [229; 234) 'v + 1': i32
+    [233; 234) '1': i32
+    [241; 242) 'x': Option<u32>
+    [241; 257) 'x.map(... 1u64)': {unknown}
+    [247; 256) '|_v| 1u64': |u32| -> u64
+    [248; 250) '_v': u32
+    [252; 256) '1u64': u64
+    [267; 268) 'y': Option<i64>
+    [284; 285) 'x': Option<u32>
+    [284; 297) 'x.map(|_v| 1)': Option<i64>
+    [290; 296) '|_v| 1': |u32| -> i32
+    [291; 293) '_v': u32
+    [295; 296) '1': i32
     "###
     );
 }
diff --git a/crates/ra_hir/src/ty/traits.rs b/crates/ra_hir/src/ty/traits.rs
index c0c132809..8e256341d 100644
--- a/crates/ra_hir/src/ty/traits.rs
+++ b/crates/ra_hir/src/ty/traits.rs
@@ -9,7 +9,7 @@ use ra_prof::profile;
 use rustc_hash::FxHashSet;
 
 use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};
-use crate::{db::HirDatabase, Crate, ImplBlock, Trait};
+use crate::{db::HirDatabase, expr::ExprId, Crate, DefWithBody, ImplBlock, Trait};
 
 use self::chalk::{from_chalk, ToChalk};
 
@@ -252,3 +252,37 @@ pub enum Guidance {
     /// There's no useful information to feed back to type inference
     Unknown,
 }
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum FnTrait {
+    FnOnce,
+    FnMut,
+    Fn,
+}
+
+impl FnTrait {
+    fn lang_item_name(self) -> &'static str {
+        match self {
+            FnTrait::FnOnce => "fn_once",
+            FnTrait::FnMut => "fn_mut",
+            FnTrait::Fn => "fn",
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct ClosureFnTraitImplData {
+    def: DefWithBody,
+    expr: ExprId,
+    fn_trait: FnTrait,
+}
+
+/// An impl. Usually this comes from an impl block, but some built-in types get
+/// synthetic impls.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum Impl {
+    /// A normal impl from an impl block.
+    ImplBlock(ImplBlock),
+    /// Closure types implement the Fn traits synthetically.
+    ClosureFnTraitImpl(ClosureFnTraitImplData),
+}
diff --git a/crates/ra_hir/src/ty/traits/chalk.rs b/crates/ra_hir/src/ty/traits/chalk.rs
index 34e623931..f229b1aef 100644
--- a/crates/ra_hir/src/ty/traits/chalk.rs
+++ b/crates/ra_hir/src/ty/traits/chalk.rs
@@ -12,7 +12,7 @@ use chalk_rust_ir::{AssociatedTyDatum, ImplDatum, StructDatum, TraitDatum};
 use ra_db::salsa::{InternId, InternKey};
 use test_utils::tested_by;
 
-use super::{Canonical, ChalkContext, Obligation};
+use super::{Canonical, ChalkContext, Impl, Obligation};
 use crate::{
     db::HirDatabase,
     generics::GenericDef,
@@ -111,7 +111,7 @@ impl ToChalk for Ty {
             }
             chalk_ir::Ty::ForAll(_) => unimplemented!(),
             chalk_ir::Ty::BoundVar(idx) => Ty::Bound(idx as u32),
-            chalk_ir::Ty::InferenceVar(_iv) => panic!("unexpected chalk infer ty"),
+            chalk_ir::Ty::InferenceVar(_iv) => Ty::Unknown,
         }
     }
 }
@@ -175,15 +175,15 @@ impl ToChalk for TypeCtor {
     }
 }
 
-impl ToChalk for ImplBlock {
+impl ToChalk for Impl {
     type Chalk = chalk_ir::ImplId;
 
     fn to_chalk(self, db: &impl HirDatabase) -> chalk_ir::ImplId {
-        db.intern_impl_block(self).into()
+        db.intern_impl(self).into()
     }
 
-    fn from_chalk(db: &impl HirDatabase, impl_id: chalk_ir::ImplId) -> ImplBlock {
-        db.lookup_intern_impl_block(impl_id.into())
+    fn from_chalk(db: &impl HirDatabase, impl_id: chalk_ir::ImplId) -> Impl {
+        db.lookup_intern_impl(impl_id.into())
     }
 }
 
@@ -388,19 +388,36 @@ where
     fn impls_for_trait(
         &self,
         trait_id: chalk_ir::TraitId,
-        _parameters: &[Parameter],
+        parameters: &[Parameter],
     ) -> Vec<ImplId> {
         debug!("impls_for_trait {:?}", trait_id);
         if trait_id == UNKNOWN_TRAIT {
             return Vec::new();
         }
         let trait_: Trait = from_chalk(self.db, trait_id);
-        let result: Vec<_> = self
+        let mut result: Vec<_> = self
             .db
             .impls_for_trait(self.krate, trait_)
             .iter()
-            .map(|impl_block| impl_block.to_chalk(self.db))
+            .copied()
+            .map(Impl::ImplBlock)
+            .map(|impl_| impl_.to_chalk(self.db))
             .collect();
+
+        let ty: Ty = from_chalk(self.db, parameters[0].assert_ty_ref().clone());
+        if let Ty::Apply(ApplicationTy { ctor: TypeCtor::Closure { def, expr }, .. }) = ty {
+            for fn_trait in
+                [super::FnTrait::FnOnce, super::FnTrait::FnMut, super::FnTrait::Fn].iter().copied()
+            {
+                if let Some(actual_trait) = get_fn_trait(self.db, self.krate, fn_trait) {
+                    if trait_ == actual_trait {
+                        let impl_ = super::ClosureFnTraitImplData { def, expr, fn_trait };
+                        result.push(Impl::ClosureFnTraitImpl(impl_).to_chalk(self.db));
+                    }
+                }
+            }
+        }
+
         debug!("impls_for_trait returned {} impls", result.len());
         result
     }
@@ -602,7 +619,21 @@ pub(crate) fn impl_datum_query(
 ) -> Arc<ImplDatum> {
     let _p = ra_prof::profile("impl_datum");
     debug!("impl_datum {:?}", impl_id);
-    let impl_block: ImplBlock = from_chalk(db, impl_id);
+    let impl_: Impl = from_chalk(db, impl_id);
+    match impl_ {
+        Impl::ImplBlock(impl_block) => impl_block_datum(db, krate, impl_id, impl_block),
+        Impl::ClosureFnTraitImpl(data) => {
+            closure_fn_trait_impl_datum(db, krate, impl_id, data).unwrap_or_else(invalid_impl_datum)
+        }
+    }
+}
+
+fn impl_block_datum(
+    db: &impl HirDatabase,
+    krate: Crate,
+    impl_id: ImplId,
+    impl_block: ImplBlock,
+) -> Arc<ImplDatum> {
     let generic_params = impl_block.generic_params(db);
     let bound_vars = Substs::bound_vars(&generic_params);
     let trait_ref = impl_block
@@ -661,6 +692,87 @@ pub(crate) fn impl_datum_query(
     Arc::new(impl_datum)
 }
 
+fn invalid_impl_datum() -> Arc<ImplDatum> {
+    let trait_ref = chalk_ir::TraitRef {
+        trait_id: UNKNOWN_TRAIT,
+        parameters: vec![chalk_ir::Ty::BoundVar(0).cast()],
+    };
+    let impl_datum_bound = chalk_rust_ir::ImplDatumBound {
+        trait_ref: chalk_rust_ir::PolarizedTraitRef::Positive(trait_ref),
+        where_clauses: Vec::new(),
+        associated_ty_values: Vec::new(),
+        impl_type: chalk_rust_ir::ImplType::External,
+    };
+    let impl_datum = ImplDatum { binders: make_binders(impl_datum_bound, 1) };
+    Arc::new(impl_datum)
+}
+
+fn closure_fn_trait_impl_datum(
+    db: &impl HirDatabase,
+    krate: Crate,
+    impl_id: ImplId,
+    data: super::ClosureFnTraitImplData,
+) -> Option<Arc<ImplDatum>> {
+    // for some closure |X, Y| -> Z:
+    // impl<T, U, V> Fn<(T, U)> for closure<fn(T, U) -> V> { Output = V }
+
+    let fn_once_trait = get_fn_trait(db, krate, super::FnTrait::FnOnce)?;
+    let trait_ = get_fn_trait(db, krate, data.fn_trait)?; // get corresponding fn trait
+
+    let num_args: u16 = match &db.body_hir(data.def)[data.expr] {
+        crate::expr::Expr::Lambda { args, .. } => args.len() as u16,
+        _ => {
+            log::warn!("closure for closure type {:?} not found", data);
+            0
+        }
+    };
+
+    let arg_ty = Ty::apply(
+        TypeCtor::Tuple { cardinality: num_args },
+        (0..num_args).map(|i| Ty::Bound(i.into())).collect::<Vec<_>>().into(),
+    );
+    let output_ty = Ty::Bound(num_args.into());
+    let sig_ty = Ty::apply(
+        TypeCtor::FnPtr { num_args },
+        (0..num_args + 1).map(|i| Ty::Bound(i.into())).collect::<Vec<_>>().into(),
+    );
+
+    let self_ty = Ty::apply_one(TypeCtor::Closure { def: data.def, expr: data.expr }, sig_ty);
+
+    let trait_ref = TraitRef { trait_, substs: vec![self_ty, arg_ty].into() };
+
+    let output_ty_id = fn_once_trait.associated_type_by_name(db, &crate::name::OUTPUT_TYPE)?;
+
+    let output_ty_value = chalk_rust_ir::AssociatedTyValue {
+        associated_ty_id: output_ty_id.to_chalk(db),
+        impl_id,
+        value: make_binders(
+            chalk_rust_ir::AssociatedTyValueBound { ty: output_ty.to_chalk(db) },
+            0,
+        ),
+    };
+
+    let impl_type = chalk_rust_ir::ImplType::External;
+
+    let impl_datum_bound = chalk_rust_ir::ImplDatumBound {
+        trait_ref: chalk_rust_ir::PolarizedTraitRef::Positive(trait_ref.to_chalk(db)),
+        where_clauses: Vec::new(),
+        associated_ty_values: vec![output_ty_value],
+        impl_type,
+    };
+    let impl_datum = ImplDatum { binders: make_binders(impl_datum_bound, num_args as usize + 1) };
+    Some(Arc::new(impl_datum))
+}
+
+fn get_fn_trait(db: &impl HirDatabase, krate: Crate, fn_trait: super::FnTrait) -> Option<Trait> {
+    let lang_items = db.lang_items(krate);
+    let target = lang_items.target(fn_trait.lang_item_name())?;
+    match target {
+        crate::lang_item::LangItemTarget::Trait(t) => Some(*t),
+        _ => None,
+    }
+}
+
 fn id_from_chalk<T: InternKey>(chalk_id: chalk_ir::RawId) -> T {
     T::from_intern_id(InternId::from(chalk_id.index))
 }