Use TrieMaps to speed up type class instance lookup

[ezyang]

Currently, type class instance resolution is performed by doing a map lookup by type class, and then linearly matching against every instance. This is not great, and for a while, simonpj has been keen on using the TrieMap data structure in GHC, which has been previously used to implement a map from CoreExprs to various things (e.g. for CSE). What makes this a little tricky is that instance lookup isn't intended to be an exact match: we should unify so-called template variables and provide a substitution; furthermore, there may be multiple matches.

After some prototyping, I think we should be able to make this work. The primary refactoring we have to do is *maintain the key associated with an entry in a TrieMap*. Unlike the current uses of TrieMaps, where it's acceptable to lose the underlying key associated with an entry in the TrieMap, we need to know the key when doing unification, because it may be reported in the substitution. There are a few knock-on effects of this:

• We should add a method foldWithKeyTM :: (Key m -> a -> b -> b) -> m a -> b -> b to the TrieMap type class.
• We need a variant of UniqFM which tracks the original key that was used. I propose we name it UniqKM (unique key map). A number of implementations of TrieMap need to be adjusted to use this instead of UniqFM.
• Why can't we just represent keyed TrieMaps as TypeMap (Type, a)? It might be possible. An insurmountable difficulty would be if we need to know the partial key PRIOR to having finished traversing the TrieMap; however, for the parts of the unification algorithm I've implemented, this does not seem to be the case. The primary actual difficulty is that Type uses a named representation, whereas TypeMap keys are on-the-fly deBruijn numbered. There would at least be some annoying fiddliness.
• Reversing the deBruijn numbered key into a Type is a bit goofy. Essentially, you need the reverse of the current CmEnv: a mapping from de Bruijn levels to the Var you've decided to allocate. (I've called this CfEnv in my prototype)
• When we represent a TrieMap binder, we have to put the binder map on the OUTSIDE, as opposed to the inside as it is currently. Otherwise, we can't update the CfEnv with the new mapping before making the recursive call to fold-with-key.

I'll attach the standalone Haskell file I used to prototype this, wherein I copy-pasted a lot of Haskell from GHC's source and implemented fuzzy map on a simplified version of Type.

[ezyang]

Standalone prototype

[simonpj]

Great stuff.

• I suggest making this a new data type, rather than simply replacing TrieMap. The latter is simple, uniform, and efficient. Let's not clutter it up!

• But perhaps TrieMap itself can be modestly improved, by giving access to the key.

Simon

[ezyang]

Simon, here is my proposal: we make a new type class TrieKeyMap whose parent is TrieMap and defines just the new foldWithKeyTM method. Most data types will be instances of both, but not, for example, UniqFM.

Update: Oh, your previous comment makes more sense if I s/TrieMap/CoreMap/ in your first bullet point. I agree: we should make a new CoreMap for type class lookups. But what is your feeling about the type class design? Here is the sub type-class version: ​https://phabricator.haskell.org/D604

[simonpj]

If we need the keys, then yes, class TrieKeyMap looks good. But I'm not yet certain that we do. I need to look at your straw-man TrieMap.hs.

See #9960 for an orthogonal issue that also concerns TrieMap.

[ezyang]

simonpj: The key (ahem) part where we need the keys is we compute substitutions. In the prototype, it's line 696. Essentially, you have a sub-trie, and want to extract all the expressions corresponding to entries in the sub-trie, since those are your substitutions.

[simonpj]

I've taken a look at your straw man. I really don't understand it.

What is TyBox??

Standing back, I'm sure that the template variables should be deBruijn-ised. And a VarMap should look like

data VarMap a = VM { vm_bvar   :: BoundVarMap a  -- Bound variable
                   , vm_fvar   :: VarEnv a       -- Free variable
                   , vm_tvar   :: TmplVarMap a } -- Template variable

type TmplVar = Int  -- Like BoundVar but a different namespace
type TmplVarMap = IntMap

That is, in a key, a variable could be bound by a forall (vm_bvar), or free (bm_fvar), or a template variable (vm_tvar). So a key (T (forall a. a -> b -> c)), where b is a template variable and c is free will effectively be treated as

  T (forall. bv1 -> tv1 -> c)

where I used bv1 for BoundVar 1 and tv1 for TmplVar 1.

When doing a matchT we must maintain a

• CmEnv, for deBruijn of the foralls
• TmplVarMap Type for the substitution so far, mapping template variables to types.

When extending the TmplVarMap must check that the Type does not mention any CmEnv-bound variables. (A kind of occurs/escape check.) When matching, if the key is say (tv1,tv1), then the second time we encounter tv1 it'll already be bound in the substitution, so we must check for equality with the equality function mentioned in #9960.

Does that make sense? Happy to talk more, or Skype.

Simon

[ezyang]

Some clarifying comments from the Skype conversation:

