WIP: Recognise result-type polymorphic join points by monomorphisation (#14620)
Implements a worker/wrapper transformation for result type polymorphic join points:
Note [Join point worker/wrapper]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Some let bindings that are 'AlwaysTailCalled' still need a bit of work to become
a join point. Here's an example:
let f :: forall a b. [a] -> forall c. b -> Maybe c -> [(a,c)]
f = <rhs>
in (<body> :: [(Bool, Char)])
Suppose @f@ is 'AlwaysTailCalled' in <body>. Its result type is polymorphic,
which is exactly the situation as in
Historic Note [The polymorphism rule of join points].
That Note also explains why we can't just turn @f@ into join point untouched.
So we need a transformation that monomorphises @f@ for its result type. Since we
have available the type of (the soon to be join-) <body>, we can match @[(a,c)]@
against @[(Bool, Char)]@ to get a substitution @[a ↦ Bool, c ↦ Char]@. We then
apply this substitution to the tyco binders in the type of @f@, as we ascend
from the result type:
* @Maybe c@ is substituted to @Maybe Char@ ('SubstBinder')
* @b@ is substituted to @b@ (unchanged, likewise 'SubstBinder')
* @forall c@ is in the domain of the substitution and thus will be
instantiated ('InstBinder')
* @[a]@ is substituted to @[Bool]@ ('SubstBinder')
* @forall b@ is substituted (note that in general its kind might mention @a@)
to @forall b@ ('SubstBinder')
* @forall a@ is in the domain of the substitution and thus will be
instantiated ('InstBinder')
Figuring out this list of 'JoinWorkerBinder's (which is the
'SubstBinder'/'InstBinder') is the job of 'matchJoinResTy'.
In the simple, non-polymorphic case, it returns a list of 'SubstBinder's, one
for each join binder. Otherwise, there is at least one 'InstBinder' that
indicates monorphisation of a polymorphic join result type.
The challenge is in rewriting all call sites of @f@ to match its new type,
dropping the instantiated type arguments. A typical use case for the
worker/wrapper transformation. Thus, we make @f@ a wrapper that rewrites to the
new worker join point f':
let f :: forall a b. [a] -> forall c. b -> Maybe c -> [(a,c)]
f = \@a @b (xs :: [a]) @c b (mc :: Maybe c) -> f' @b xs b mc
f' :: forall b. [Bool] -> b -> Maybe Char -> [(Bool, Char)]
f' = \@b (xs :: [Bool]) b mb -> <rhs> @Bool @b xs @Char b mb
in (<body> :: [(Bool, Char)])
Take note that @f@'s type did not change, but its new RHS is now actually
ill-typed. This doesn't matter as long as we manage to inline the wrapper
unconditionally at its call sites in <body>, where the arguments for @a@ and @c@
will always be @Bool@ and @Char@.
The join point worker @f'@ similarly instantiates @a@ and @c@ to @Bool@ and
@Char@. Its result type is monomorphic and it can be made into a join point.
The worker/wrapper split is carried out by 'joinPointBindings_maybe', but only if
there are any 'InstBinder's at all (In which case it returns the result
@JoinPointAfterMono@).
Equipped with the 'matchJoinResTy' result (InstBinder = I, SubstBinder = S)
[I Bool, S (b::*), S (_::[Bool]), I Char, S (_::b), S (_::Maybe Char)]
It builds the wrapper body of @f@ by applying the new worker binder @f'@ to
* Nothing if the corresponding 'JoinWorkerBinder' is @I _@
* @b@ if the corresponding 'JoinWorkerBinder' is @S _@ and @b@ is the old
lambda binder
It builds the worker body of @f'@ by applying the <rhs> to
* @ty@ if the corresponding 'JoinWorkerBinder' is @I ty@
* @tv@ if the corresponding 'JoinWorkerBinder' is @S (tv::..)@ (Named binder)
* @b@ if the corresponding 'JoinWorkerBinder' is @S (_::ty)@ (Anon binder)
and @b@ is the old lambda binder with its type updated to @ty@.
The result of @joinPointBinding_maybe@ is ultimately exported via @tryJoinWW@
and is used in the simple optimiser as well as the Simplifier, which both
inline the join point wrapper unconditionally (if present).