Use indexmap throughout

use `DiGraph`specialization forest

Author: pierwill

Cargo.lock

@@ -20,3 +20,4 @@ chalk-solve = { version = "0.76.0-dev.0", path = "../chalk-solve" }
 chalk-recursive = { version = "0.76.0-dev.0", path = "../chalk-recursive" }
 chalk-engine = { version = "0.76.0-dev.0", path = "../chalk-engine" }
 chalk-parse = { version = "0.76.0-dev.0", path = "../chalk-parse" }
+indexmap = "1.7.0"

chalk-integration/Cargo.toml

@@ -9,7 +9,6 @@ use chalk_ir::{
 use chalk_parse::ast::*;
 use chalk_solve::rust_ir::{self, IntoWhereClauses};
 use program_lowerer::ProgramLowerer;
-use std::collections::BTreeMap;
 use string_cache::DefaultAtom as Atom;
 use tracing::debug;
 
@@ -564,8 +563,6 @@ impl LowerWithEnv for [QuantifiedInlineBound] {
             }
         }
 
-        auto_traits.sort_by_key(|b| b.1);
-
         regular_traits
             .iter()
             .chain(auto_traits.iter())
@@ -1014,7 +1011,7 @@ impl LowerWithEnv for (&TraitDefn, chalk_ir::TraitId<ChalkIr>) {
 
 pub fn lower_goal(goal: &Goal, program: &LoweredProgram) -> LowerResult<chalk_ir::Goal<ChalkIr>> {
     let interner = ChalkIr;
-    let associated_ty_lookups: BTreeMap<_, _> = program
+    let associated_ty_lookups: indexmap::IndexMap<_, _> = program
         .associated_ty_data
         .iter()
         .map(|(&associated_ty_id, datum)| {
@@ -1052,7 +1049,7 @@ pub fn lower_goal(goal: &Goal, program: &LoweredProgram) -> LowerResult<chalk_ir
         opaque_ty_kinds: &program.opaque_ty_kinds,
         associated_ty_lookups: &associated_ty_lookups,
         foreign_ty_ids: &program.foreign_ty_ids,
-        parameter_map: BTreeMap::new(),
+        parameter_map: indexmap::IndexMap::new(),
         auto_traits: &auto_traits,
     };
 

chalk-integration/src/lowering.rs

@@ -6,31 +6,32 @@ use chalk_ir::{
 use chalk_ir::{cast::Cast, ForeignDefId, WithKind};
 use chalk_parse::ast::*;
 use chalk_solve::rust_ir::AssociatedTyValueId;
-use std::collections::BTreeMap;
+// use std::collections::BTreeMap;
 
 use crate::error::RustIrError;
 use crate::interner::ChalkIr;
 use crate::{Identifier as Ident, TypeKind};
 
-pub type AdtIds = BTreeMap<Ident, chalk_ir::AdtId<ChalkIr>>;
-pub type FnDefIds = BTreeMap<Ident, chalk_ir::FnDefId<ChalkIr>>;
-pub type ClosureIds = BTreeMap<Ident, chalk_ir::ClosureId<ChalkIr>>;
-pub type TraitIds = BTreeMap<Ident, chalk_ir::TraitId<ChalkIr>>;
-pub type GeneratorIds = BTreeMap<Ident, chalk_ir::GeneratorId<ChalkIr>>;
-pub type OpaqueTyIds = BTreeMap<Ident, chalk_ir::OpaqueTyId<ChalkIr>>;
-pub type AdtKinds = BTreeMap<chalk_ir::AdtId<ChalkIr>, TypeKind>;
-pub type FnDefKinds = BTreeMap<chalk_ir::FnDefId<ChalkIr>, TypeKind>;
-pub type ClosureKinds = BTreeMap<chalk_ir::ClosureId<ChalkIr>, TypeKind>;
-pub type TraitKinds = BTreeMap<chalk_ir::TraitId<ChalkIr>, TypeKind>;
-pub type AutoTraits = BTreeMap<chalk_ir::TraitId<ChalkIr>, bool>;
-pub type OpaqueTyVariableKinds = BTreeMap<chalk_ir::OpaqueTyId<ChalkIr>, TypeKind>;
-pub type GeneratorKinds = BTreeMap<chalk_ir::GeneratorId<ChalkIr>, TypeKind>;
-pub type AssociatedTyLookups = BTreeMap<(chalk_ir::TraitId<ChalkIr>, Ident), AssociatedTyLookup>;
+pub type AdtIds = indexmap::IndexMap<Ident, chalk_ir::AdtId<ChalkIr>>;
+pub type FnDefIds = indexmap::IndexMap<Ident, chalk_ir::FnDefId<ChalkIr>>;
+pub type ClosureIds = indexmap::IndexMap<Ident, chalk_ir::ClosureId<ChalkIr>>;
+pub type TraitIds = indexmap::IndexMap<Ident, chalk_ir::TraitId<ChalkIr>>;
+pub type GeneratorIds = indexmap::IndexMap<Ident, chalk_ir::GeneratorId<ChalkIr>>;
+pub type OpaqueTyIds = indexmap::IndexMap<Ident, chalk_ir::OpaqueTyId<ChalkIr>>;
+pub type AdtKinds = indexmap::IndexMap<chalk_ir::AdtId<ChalkIr>, TypeKind>;
+pub type FnDefKinds = indexmap::IndexMap<chalk_ir::FnDefId<ChalkIr>, TypeKind>;
+pub type ClosureKinds = indexmap::IndexMap<chalk_ir::ClosureId<ChalkIr>, TypeKind>;
+pub type TraitKinds = indexmap::IndexMap<chalk_ir::TraitId<ChalkIr>, TypeKind>;
+pub type AutoTraits = indexmap::IndexMap<chalk_ir::TraitId<ChalkIr>, bool>;
+pub type OpaqueTyVariableKinds = indexmap::IndexMap<chalk_ir::OpaqueTyId<ChalkIr>, TypeKind>;
+pub type GeneratorKinds = indexmap::IndexMap<chalk_ir::GeneratorId<ChalkIr>, TypeKind>;
+pub type AssociatedTyLookups =
+    indexmap::IndexMap<(chalk_ir::TraitId<ChalkIr>, Ident), AssociatedTyLookup>;
 pub type AssociatedTyValueIds =
-    BTreeMap<(chalk_ir::ImplId<ChalkIr>, Ident), AssociatedTyValueId<ChalkIr>>;
-pub type ForeignIds = BTreeMap<Ident, chalk_ir::ForeignDefId<ChalkIr>>;
+    indexmap::IndexMap<(chalk_ir::ImplId<ChalkIr>, Ident), AssociatedTyValueId<ChalkIr>>;
+pub type ForeignIds = indexmap::IndexMap<Ident, chalk_ir::ForeignDefId<ChalkIr>>;
 
-pub type ParameterMap = BTreeMap<Ident, chalk_ir::WithKind<ChalkIr, BoundVar>>;
+pub type ParameterMap = indexmap::IndexMap<Ident, chalk_ir::WithKind<ChalkIr, BoundVar>>;
 
 pub type LowerResult<T> = Result<T, RustIrError>;
 

chalk-integration/src/lowering/env.rs

@@ -9,7 +9,7 @@ use chalk_solve::rust_ir::{
     GeneratorWitnessDatum, GeneratorWitnessExistential, OpaqueTyDatum, OpaqueTyDatumBound,
 };
 use rust_ir::IntoWhereClauses;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::HashSet;
 use std::sync::Arc;
 use string_cache::DefaultAtom as Atom;
 
@@ -141,23 +141,23 @@ impl ProgramLowerer {
     }
 
     pub fn lower(self, program: &Program, raw_ids: &Vec<RawId>) -> LowerResult<LoweredProgram> {
-        let mut adt_data = BTreeMap::new();
-        let mut adt_reprs = BTreeMap::new();
-        let mut adt_variances = BTreeMap::new();
-        let mut fn_def_data = BTreeMap::new();
-        let mut fn_def_variances = BTreeMap::new();
-        let mut closure_inputs_and_output = BTreeMap::new();
-        let mut closure_closure_kind = 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();
-        let mut associated_ty_data = BTreeMap::new();
-        let mut associated_ty_values = BTreeMap::new();
-        let mut opaque_ty_data = BTreeMap::new();
-        let mut generator_data = BTreeMap::new();
-        let mut generator_witness_data = BTreeMap::new();
-        let mut hidden_opaque_types = BTreeMap::new();
+        let mut adt_data = indexmap::IndexMap::new();
+        let mut adt_reprs = indexmap::IndexMap::new();
+        let mut adt_variances = indexmap::IndexMap::new();
+        let mut fn_def_data = indexmap::IndexMap::new();
+        let mut fn_def_variances = indexmap::IndexMap::new();
+        let mut closure_inputs_and_output = indexmap::IndexMap::new();
+        let mut closure_closure_kind = indexmap::IndexMap::new();
+        let mut closure_upvars = indexmap::IndexMap::new();
+        let mut trait_data = indexmap::IndexMap::new();
+        let mut well_known_traits = indexmap::IndexMap::new();
+        let mut impl_data = indexmap::IndexMap::new();
+        let mut associated_ty_data = indexmap::IndexMap::new();
+        let mut associated_ty_values = indexmap::IndexMap::new();
+        let mut opaque_ty_data = indexmap::IndexMap::new();
+        let mut generator_data = indexmap::IndexMap::new();
+        let mut generator_witness_data = indexmap::IndexMap::new();
+        let mut hidden_opaque_types = indexmap::IndexMap::new();
         let mut custom_clauses = Vec::new();
 
         for (item, &raw_id) in program.items.iter().zip(raw_ids) {
@@ -175,7 +175,7 @@ impl ProgramLowerer {
                 generator_ids: &self.generator_ids,
                 generator_kinds: &self.generator_kinds,
                 associated_ty_lookups: &self.associated_ty_lookups,
-                parameter_map: BTreeMap::new(),
+                parameter_map: indexmap::IndexMap::new(),
                 auto_traits: &self.auto_traits,
                 foreign_ty_ids: &self.foreign_ty_ids,
             };

chalk-integration/src/lowering/program_lowerer.rs

@@ -15,87 +15,89 @@ use chalk_solve::rust_ir::{
 };
 use chalk_solve::split::Split;
 use chalk_solve::RustIrDatabase;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::HashSet;
 use std::fmt;
 use std::sync::Arc;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Program {
     /// From ADT name to item-id. Used during lowering only.
-    pub adt_ids: BTreeMap<Identifier, AdtId<ChalkIr>>,
+    pub adt_ids: indexmap::IndexMap<Identifier, AdtId<ChalkIr>>,
 
     /// For each ADT:
-    pub adt_kinds: BTreeMap<AdtId<ChalkIr>, TypeKind>,
+    pub adt_kinds: indexmap::IndexMap<AdtId<ChalkIr>, TypeKind>,
 
-    pub adt_variances: BTreeMap<AdtId<ChalkIr>, Vec<Variance>>,
+    pub adt_variances: indexmap::IndexMap<AdtId<ChalkIr>, Vec<Variance>>,
 
-    pub fn_def_ids: BTreeMap<Identifier, FnDefId<ChalkIr>>,
+    pub fn_def_ids: indexmap::IndexMap<Identifier, FnDefId<ChalkIr>>,
 
-    pub fn_def_kinds: BTreeMap<FnDefId<ChalkIr>, TypeKind>,
+    pub fn_def_kinds: indexmap::IndexMap<FnDefId<ChalkIr>, TypeKind>,
 
-    pub fn_def_variances: BTreeMap<FnDefId<ChalkIr>, Vec<Variance>>,
+    pub fn_def_variances: indexmap::IndexMap<FnDefId<ChalkIr>, Vec<Variance>>,
 
-    pub closure_ids: BTreeMap<Identifier, ClosureId<ChalkIr>>,
+    pub closure_ids: indexmap::IndexMap<Identifier, ClosureId<ChalkIr>>,
 
-    pub closure_upvars: BTreeMap<ClosureId<ChalkIr>, Binders<Ty<ChalkIr>>>,
+    pub closure_upvars: indexmap::IndexMap<ClosureId<ChalkIr>, Binders<Ty<ChalkIr>>>,
 
-    pub closure_kinds: BTreeMap<ClosureId<ChalkIr>, TypeKind>,
+    pub closure_kinds: indexmap::IndexMap<ClosureId<ChalkIr>, TypeKind>,
 
     /// For each generator
-    pub generator_ids: BTreeMap<Identifier, GeneratorId<ChalkIr>>,
+    pub generator_ids: indexmap::IndexMap<Identifier, GeneratorId<ChalkIr>>,
 
-    pub generator_kinds: BTreeMap<GeneratorId<ChalkIr>, TypeKind>,
+    pub generator_kinds: indexmap::IndexMap<GeneratorId<ChalkIr>, TypeKind>,
 
-    pub generator_data: BTreeMap<GeneratorId<ChalkIr>, Arc<GeneratorDatum<ChalkIr>>>,
+    pub generator_data: indexmap::IndexMap<GeneratorId<ChalkIr>, Arc<GeneratorDatum<ChalkIr>>>,
 
-    pub generator_witness_data: BTreeMap<GeneratorId<ChalkIr>, Arc<GeneratorWitnessDatum<ChalkIr>>>,
+    pub generator_witness_data:
+        indexmap::IndexMap<GeneratorId<ChalkIr>, Arc<GeneratorWitnessDatum<ChalkIr>>>,
 
     /// From trait name to item-id. Used during lowering only.
-    pub trait_ids: BTreeMap<Identifier, TraitId<ChalkIr>>,
+    pub trait_ids: indexmap::IndexMap<Identifier, TraitId<ChalkIr>>,
 
     /// For each trait:
-    pub trait_kinds: BTreeMap<TraitId<ChalkIr>, TypeKind>,
+    pub trait_kinds: indexmap::IndexMap<TraitId<ChalkIr>, TypeKind>,
 
     /// For each ADT:
-    pub adt_data: BTreeMap<AdtId<ChalkIr>, Arc<AdtDatum<ChalkIr>>>,
+    pub adt_data: indexmap::IndexMap<AdtId<ChalkIr>, Arc<AdtDatum<ChalkIr>>>,
 
-    pub adt_reprs: BTreeMap<AdtId<ChalkIr>, Arc<AdtRepr<ChalkIr>>>,
+    pub adt_reprs: indexmap::IndexMap<AdtId<ChalkIr>, Arc<AdtRepr<ChalkIr>>>,
 
-    pub fn_def_data: BTreeMap<FnDefId<ChalkIr>, Arc<FnDefDatum<ChalkIr>>>,
+    pub fn_def_data: indexmap::IndexMap<FnDefId<ChalkIr>, Arc<FnDefDatum<ChalkIr>>>,
 
     pub closure_inputs_and_output:
-        BTreeMap<ClosureId<ChalkIr>, Binders<FnDefInputsAndOutputDatum<ChalkIr>>>,
+        indexmap::IndexMap<ClosureId<ChalkIr>, Binders<FnDefInputsAndOutputDatum<ChalkIr>>>,
 
     // Weird name, but otherwise would overlap with `closure_kinds` above.
-    pub closure_closure_kind: BTreeMap<ClosureId<ChalkIr>, ClosureKind>,
+    pub closure_closure_kind: indexmap::IndexMap<ClosureId<ChalkIr>, ClosureKind>,
 
     /// For each impl:
-    pub impl_data: BTreeMap<ImplId<ChalkIr>, Arc<ImplDatum<ChalkIr>>>,
+    pub impl_data: indexmap::IndexMap<ImplId<ChalkIr>, Arc<ImplDatum<ChalkIr>>>,
 
     /// For each associated ty value `type Foo = XXX` found in an impl:
     pub associated_ty_values:
-        BTreeMap<AssociatedTyValueId<ChalkIr>, Arc<AssociatedTyValue<ChalkIr>>>,
+        indexmap::IndexMap<AssociatedTyValueId<ChalkIr>, Arc<AssociatedTyValue<ChalkIr>>>,
 
     // From opaque type name to item-id. Used during lowering only.
-    pub opaque_ty_ids: BTreeMap<Identifier, OpaqueTyId<ChalkIr>>,
+    pub opaque_ty_ids: indexmap::IndexMap<Identifier, OpaqueTyId<ChalkIr>>,
 
     /// For each opaque type:
-    pub opaque_ty_kinds: BTreeMap<OpaqueTyId<ChalkIr>, TypeKind>,
+    pub opaque_ty_kinds: indexmap::IndexMap<OpaqueTyId<ChalkIr>, TypeKind>,
 
     /// For each opaque type:
-    pub opaque_ty_data: BTreeMap<OpaqueTyId<ChalkIr>, Arc<OpaqueTyDatum<ChalkIr>>>,
+    pub opaque_ty_data: indexmap::IndexMap<OpaqueTyId<ChalkIr>, Arc<OpaqueTyDatum<ChalkIr>>>,
 
     /// Stores the hidden types for opaque types
-    pub hidden_opaque_types: BTreeMap<OpaqueTyId<ChalkIr>, Arc<Ty<ChalkIr>>>,
+    pub hidden_opaque_types: indexmap::IndexMap<OpaqueTyId<ChalkIr>, Arc<Ty<ChalkIr>>>,
 
     /// For each trait:
-    pub trait_data: BTreeMap<TraitId<ChalkIr>, Arc<TraitDatum<ChalkIr>>>,
+    pub trait_data: indexmap::IndexMap<TraitId<ChalkIr>, Arc<TraitDatum<ChalkIr>>>,
 
     /// For each trait lang item
-    pub well_known_traits: BTreeMap<WellKnownTrait, TraitId<ChalkIr>>,
+    pub well_known_traits: indexmap::IndexMap<WellKnownTrait, TraitId<ChalkIr>>,
 
     /// For each associated ty declaration `type Foo` found in a trait:
-    pub associated_ty_data: BTreeMap<AssocTypeId<ChalkIr>, Arc<AssociatedTyDatum<ChalkIr>>>,
+    pub associated_ty_data:
+        indexmap::IndexMap<AssocTypeId<ChalkIr>, Arc<AssociatedTyDatum<ChalkIr>>>,
 
     /// For each user-specified clause
     pub custom_clauses: Vec<ProgramClause<ChalkIr>>,
@@ -104,7 +106,7 @@ pub struct Program {
     pub object_safe_traits: HashSet<TraitId<ChalkIr>>,
 
     /// For each foreign type `extern { type A; }`
-    pub foreign_ty_ids: BTreeMap<Identifier, ForeignDefId<ChalkIr>>,
+    pub foreign_ty_ids: indexmap::IndexMap<Identifier, ForeignDefId<ChalkIr>>,
 }
 
 impl Program {

chalk-integration/src/program.rs

@@ -19,7 +19,7 @@ use chalk_solve::Solver;
 use salsa::Database;
 use std::clone::Clone;
 use std::cmp::{Eq, PartialEq};
-use std::collections::BTreeMap;
+// use std::collections::BTreeMap;
 use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 use std::sync::Mutex;
@@ -40,7 +40,10 @@ pub trait LoweringDatabase:
     /// one another (the "specialization priorities").
     fn coherence(
         &self,
-    ) -> Result<BTreeMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>, ChalkError>;
+    ) -> Result<
+        indexmap::IndexMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>,
+        ChalkError,
+    >;
 
     fn orphan_check(&self) -> Result<(), ChalkError>;
 
@@ -142,12 +145,13 @@ fn orphan_check(db: &dyn LoweringDatabase) -> Result<(), ChalkError> {
 
 fn coherence(
     db: &dyn LoweringDatabase,
-) -> Result<BTreeMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>, ChalkError> {
+) -> Result<indexmap::IndexMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>, ChalkError>
+{
     let program = db.program_ir()?;
     let solver_choice = db.solver_choice();
     let priorities_map = tls::set_current_program(&program, || -> Result<_, ChalkError> {
         let solver_builder = || solver_choice.into_solver();
-        let priorities_map: Result<BTreeMap<_, _>, ChalkError> = program
+        let priorities_map: Result<indexmap::IndexMap<_, _>, ChalkError> = program
             .trait_data
             .keys()
             .map(|&trait_id| {

chalk-integration/src/query.rs

@@ -109,16 +109,19 @@ where
     fn build_specialization_forest(&self) -> Result<Graph<ImplId<I>, ()>, CoherenceError<I>> {
         // The forest is returned as a graph but built as a GraphMap; this is
         // so that we never add multiple nodes with the same ItemId.
-        let mut forest = DiGraphMap::new();
+        let mut forest = DiGraph::new();
 
         // Find all specializations (implemented in coherence/solve)
         // Record them in the forest by adding an edge from the less special
         // to the more special.
         self.visit_specializations_of_trait(|less_special, more_special| {
-            forest.add_edge(less_special, more_special, ());
+            let l = forest.add_node(less_special);
+            let m = forest.add_node(more_special);
+
+            forest.add_edge(l, m, ());
         })?;
 
-        Ok(forest.into_graph())
+        Ok(forest)
     }
 
     // Recursively set priorities for those node and all of its children.