Flattening the contexts for fun and profit

I wanted to write up some intermediate thoughts on a proposed reorganization for the front-end of the type checker. The eventual goal is of course cleaner, faster, and less buggy code – particularly around the area of associated type projections and normalizations, which are currently quite a pain to get right.

@jroesch has already done some of the steps I describe here in PR #26582. I started this write-up a week or two ago with the goal of posting it before he got too far, since I’d love to get feedback on this approach! Unfortunately, I didn’t quite finish it up before I had to leave for the Mozilla work week, but anyway I’m posting it now, and I’d still to get feedback. :blush: After all, PR #26582 is only the first few steps.

Brief summary of how things work today

The current model for normalization is “eager” in the sense that it tries to normalize types that may contain projections as soon as we do any substitutions. So, for example, if you call a function like:

fn get_next<I:Iterator>(iter: I) -> Option<I::Item> { i.next() }

The compiler will handle this by first substituting a fresh type variable for I, let’s call it $I, and then trying to normalize the types in the signature. In this case, that includes $I::Item – given that we don’t know what $I is yet, we will wind up creating a fresh variable $J and adding the pending constraint $I::Item == $J. This means that the type and pending obligations which result are roughly as follows:

fn($I) -> Option<$J>
$I: Iterator
$I::Item == $J

Now, until we know that $I is, those two obligations cannot be satisfied. It’s important for overall performance that we don’t keep trying to satisfy them. We currently have a pretty hacky scheme for this – basically we add the obligations to a vector, and we track how many of the items in that vector have already been tested. We then avoid testing them again.

The problem with this is that we don’t have an easy way of knowing when we might make progress again. Of course, we can always wait till the end and then do a final sweep, but sometimes it’s important to resolve the pending traits in order to find out the value of a type variable. Say for example that we have unified $I with vec::IntoIter<String>. In that case, if we tried to resolve $I::Item == $J, we would find out that $J == String. But if we don’t try, we’ll never know.

Currently we handle this by waiting until we get in a situation where it is important to know what $J is (for example, when the program tries to call a method on a variable of type $J). In that case, we go back and revisit the deferred work items. This is effective at edging off the worst case, but could be a lot better.

Some problems I see and their cause

There are a few problems I see with the current setup:

  1. Eagerly normalizing is a pain. It’s easy to forget, and it requires us to thread a lot of plumbing through everywhere, since normalizing can create new obligations and those have to make their way back to the central listing.
  2. When an obligation turns out to be ambiguous due to the presence of an inference variable, it would be great if we could “park” it permanently until that inference variable is unified, and then unpark it. The current counter scheme is trying to get at this, but it does so inefficiently.

I think a big part of the problem is due to the current decomposition of the type-checking state into a variety of random little contexts. I’m thinking of ty::ctxt, InferCtxt, FulfillmentContext, ParamEnv, Typer, SelectionContext, FnCtxt, Inherited, etc. Having written a good number of those, I can tell you that their structure, while I put a lot of work into it, is largely influenced by the order in which they were written, rather than a logical "top-down" decomposition.

In particular, the current breakdown makes it hard to (a) uniformly apply transactions and (b) “flow” data nicely from part to part. For example, if a subtyping relation produced an obligation (as some future type system extensions call for), it would be kind of a pain to get it registered into the appropriate FulfillmentContext.

Lazy normalization

I would like to change from eager normalization to lazy normalization. The idea is to work roughly like we handle inference variables right now. For the most part, the inference variable type – even when unified – can float around and be propagated. But at various points in the compiler, such as when relating types in middle::infer::Sub or middle::infer::Glb etc, we will “force” the variable, either unifying it with something or else substituting the previously unified value.

We can take a similar approach with projections. Basically, we leave them alone most of the time, but in some contexts we will normalize them before proceeding. (Obviously we’ll need an efficient cache to avoid endlessly repeating work.)

I am not sure of the full set of contexts where we would have to normalize, I think it will be roughly the same as the places where we choose to resolve inference variables (in fact, I’d probably make the "inference variable resolution" functions also normalize projections as part of what they do), and of course middle::infer::Sub and friends.

However, for lazy normalization to work, the inference context needs to be able to easily call into the project code. To do this, it must (currently) be able to supply a ParamEnv, Typer, etc. This is a pain given the way that the contexts are currently layered, which leads us to the next bullet point.

Reshuffling the contexts

So how can we reorganize the various contexts to improve the situation? I’d basically like to wind up with one unified TypingContext that more-or-less replaces FnCtxt, and which can easily be used from anywhere that needs it.

I went through the code and came up with this organization. @jroesch has been hard at work on implementing it, and has made good progress towards the first part: https://github.com/rust-lang/rust/pull/26582

One of the goals here is to remove the Typer and ClosureTyper traits. The only reason that these traits exist is because, during typeck, we want to draw our data from the FnCtxt, whereas during the rest of compilation we want to draw our data from the global tables in the tcx. This is because we always ensure that unresolved inference variables are confined to the FnCtxt (but see the next section).

Rather than having these traits, I would like to refactor all the relevant tables into a Tables struct. We would then take an &Tables to choose where we draw our data from, rather than an Typer object. Eventually, this fits in the plans to have separate tables per fn, because if we had a MethodData struct that stored everything relevant to a particular function, this MethodData struct would be exactly the set of Tables we need.