1. The reason why the first prototype doesn't make sense was because I accidentally thought the template variables were in the expression I was matching with, rather than in the *keys* of the TrieMap. This was also why I believed a key was necessary.

2. TyBox in the prototype is a hack to avoid having circularity when specifying the kind of box. We don't really need kinds for the prototype but I was trying to keep the data structures as close as possible to GHC proper.

[ezyang]

So, here is an irritating problem with trying to deBruijn number template variables on the fly.

First, let's establish that it is desirable for template variables to be deBruijn numbered in order of appearance. This is because if you have instanc heads C a b and C b a (a and b being template variable), it would be rather silly if we de Bruijinized these as C 0 1 and C 1 0, when really they're the same. We don't know, a priori, what order the "template quantifiers" should be, so it would be great to defer this choice until we've explored the expression deeply enough to tell.

The fly in the ointment is this: with template variables, environment extension happens when we encounter a template variable (as opposed to de Bruijn numbering lambda/foralls, where extension occurs when we encounter a lambda/forall). So, if I am handling C ????? a (where a is a template variable), I do not know what the index of a will be until I have processed ????? and determined how many template variables show up for the first time there. Accordingly, my lookup function must return the new template variable environment, in the same way the unifying match function returns the substitution.

This is not completely terrible, since we already were going to need to return extra information on lookup for the unifying match, but it does mean we can't use any of the existing lookup/insert functions to manage template variables. C'est la vie?

[ezyang]

I decided to solve this problem by assuming template var quantifiers are in the same order as the list of template variables.

[simonpj]

Edward, let's Skype about this. I don't understand the ramifications yet. Ping me.

Simon

[ezyang]

So, I finished the InstMap implementation, and was about to wire up it with InstEnv when I realized there is one more thing to implement.

You see, if the template variable doesn't match, the instance lookup code will TRY AGAIN to unify, this time properly unifying all variables and not just matching template variables, in order to provide the list of non-matching but unifying instances. So it looks like we'll need to implement this too, because otherwise the TrieMap lookup won't provide enough candidates.

An added complication: since we might unify variables from the key, we will sometimes need to be able to pull up the type embedded in the map, which means we will need to revive foldWithKey in some situations.

[simonpj]

Yes, and that is indeed a pain.

Perhaps we should always do unification, which generalises matching, and distinguish the match cases post-hoc.

Re the key thing, I think we may just want to keep (key,value) pairs at the leaves, rather than just values. Better than reconstructing the key from the tree.

Simon

[ezyang]

It's been two years since I last poked this, and since then Richard Eisenberg merged kind equalities. This means that the match implementation I prototyped no longer works: matching has been rolled up with a very general unification algorithm, which keeps track of kind coercions so that the resulting substitutions are well typed. So I am uncertain whether or not we can still carve out the "matching" portion of the algorithm without having to handle all of this extra logic.

[ezyang]

There is a more up-to-date version of the code at ​https://github.com/ezyang/triemap/blob/master/TrieMap.hs

[goldfire]

The new algorithm should be just as amenable to this new implementation as the old one. Yes, it's all tied in with unification (and there are kind equalities) but don't let that scare you. In particular, the kind coercion passed down through the algorithm is only important in the TyVar case (and, perhaps, in the ForAllTy case, where it's propagated. Does this ever actually trigger? Who knows). So I believe this shouldn't cause any undue difficulty in realizing your idea (which I think is great).

[ezyang]

I picked this up again. I looked carefully at where the kind coercion is used. Specifically:

• (TyVar case) When matching in uVar, it is used to get the two types to the same kind so that eqType may succeed.
• (TyVar case) In uUnrefined, it is used to ensure that the TvSubst is well kinded
• (CoercionTy case) It is used to extend the CvSubst

So, although you don't say it explicitly, I think that I will need to pass kco around in the TrieMap implementation so that I can implement these methods correctly. One source of awkwardness here is that the TypeMap doesn't currently store CastTy, which means any coercions induced by them will be lost to the ether (I think what needs to be done is just to add back a map for CastTy constructor in TypeMap.)

[ezyang]

It also looks like I'm going to have to undo the trieMapView trick, because when unifying I need to distinguish between injective and non-injective parameters, and if it's a pile of nested AppTys that is much more difficult to do. (It also doesn't make much sense semantically, since type family applications are supposed to be saturated. Not that you're ever supposed to have a type family in an instance head, but mumble mumble #13262, and note that TypeMap is used by CoreMap, where type families definitely could show up.)

[goldfire]

There shouldn't be a map for CastTy in the trie. See Note [Non-trivial definitional equality] in TyCoRep. Remember that ty |> co is equal to ty if the kinds are the same. A more interesting scenario: (t1 |> co1) (t2 |> co2) is equal to t1 t2 if the kinds are the same. Note that it might be that neither co1 nor co2 is reflexive!