Reduce next-solver memory usage by interning CanonicalQueryInput

Project: rust-lang/rust | Language: Rust | Concepts: memory-management, data-structures | Difficulty: senior
Link: https://github.com/rust-lang/rust/pull/162031

The problem

The next-generation trait solver in the Rust compiler was using a large amount of memory, especially when compiling crates that hit the trait solver recursion limit. The main culprit was the storage of CanonicalQueryInput values in large hash maps (such as in the search graph), where each input was stored as a full, heap-allocated structure. This led to excessive memory usage (e.g., ~15GiB for bevy_render), making the compiler prone to out-of-memory errors on large or complex crates.

Why it's non-trivial

A naive fix would be to simply try to reduce the size of the data structures or clear them more aggressively. However, the real issue is that many CanonicalQueryInput values are deeply duplicated, and deduplication is non-trivial: you need to ensure that logically identical inputs are physically shared, which requires interning. Doing this safely and efficiently, and threading the new interned type through all the right places in the codebase, is complex and error-prone.

The approach

The PR introduces a new interned wrapper for CanonicalQueryInput, threads this interned type through the trait solver, and updates the interner trait and its implementations accordingly.

1. Define and use an interned CanonicalInput type:

#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash, StableHash)]
pub struct CanonicalInput<'tcx>(pub(crate) Interned<'tcx, CanonicalInputData<TyCtxt<'tcx>>>);

impl<'tcx> std::ops::Deref for CanonicalInput<'tcx> {
    type Target = CanonicalInputData<TyCtxt<'tcx>>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

This defines a newtype wrapper around an interned pointer to the canonical input data, enabling cheap comparison, hashing, and sharing.

2. Add a new interner slot and method for canonical inputs:

pub struct CtxtInterners<'tcx> {
    ...
+    canonical_inputs: InternedSet<'tcx, CanonicalInputData<TyCtxt<'tcx>>>,
}

direct_interners! {
    ...
+    canonical_inputs: intern_canonical_input(CanonicalInputData<TyCtxt<'tcx>>): CanonicalInput -> CanonicalInput<'tcx>,
}

This adds a dedicated interner for canonical inputs, so all identical inputs are stored only once.

3. Extend the Interner trait to support interning canonical inputs:

pub trait Interner: ... {
    ...
+    type CanonicalInput: Copy + Debug + Hash + Eq + Deref<Target = CanonicalInputData<Self>>;
+    fn mk_canonical_input(self, data: CanonicalInputData<Self>) -> Self::CanonicalInput;
}

This makes canonical input interning a first-class part of the interner abstraction, so it can be used generically.

4. Update all code to use the interned type instead of the raw data:

For example, in the canonicalization logic:

-pub(super) fn canonicalize_goal<D, I>(...) -> (ThinVec<I::GenericArg>, CanonicalInput<I, I::Predicate>)
+pub(super) fn canonicalize_goal<D, I>(...) -> (ThinVec<I::GenericArg>, I::CanonicalInput)
...
-    let query_input =
-        ty::CanonicalQueryInput { canonical, typing_mode: TypingModeEqWrapper(typing_mode) };
+    let query_input = delegate.cx().mk_canonical_input(ty::CanonicalQueryInput {
+        canonical,
+        typing_mode: TypingModeEqWrapper(typing_mode),
+    });

This ensures that every canonical input is interned at creation, and all downstream consumers use the interned type.

5. Update all trait bounds and function signatures to use the new interned type:

For example, in the search graph and evaluation contexts:

-    fn compute_goal(
-        search_graph: &mut SearchGraph<D>,
-        cx: I,
-        input: CanonicalInput<I>,
-        inspect: &mut Self::ProofTreeBuilder,
-    ) -> (QueryResult<I>, AccessedOpaques<I>) {
+    fn compute_goal(
+        search_graph: &mut SearchGraph<D>,
+        cx: I,
+        input: I::CanonicalInput,
+        inspect: &mut Self::ProofTreeBuilder,
+    ) -> (QueryResult<I>, AccessedOpaques<I>) {

All trait and function signatures that previously took a raw canonical input now take the interned version.

From the review thread

laundmo: I'm not so sure anymore. I've been considering more and i think the fact NestedGoals stores a reference now is a drawback over, for example, storing CanonicalInputs in a Vec and handing out indices, kinda like you mentioned in the zulip at some point.

To me it seems likely that's the direction where other improvements can be made.

lcnr: This PR feels small enough for us to easily revert this as part of a PR using indices, and using indices likely takes more than a few days, so preventing people compiling bevy from getting an OOM is probably worth it anyways

lcnr: I do think we should be able to remove some QueryKey impls, but that's fine and can happen separately

Takeaways

• Interning is a powerful technique for deduplicating large, frequently-repeated data structures, drastically reducing memory use at the cost of some indirection.
• Threading a new interned type through a large codebase requires careful updates to all type signatures and trait bounds to ensure correctness and avoid accidental mixing of interned and non-interned data.
• It is sometimes worth landing a simple, high-impact fix (like interning) even if further optimizations (like using indices instead of references) are possible, especially when the immediate benefit is preventing critical failures such as out-of-memory errors.