Meanwhile, all the code that today takes a Typer (or ClosureTyper), tcx, and InferCtxt, would instead just take an InferCtxt. InferCtxt would wrap a given &Tables object and ParameterEnvironment (and tcx, naturally).

Question: Keep the “pre-typeck” vs “post-typeck” split?

One mildly open question is whether we should just give up on having data that contains unresolved inference variables confined to the FnCtxt. The original motivation here was that we would use a different type for types that may contain TyInfer. This would guarantee that we have completely resolved all types. I am no longer motivated by that particular goal, in that I don’t think this is a large source of bugs relative to the pain caused by the split. So I was thinking of advocating that we just ALWAYS store the data in the tcx.

But I’ve been rethinking that because it’d be really useful to be able to have a tree of type arenas, so that types containing "ephmeral" content such as inference variables or skolemization by-products could be kept separate from the “global” types and then freed. This would also be important for parallelism, because we’re going to need to be able to create types from distinct threads and merge them back into the global context. To do that, we probably want the types during typeck to be Ty<'typeck> whereas the global types are Ty<'tcx>.

If we wanted to do this more limited sort of thing, it would mean that we’d have to rewrite a lot of code to be generic over the lifetime of the types, trait-refs, etc:

fn resolve_traits<'pass,'tcx>(data: Ty<'pass>, ...)

Inference contexts would then create

@eddyb, I know you’ve given some thought to this, does this sound like roughly what you had in mind?

1 Like

I’ve already started!
The type pretty-printing was a step towards it and now https://github.com/rust-lang/rust/pull/26575 is waiting for a review.
I also have some changes that are not yet in the PR, trying to get struct TyCx<'a, 'global: 'local, 'local: 'a> working, without an actual pair of contexts backing it.
So far it looks easier than I thought, though to actually use some API with something other than TyCx<'a, 'tcx, 'tcx>, it has to be converted first.
I’m hoping we can get away without a lot of churn, only more lifetimes in the right contexts.

Although, being able to use '_ inside types in function arguments would greatly simplify a lot of the changes, including some in that PR.

Yes, I saw, I'll try to do that today!

Right, so we can work this out over time, but I imagine that TyCx here might eventually be the TypingContext (nee InferCtxt) that I am referring to. That is, the idea would be that you have a global type context (tcx today). This represents the program as a whole. When you want to do inference, checking subtyping, and generally work with types, you create a typing context. This brings together:

  • the initial environment
  • an appropriate set of tables for things like closure types etc (always the tables from some particular fn)
  • a local arena for temporary types and so forth

This local context (and, notably, its arena) will have a lifetime 'env (vs the global lifetime 'tcx). We have to be careful and ensure that the lifetimes for Ty etc are contravariant, so that any types, trait-references, etc that we fetch from the global context (e.g. TraitRef<'tcx>) can be upcast to TraitRef<'env> and thus intermingle with the shorter-lived types and so forth that we'll be producing as part of this inference context.

In the case of typeck in particular, we'll take these shorter lived types and migrate them out to the global type context at the end, which obviously builds on the code you wrote to dynamically check whether a Ty<'tcx> belongs to a particular table.

Sound about right?

Can you elaborate here on just what you expect '_ to do -- a fresh lifetime parameter?

Yes, the same as omitting it in an optional location. &'a ty::ctxt<'tcx> became TyCx<'a, 'tcx> but there is no TyCx<'_, 'tcx> to replace &ty::ctxt<'tcx>.
It only gets worse, as TyCx<'_, '_, 'tcx> is the right type to use for working with Ty<'tcx> and TyCx<'_, 'tcx, '_> is the right type to use when modifying the global context (passes other than typeck would always deal with TyCx<'_, 'tcx, 'tcx>).

For RFC and initial implementation on '_, see https://github.com/rust-lang/rfcs/pull/1177.

I wouldn’t really like to merge FnCtxt with the typing contexts - it contains lots of function-related stuff, and I like the current state where the value-level - i.e. Expr, Stmt, etc. - is separate from the type-level. Also SelectionContext does not really have any state (only the freshener, and I don’t think we need to keep the same freshener between different calls to select). We should probably merge ParameterEnvironment, InferCtxt, FulfillmentContext, and ClosureTyper (not that sure about mc::Typer) into a single TypingContext through.

By the way, only_new_obligations does not seem to be that much of a help - when compiling rustc it seems to save less than 10% of obligations processed (it does prevent O(n^2)-ness in some cases through).

I know that only_new_obligations isn’t a big help – it just takes the edge off certain bad cases. I think we can stop the bad behavior in a more uniform way if we do things the way I described. I cannot say if that will be a big help overall.

Another thing we could do for improving general performance, I suspect, is to stop checking the Sized obligations during typeck. I moved them in there so that things would be uniform, but I now regret that. I think it would be better to check them in a separate pass once all types are known, so that we don’t have to queue a lot of information waiting for inference. This is how things used to be done, and I think there still exists a pass doing some of that checking – used to be called kindck but i think it now has a more limited name.

_: Sized obligations aren’t that expensive to deal with, even if there are O(n) of them.

This topic was automatically closed 90 days after the last reply. New replies are no longer allowed.