More closures work

Author: jackh726

chalk-integration/src/db.rs

@@ -7,9 +7,9 @@ use crate::{
     tls,
 };
 use chalk_ir::{
-    AdtId, AssocTypeId, Canonical, ClosureId, ConstrainedSubst, Environment, FnDefId, GenericArg,
-    Goal, ImplId, InEnvironment, OpaqueTyId, ProgramClause, ProgramClauses, Substitution, TraitId,
-    Ty, UCanonical,
+    AdtId, AssocTypeId, Binders, Canonical, ClosureId, ConstrainedSubst, Environment, FnDefId,
+    GenericArg, Goal, ImplId, InEnvironment, OpaqueTyId, ProgramClause, ProgramClauses,
+    Substitution, TraitId, Ty, UCanonical,
 };
 use chalk_solve::rust_ir::{
     AdtDatum, AssociatedTyDatum, AssociatedTyValue, AssociatedTyValueId, ClosureDatum, FnDefDatum,
@@ -168,7 +168,7 @@ impl RustIrDatabase<ChalkIr> for ChalkDatabase {
         &self,
         closure_id: ClosureId<ChalkIr>,
         substs: &Substitution<ChalkIr>,
-    ) -> Ty<ChalkIr> {
+    ) -> Binders<Ty<ChalkIr>> {
         self.program_ir()
             .unwrap()
             .closure_upvars(closure_id, substs)

chalk-integration/src/lib.rs

@@ -20,6 +20,7 @@ pub use interner::{Identifier, RawId};
 pub enum TypeSort {
     Struct,
     FnDef,
+    Closure,
     Trait,
     Opaque,
 }

chalk-integration/src/lowering.rs

@@ -26,6 +26,7 @@ type OpaqueTyIds = BTreeMap<Ident, chalk_ir::OpaqueTyId<ChalkIr>>;
 type AdtKinds = BTreeMap<chalk_ir::AdtId<ChalkIr>, TypeKind>;
 type FnDefKinds = BTreeMap<chalk_ir::FnDefId<ChalkIr>, TypeKind>;
 type FnDefAbis = BTreeMap<FnDefId<ChalkIr>, <ChalkIr as Interner>::FnAbi>;
+type ClosureKinds = BTreeMap<chalk_ir::ClosureId<ChalkIr>, TypeKind>;
 type TraitKinds = BTreeMap<chalk_ir::TraitId<ChalkIr>, TypeKind>;
 type OpaqueTyKinds = BTreeMap<chalk_ir::OpaqueTyId<ChalkIr>, TypeKind>;
 type AssociatedTyLookups = BTreeMap<(chalk_ir::TraitId<ChalkIr>, Ident), AssociatedTyLookup>;
@@ -44,6 +45,7 @@ struct Env<'k> {
     fn_def_kinds: &'k FnDefKinds,
     fn_def_abis: &'k FnDefAbis,
     closure_ids: &'k ClosureIds,
+    closure_kinds: &'k ClosureKinds,
     trait_ids: &'k TraitIds,
     trait_kinds: &'k TraitKinds,
     opaque_ty_ids: &'k OpaqueTyIds,
@@ -82,6 +84,7 @@ struct AssociatedTyLookup {
 enum ApplyTypeLookup {
     Adt(AdtId<ChalkIr>),
     FnDef(FnDefId<ChalkIr>),
+    Closure(ClosureId<ChalkIr>),
     Opaque(OpaqueTyId<ChalkIr>),
 }
 
@@ -146,18 +149,8 @@ impl<'k> Env<'k> {
         if let Some(id) = self.closure_ids.get(&name.str) {
             return Ok(chalk_ir::TyData::Apply(chalk_ir::ApplicationTy {
                 name: chalk_ir::TypeName::Closure(id.clone()),
-                // XXX closure upvars
-                substitution: chalk_ir::Substitution::from(
-                    interner,
-                    Some(
-                        chalk_ir::TyData::Apply(chalk_ir::ApplicationTy {
-                            name: chalk_ir::TypeName::Tuple(0),
-                            substitution: chalk_ir::Substitution::empty(interner),
-                        })
-                        .intern(interner)
-                        .cast(interner),
-                    ),
-                ),
+                // See note in `program`. Unlike rustc, we store upvars separately.
+                substitution: chalk_ir::Substitution::empty(interner),
             })
             .intern(interner)
             .cast(interner));
@@ -193,6 +186,10 @@ impl<'k> Env<'k> {
             return Ok(ApplyTypeLookup::FnDef(*id));
         }
 
+        if let Some(id) = self.closure_ids.get(&name.str) {
+            return Ok(ApplyTypeLookup::Closure(*id));
+        }
+
         if let Some(id) = self.opaque_ty_ids.get(&name.str) {
             return Ok(ApplyTypeLookup::Opaque(*id));
         }
@@ -228,6 +225,10 @@ impl<'k> Env<'k> {
         &self.fn_def_kinds[&id]
     }
 
+    fn closure_kind(&self, id: chalk_ir::ClosureId<ChalkIr>) -> &TypeKind {
+        &self.closure_kinds[&id]
+    }
+
     fn opaque_kind(&self, id: chalk_ir::OpaqueTyId<ChalkIr>) -> &TypeKind {
         &self.opaque_ty_kinds[&id]
     }
@@ -346,6 +347,7 @@ impl LowerProgram for Program {
         let mut adt_kinds = BTreeMap::new();
         let mut fn_def_kinds = BTreeMap::new();
         let mut fn_def_abis = BTreeMap::new();
+        let mut closure_kinds = BTreeMap::new();
         let mut trait_kinds = BTreeMap::new();
         let mut opaque_ty_kinds = BTreeMap::new();
         let mut object_safe_traits = HashSet::new();
@@ -365,8 +367,10 @@ impl LowerProgram for Program {
                     fn_def_abis.insert(id, defn.abi.lower()?);
                 }
                 Item::ClosureDefn(defn) => {
+                    let type_kind = defn.lower_type_kind()?;
                     let id = ClosureId(raw_id);
                     closure_ids.insert(defn.name.str.clone(), id);
+                    closure_kinds.insert(id, type_kind);
                 }
                 Item::TraitDefn(defn) => {
                     let type_kind = defn.lower_type_kind()?;
@@ -392,6 +396,7 @@ impl LowerProgram for Program {
         let mut adt_data = BTreeMap::new();
         let mut fn_def_data = BTreeMap::new();
         let mut closure_data = BTreeMap::new();
+        let mut closure_upvars = BTreeMap::new();
         let mut trait_data = BTreeMap::new();
         let mut well_known_traits = BTreeMap::new();
         let mut impl_data = BTreeMap::new();
@@ -408,6 +413,7 @@ impl LowerProgram for Program {
                 fn_def_kinds: &fn_def_kinds,
                 fn_def_abis: &fn_def_abis,
                 closure_ids: &closure_ids,
+                closure_kinds: &closure_kinds,
                 trait_ids: &trait_ids,
                 trait_kinds: &trait_kinds,
                 opaque_ty_ids: &opaque_ty_ids,
@@ -434,6 +440,20 @@ impl LowerProgram for Program {
                         closure_def_id,
                         Arc::new(defn.lower_closure(closure_def_id, &empty_env)?),
                     );
+                    let upvars = empty_env.in_binders(defn.all_parameters(), |env| {
+                        let upvar_tys: LowerResult<Vec<chalk_ir::Ty<ChalkIr>>> =
+                            defn.upvars.iter().map(|ty| ty.lower(&env)).collect();
+                        let substitution = chalk_ir::Substitution::from(
+                            &ChalkIr,
+                            upvar_tys?.into_iter().map(|ty| ty.cast(&ChalkIr)),
+                        );
+                        Ok(chalk_ir::TyData::Apply(chalk_ir::ApplicationTy {
+                            name: chalk_ir::TypeName::Tuple(defn.upvars.len()),
+                            substitution,
+                        })
+                        .intern(&ChalkIr))
+                    })?;
+                    closure_upvars.insert(closure_def_id, upvars);
                 }
                 Item::TraitDefn(ref trait_defn) => {
                     let trait_id = TraitId(raw_id);
@@ -593,6 +613,8 @@ impl LowerProgram for Program {
             adt_ids,
             fn_def_ids,
             closure_ids,
+            closure_upvars,
+            closure_kinds,
             trait_ids,
             adt_kinds,
             fn_def_kinds,
@@ -695,6 +717,16 @@ impl LowerParameterMap for FnDefn {
     }
 }
 
+impl LowerParameterMap for ClosureDefn {
+    fn synthetic_parameters(&self) -> Option<chalk_ir::WithKind<ChalkIr, Ident>> {
+        None
+    }
+
+    fn declared_parameters(&self) -> &[VariableKind] {
+        &self.variable_kinds
+    }
+}
+
 impl LowerParameterMap for Impl {
     fn synthetic_parameters(&self) -> Option<chalk_ir::WithKind<ChalkIr, Ident>> {
         None
@@ -832,6 +864,20 @@ impl LowerWhereClauses for FnDefn {
     }
 }
 
+impl LowerTypeKind for ClosureDefn {
+    fn lower_type_kind(&self) -> LowerResult<TypeKind> {
+        let interner = &ChalkIr;
+        Ok(TypeKind {
+            sort: TypeSort::Closure,
+            name: self.name.str.clone(),
+            binders: chalk_ir::Binders::new(
+                chalk_ir::VariableKinds::from(interner, self.all_parameters().anonymize()),
+                crate::Unit,
+            ),
+        })
+    }
+}
+
 impl LowerWhereClauses for StructDefn {
     fn where_clauses(&self) -> &[QuantifiedWhereClause] {
         &self.where_clauses
@@ -1123,7 +1169,7 @@ impl LowerClosureDefn for ClosureDefn {
         closure_id: chalk_ir::ClosureId<ChalkIr>,
         env: &Env,
     ) -> LowerResult<rust_ir::ClosureDatum<ChalkIr>> {
-        let inputs_and_output = env.in_binders(vec![], |env| {
+        let inputs_and_output = env.in_binders(self.all_parameters(), |env| {
             let args: LowerResult<_> = self.argument_types.iter().map(|t| t.lower(env)).collect();
             let return_type = self.return_type.lower(env)?;
             Ok(rust_ir::FnDefInputsAndOutputDatum {
@@ -1453,6 +1499,9 @@ impl LowerTy for Ty {
                     ApplyTypeLookup::FnDef(id) => {
                         (chalk_ir::TypeName::FnDef(id), env.fn_def_kind(id))
                     }
+                    ApplyTypeLookup::Closure(id) => {
+                        (chalk_ir::TypeName::Closure(id), env.closure_kind(id))
+                    }
                     ApplyTypeLookup::Opaque(id) => {
                         (chalk_ir::TypeName::OpaqueType(id), env.opaque_kind(id))
                     }
@@ -1852,6 +1901,7 @@ impl LowerGoal<LoweredProgram> for Goal {
             adt_kinds: &program.adt_kinds,
             fn_def_kinds: &program.fn_def_kinds,
             fn_def_abis: &fn_def_abis,
+            closure_kinds: &program.closure_kinds,
             trait_kinds: &program.trait_kinds,
             opaque_ty_kinds: &program.opaque_ty_kinds,
             associated_ty_lookups: &associated_ty_lookups,

chalk-integration/src/program.rs

@@ -3,8 +3,8 @@ use crate::{tls, Identifier, TypeKind};
 use chalk_ir::could_match::CouldMatch;
 use chalk_ir::debug::Angle;
 use chalk_ir::{
-    debug::SeparatorTraitRef, AdtId, AliasTy, ApplicationTy, AssocTypeId, ClosureId, FnDefId,
-    GenericArg, Goal, Goals, ImplId, Lifetime, OpaqueTy, OpaqueTyId, ProgramClause,
+    debug::SeparatorTraitRef, AdtId, AliasTy, ApplicationTy, AssocTypeId, Binders, ClosureId,
+    FnDefId, GenericArg, Goal, Goals, ImplId, Lifetime, OpaqueTy, OpaqueTyId, ProgramClause,
     ProgramClauseImplication, ProgramClauses, ProjectionTy, Substitution, TraitId, Ty,
 };
 use chalk_solve::rust_ir::{
@@ -31,6 +31,10 @@ pub struct Program {
 
     pub closure_ids: BTreeMap<Identifier, ClosureId<ChalkIr>>,
 
+    pub closure_upvars: BTreeMap<ClosureId<ChalkIr>, Binders<Ty<ChalkIr>>>,
+
+    pub closure_kinds: BTreeMap<ClosureId<ChalkIr>, TypeKind>,
+
     /// From trait name to item-id. Used during lowering only.
     pub trait_ids: BTreeMap<Identifier, TraitId<ChalkIr>>,
 
@@ -420,21 +424,20 @@ impl RustIrDatabase<ChalkIr> for Program {
     fn closure_datum(
         &self,
         closure_id: ClosureId<ChalkIr>,
-        substs: &Substitution<ChalkIr>,
+        _substs: &Substitution<ChalkIr>,
     ) -> Arc<ClosureDatum<ChalkIr>> {
         self.closure_data[&closure_id].clone()
     }
 
     fn closure_upvars(
         &self,
         closure_id: ClosureId<ChalkIr>,
-        substs: &Substitution<ChalkIr>,
-    ) -> Ty<ChalkIr> {
-        substs
-            .parameters(self.interner())
-            .last()
-            .unwrap()
-            .assert_ty_ref(self.interner())
-            .clone()
+        _substs: &Substitution<ChalkIr>,
+    ) -> Binders<Ty<ChalkIr>> {
+        // This is different to how rustc does it. In rustc,
+        // upvars are stored in the last substitution parameter.
+        // It's just easier this way, and from `chalk-solve`s POV,
+        // this is opaque anyways.
+        self.closure_upvars[&closure_id].clone()
     }
 }

chalk-parse/src/ast.rs

@@ -59,8 +59,10 @@ pub struct FnDefn {
 pub struct ClosureDefn {
     pub name: Identifier,
     pub kind: ClosureKind,
+    pub variable_kinds: Vec<VariableKind>,
     pub argument_types: Vec<Ty>,
     pub return_type: Ty,
+    pub upvars: Vec<Ty>,
 }
 
 #[derive(Clone, Eq, PartialEq, Debug)]

chalk-parse/src/parser.lalrpop

@@ -104,18 +104,25 @@ FnArgs: Vec<Ty> = {
 };
 
 ClosureDefn: ClosureDefn = {
-    "closure" <n:Id> "(" <s:ClosureSelf> ")" ";" => ClosureDefn {
+    "closure" <n:Id> <p:Angle<VariableKind>> "(" <s:ClosureSelf> <args:ClosureArgs> ")" <ret_ty:FnReturn?>
+        "{" <upvars:SemiColon<Ty>> "}" => ClosureDefn {
         name: n,
         kind: s,
-        argument_types: vec![],
-        return_type: Ty::Tuple { types: Vec::new() },
+        variable_kinds: p,
+        argument_types: args,
+        return_type: ret_ty.unwrap_or_else(|| Ty::Tuple { types: Vec::new() }),
+        upvars: upvars,
     }
 }
 
 ClosureSelf: ClosureKind = {
     "self" => ClosureKind::FnOnce,
-    "mutself" => ClosureKind::FnMut,
-    "refself" => ClosureKind::Fn,
+    "&" "mut" "self" => ClosureKind::FnMut,
+    "&" "self" => ClosureKind::Fn,
+}
+
+ClosureArgs: Vec<Ty> = {
+    "," <args:FnArgs> => args,
 }
 
 TraitDefn: TraitDefn = {

chalk-solve/src/clauses/builtin_traits/copy.rs

@@ -56,7 +56,9 @@ pub fn add_copy_program_clauses<I: Interner>(
             }
             TypeName::Closure(closure_id) => {
                 let upvars = db.closure_upvars(*closure_id, substitution);
-                needs_impl_for_tys(db, builder, trait_ref, Some(upvars).into_iter());
+                builder.push_binders(&upvars, |builder, upvars| {
+                    needs_impl_for_tys(db, builder, trait_ref, Some(upvars).into_iter());
+                })
             }
             _ => return,
         },