Fix retrieval of AsyncFnOnce::Output

We shouldn't just assume it's always the second associated type of `AsyncFnOnce`.

Author: ChayimFriedman2

chalk-integration/src/db.rs

@@ -15,7 +15,7 @@ use chalk_ir::{
 use chalk_solve::rust_ir::{
     AdtDatum, AdtRepr, AdtSizeAlign, AssociatedTyDatum, AssociatedTyValue, AssociatedTyValueId,
     ClosureKind, CoroutineDatum, CoroutineWitnessDatum, FnDefDatum, FnDefInputsAndOutputDatum,
-    ImplDatum, OpaqueTyDatum, TraitDatum, WellKnownTrait,
+    ImplDatum, OpaqueTyDatum, TraitDatum, WellKnownAssocType, WellKnownTrait,
 };
 use chalk_solve::{RustIrDatabase, Solution, SubstitutionResult};
 use salsa::Database;
@@ -183,6 +183,15 @@ impl RustIrDatabase<ChalkIr> for ChalkDatabase {
             .well_known_trait_id(well_known_trait)
     }
 
+    fn well_known_assoc_type_id(
+        &self,
+        assoc_type: WellKnownAssocType,
+    ) -> Option<AssocTypeId<ChalkIr>> {
+        self.program_ir()
+            .unwrap()
+            .well_known_assoc_type_id(assoc_type)
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &Environment<ChalkIr>,

chalk-integration/src/program.rs

@@ -537,6 +537,14 @@ impl RustIrDatabase<ChalkIr> for Program {
         self.well_known_traits.get(&well_known_trait).copied()
     }
 
+    fn well_known_assoc_type_id(
+        &self,
+        _assoc_type: chalk_solve::rust_ir::WellKnownAssocType,
+    ) -> Option<AssocTypeId<ChalkIr>> {
+        // FIXME
+        None
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &chalk_ir::Environment<ChalkIr>,

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

@@ -1,5 +1,5 @@
 use crate::clauses::ClauseBuilder;
-use crate::rust_ir::{ClosureKind, FnDefInputsAndOutputDatum, WellKnownTrait};
+use crate::rust_ir::{ClosureKind, FnDefInputsAndOutputDatum, WellKnownAssocType, WellKnownTrait};
 use crate::{Interner, RustIrDatabase, TraitRef};
 use chalk_ir::cast::Cast;
 use chalk_ir::{
@@ -83,14 +83,9 @@ fn push_clauses<I: Interner>(
 
         if let WellKnownTrait::AsyncFnOnce = well_known {
             builder.push_bound_ty(|builder, ty| {
-                let trait_datum = db.trait_datum(trait_id);
-                assert_eq!(
-                    trait_datum.associated_ty_ids.len(),
-                    2,
-                    "AsyncFnOnce trait should have exactly two associated types, found {:?}",
-                    trait_datum.associated_ty_ids
-                );
-                let output_id = trait_datum.associated_ty_ids[1];
+                let output_id = db
+                    .well_known_assoc_type_id(WellKnownAssocType::AsyncFnOnceOutput)
+                    .unwrap();
                 let async_alias = AliasTy::Projection(ProjectionTy {
                     associated_ty_id: output_id,
                     substitution,

chalk-solve/src/display/stub.rs

@@ -188,6 +188,13 @@ impl<I: Interner, DB: RustIrDatabase<I>> RustIrDatabase<I> for StubWrapper<'_, D
         self.db.well_known_trait_id(well_known_trait)
     }
 
+    fn well_known_assoc_type_id(
+        &self,
+        assoc_type: crate::rust_ir::WellKnownAssocType,
+    ) -> Option<chalk_ir::AssocTypeId<I>> {
+        self.db.well_known_assoc_type_id(assoc_type)
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &chalk_ir::Environment<I>,

chalk-solve/src/lib.rs

@@ -128,6 +128,9 @@ pub trait RustIrDatabase<I: Interner>: Debug {
     /// Returns id of a trait lang item, if found
     fn well_known_trait_id(&self, well_known_trait: WellKnownTrait) -> Option<TraitId<I>>;
 
+    /// Returns id of a associated type, if found.
+    fn well_known_assoc_type_id(&self, assoc_type: WellKnownAssocType) -> Option<AssocTypeId<I>>;
+
     /// Calculates program clauses from an env. This is intended to call the
     /// `program_clauses_for_env` function and then possibly cache the clauses.
     fn program_clauses_for_env(&self, environment: &Environment<I>) -> ProgramClauses<I>;

chalk-solve/src/logging_db.rs

@@ -229,6 +229,14 @@ where
         trait_id
     }
 
+    fn well_known_assoc_type_id(&self, assoc_type: WellKnownAssocType) -> Option<AssocTypeId<I>> {
+        let assoc_type_id = self.ws.db().well_known_assoc_type_id(assoc_type);
+        if let Some(id) = assoc_type_id {
+            self.record(self.ws.db().associated_ty_data(id).trait_id);
+        }
+        assoc_type_id
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &chalk_ir::Environment<I>,
@@ -489,6 +497,10 @@ where
         self.db.well_known_trait_id(well_known_trait)
     }
 
+    fn well_known_assoc_type_id(&self, assoc_type: WellKnownAssocType) -> Option<AssocTypeId<I>> {
+        self.db.well_known_assoc_type_id(assoc_type)
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &chalk_ir::Environment<I>,

chalk-solve/src/rust_ir.rs

@@ -287,6 +287,15 @@ pub enum WellKnownTrait {
 
 chalk_ir::const_visit!(WellKnownTrait);
 
+/// A list of the associated types that are "well known" to chalk, which means that
+/// the chalk-solve crate has special, hard-coded impls for them.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Ord, PartialOrd, Hash)]
+pub enum WellKnownAssocType {
+    AsyncFnOnceOutput,
+}
+
+chalk_ir::const_visit!(WellKnownAssocType);
+
 impl<I: Interner> TraitDatum<I> {
     pub fn is_auto_trait(&self) -> bool {
         self.flags.auto

tests/display/unique_names.rs

@@ -155,6 +155,12 @@ where
     ) -> Option<chalk_ir::TraitId<I>> {
         self.db.well_known_trait_id(well_known_trait)
     }
+    fn well_known_assoc_type_id(
+        &self,
+        assoc_type: chalk_solve::rust_ir::WellKnownAssocType,
+    ) -> Option<chalk_ir::AssocTypeId<I>> {
+        self.db.well_known_assoc_type_id(assoc_type)
+    }
     fn program_clauses_for_env(
         &self,
         environment: &chalk_ir::Environment<I>,

tests/integration/panic.rs

@@ -211,6 +211,13 @@ impl RustIrDatabase<ChalkIr> for MockDatabase {
         unimplemented!()
     }
 
+    fn well_known_assoc_type_id(
+        &self,
+        assoc_type: WellKnownAssocType,
+    ) -> Option<AssocTypeId<ChalkIr>> {
+        unimplemented!()
+    }
+
     fn program_clauses_for_env(
         &self,
         environment: &Environment<ChalkIr>,