vendor/chalk-solve/src/clauses/builtin_traits/fn_family.rs

   1 use crate::clauses::ClauseBuilder;
   2 use crate::infer::instantiate::IntoBindersAndValue;
   3 use crate::rust_ir::{ClosureKind, FnDefInputsAndOutputDatum, WellKnownTrait};
   4 use crate::{Interner, RustIrDatabase, TraitRef};
   5 use chalk_ir::cast::Cast;
   6 use chalk_ir::{
   7     AliasTy, ApplicationTy, Binders, Floundered, Normalize, ProjectionTy, Substitution, TraitId,
   8     Ty, TyData, TypeName, VariableKinds,
   9 };
  10
  11 fn push_clauses<I: Interner>(
  12     db: &dyn RustIrDatabase<I>,
  13     builder: &mut ClauseBuilder<'_, I>,
  14     well_known: WellKnownTrait,
  15     trait_id: TraitId<I>,
  16     self_ty: Ty<I>,
  17     arg_sub: Substitution<I>,
  18     return_type: Ty<I>,
  19 ) {
  20     let interner = db.interner();
  21     let tupled = TyData::Apply(ApplicationTy {
  22         name: TypeName::Tuple(arg_sub.len(interner)),
  23         substitution: arg_sub,
  24     })
  25     .intern(interner);
  26     let substitution =
  27         Substitution::from(interner, &[self_ty.cast(interner), tupled.cast(interner)]);
  28     builder.push_fact(TraitRef {
  29         trait_id,
  30         substitution: substitution.clone(),
  31     });
  32
  33     // The `Output` type is defined on the `FnOnce`
  34     if let WellKnownTrait::FnOnce = well_known {
  35         let trait_datum = db.trait_datum(trait_id);
  36         assert_eq!(
  37             trait_datum.associated_ty_ids.len(),
  38             1,
  39             "FnOnce trait should have exactly one associated type, found {:?}",
  40             trait_datum.associated_ty_ids
  41         );
  42         // Constructs the alias. For `Fn`, for example, this would look like
  43         // `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
  44         let output_id = trait_datum.associated_ty_ids[0];
  45         let alias = AliasTy::Projection(ProjectionTy {
  46             associated_ty_id: output_id,
  47             substitution,
  48         });
  49         builder.push_fact(Normalize {
  50             alias,
  51             ty: return_type,
  52         });
  53     }
  54 }
  55
  56 fn push_clauses_for_apply<I: Interner>(
  57     db: &dyn RustIrDatabase<I>,
  58     builder: &mut ClauseBuilder<'_, I>,
  59     well_known: WellKnownTrait,
  60     trait_id: TraitId<I>,
  61     self_ty: Ty<I>,
  62     inputs_and_output: &Binders<FnDefInputsAndOutputDatum<I>>,
  63 ) {
  64     let interner = db.interner();
  65     builder.push_binders(inputs_and_output, |builder, inputs_and_output| {
  66         let arg_sub = inputs_and_output
  67             .argument_types
  68             .iter()
  69             .cloned()
  70             .map(|ty| ty.cast(interner));
  71         let arg_sub = Substitution::from(interner, arg_sub);
  72         let output_ty = inputs_and_output.return_type;
  73
  74         push_clauses(
  75             db, builder, well_known, trait_id, self_ty, arg_sub, output_ty,
  76         );
  77     });
  78 }
  79
  80 /// Handles clauses for FnOnce/FnMut/Fn.
  81 /// If `self_ty` is a function, we push a clause of the form
  82 /// `fn(A1, A2, ..., AN) -> O: FnTrait<(A1, A2, ..., AN)>`, where `FnTrait`
  83 /// is the trait corresponding to `trait_id` (FnOnce/FnMut/Fn)
  84 ///
  85 /// If `trait_id` is `FnOnce`, we also push a clause for the output type of the form:
  86 /// `Normalize(<fn(A) -> B as FnOnce<(A,)>>::Output -> B)`
  87 /// We do not add the usual `Implemented(fn(A) -> b as FnOnce<(A,)>` clause
  88 /// as a condition, since we already called `push_fact` with it
  89 pub fn add_fn_trait_program_clauses<I: Interner>(
  90     db: &dyn RustIrDatabase<I>,
  91     builder: &mut ClauseBuilder<'_, I>,
  92     well_known: WellKnownTrait,
  93     self_ty: Ty<I>,
  94 ) -> Result<(), Floundered> {
  95     let interner = db.interner();
  96     let trait_id = db.well_known_trait_id(well_known).unwrap();
  97
  98     match self_ty.data(interner) {
  99         TyData::Apply(apply) => match apply.name {
 100             TypeName::FnDef(fn_def_id) => {
 101                 let fn_def_datum = builder.db.fn_def_datum(fn_def_id);
 102                 let bound = fn_def_datum
 103                     .binders
 104                     .substitute(builder.interner(), &apply.substitution);
 105                 push_clauses_for_apply(
 106                     db,
 107                     builder,
 108                     well_known,
 109                     trait_id,
 110                     self_ty,
 111                     &bound.inputs_and_output,
 112                 );
 113                 Ok(())
 114             }
 115             TypeName::Closure(closure_id) => {
 116                 let closure_kind = db.closure_kind(closure_id, &apply.substitution);
 117                 let trait_matches = match (well_known, closure_kind) {
 118                     (WellKnownTrait::Fn, ClosureKind::Fn) => true,
 119                     (WellKnownTrait::FnMut, ClosureKind::FnMut)
 120                     | (WellKnownTrait::FnMut, ClosureKind::Fn) => true,
 121                     (WellKnownTrait::FnOnce, _) => true,
 122                     _ => false,
 123                 };
 124                 if !trait_matches {
 125                     return Ok(());
 126                 }
 127                 let closure_inputs_and_output =
 128                     db.closure_inputs_and_output(closure_id, &apply.substitution);
 129                 push_clauses_for_apply(
 130                     db,
 131                     builder,
 132                     well_known,
 133                     trait_id,
 134                     self_ty,
 135                     &closure_inputs_and_output,
 136                 );
 137                 Ok(())
 138             }
 139             _ => Ok(()),
 140         },
 141         TyData::Function(fn_val) => {
 142             let (binders, orig_sub) = fn_val.into_binders_and_value(interner);
 143             let bound_ref = Binders::new(VariableKinds::from(interner, binders), orig_sub);
 144             builder.push_binders(&bound_ref, |builder, orig_sub| {
 145                 // The last parameter represents the function return type
 146                 let (arg_sub, fn_output_ty) = orig_sub
 147                     .parameters(interner)
 148                     .split_at(orig_sub.len(interner) - 1);
 149                 let arg_sub = Substitution::from(interner, arg_sub);
 150                 let output_ty = fn_output_ty[0].assert_ty_ref(interner).clone();
 151
 152                 push_clauses(
 153                     db,
 154                     builder,
 155                     well_known,
 156                     trait_id,
 157                     self_ty.clone(),
 158                     arg_sub,
 159                     output_ty,
 160                 );
 161             });
 162             Ok(())
 163         }
 164         // Function traits are non-enumerable
 165         TyData::InferenceVar(..) | TyData::Alias(..) => Err(Floundered),
 166         _ => Ok(()),
 167     }
 168 }