Fusing loops by specializing on functions with SpecConstr?

Alexis King lexi.lambda at gmail.com
Sat Apr 4 01:45:36 UTC 2020


I fiddled with alternative representations for a while and didn’t make
any progress—it was too easy to end up with code explosion in the
presence of any unknown calls—but I seem to have found a RULES-based
approach that works very well on the examples I’ve tried. It’s quite
simple, which makes it especially appealing!

I started by defining a wrapper around the `SF` constructor to attach
rules to:

    mkSF :: (a -> s -> Step s b) -> s -> SF a b
    mkSF = SF
    {-# INLINE CONLIKE [1] mkSF #-}

I  then changed the definitions of (.), (***), (&&&), (+++), and (&&&)
to use `mkSF` instead of `SF`, but I left the other methods alone, so
they just use `SF` directly. Then I defined two rewrite rules:

    {-# RULES
    "mkSF @((), _)" forall f s. mkSF f ((), s) =
      SF (\a s1 -> case f a ((), s1) of Step ((), s2) b -> Step s2 b) s
    "mkSF @(_, ())" forall f s. mkSF f (s, ()) =
      SF (\a s1 -> case f a (s1, ()) of Step (s2, ()) b -> Step s2 b) s
    #-}

That’s it. These two rules alone are enough to eliminate the redundant
tupling. Now the optimized version of `mapMaybeSF` is beautiful!

    mapMaybeSF = \ @ a @ b f -> case f of { SF @ s f2 s2 ->
      SF (\ a1 s1 -> case a1 of {
           Nothing -> case s1 of dt { __DEFAULT -> Step dt Nothing }
           Just x -> case f2 x s1 of {
             Step s2' c1 -> Step s2' (Just c1) }})
         s2 }

So unless this breaks down in some larger situation I’m not aware of, I
think this solves my problem without the need for any fancy SpecConstr
shenanigans. Many thanks to you, Sebastian, for pointing me in the right
direction!

Alexis


More information about the ghc-devs mailing list