Shortcut fusion for triemaps - haskell

Question

This problem arose when attempting to fuse away intermediate triemaps in Haskell.
Consider the trie for Peano natural numbers:
data Nat = Zero | Succ Nat
data ExpoNat a = ExpoNat (Maybe a) (ExpoNat a)
| NoExpoNat
We can easily define a fold on ExpoNat (it is essentially a list) and use foldr/build (a.k.a. finally tagless) to fuse away intermediate occurrencess of ExpoNat:
{-# NOINLINE fold #-}
fold :: (Maybe a -> b -> b) -> b -> ExpoNat a -> b
fold f z (ExpoNat x y) = f x (fold f z y)
fold f z NoExpoNat = z
{-# NOINLINE build #-}
build :: (forall b. (Maybe a -> b -> b) -> b -> b) -> ExpoNat a
build f = f ExpoNat NoExpoNat
{-# RULES "fold/build" forall f n (g :: forall b. (Maybe a -> b -> b) -> b -> b). fold f n (build g) = g f n #-}
As an example, we take match and appl from "Is there a way to generalize this TrieMap code?" and compose them such that ExpoNat is fused away. (Note that we must "strengthen the induction hypothesis" in appl.)
{-# INLINE match #-}
match :: Nat -> ExpoNat ()
match n = build $ \f z ->
let go Zero = f (Just ()) z
go (Succ n) = f Nothing (go n)
in go n
{-# INLINE appl #-}
appl :: ExpoNat a -> (Nat -> Maybe a)
appl
= fold (\f z -> \n ->
case n of Zero -> f
Succ n' -> z n')
(\n -> Nothing)
applmatch :: Nat -> Nat -> Maybe ()
applmatch x = appl (match x)
The fusion can be verified by inspecting Core with -ddump-simpl.
Now we would like to do the same for Tree.
data Tree = Leaf | Node Tree Tree
data TreeMap a
= TreeMap {
tm_leaf :: Maybe a,
tm_node :: TreeMap (TreeMap a)
}
| EmptyTreeMap
We are in trouble: TreeMap is a non-regular data type, and so it is not obvious how to write its corresponding fold/build pair.
Haskell Programming with Nested Types: A Principled Approach seems to have the answer (see the Bush type) but 4:30 AM seems to be too late for me to get it working. How is one supposed to write hfmap? Have there been further developments since?
A similar variant of this question has been asked in What's the type of a catamorphism (fold) for non-regular recursive types?

Answer 1

I worked on it some more and I now have working fusion, without using the generic gadgets from the paper.
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
module Tree where
data Tree = Leaf | Node Tree Tree
deriving (Show)
data ExpoTree a = ExpoTree (Maybe a) (ExpoTree (ExpoTree a))
| NoExpoTree
deriving (Show, Functor)
I derived most of the specialized types by taking the generic construction and then inlining type definitions until I bottomed out. I've kept the generic construction in here for ease of comparison.
data HExpoTree f a = HExpoTree (Maybe a) (f (f a))
| HNoExpoTree
type g ~> h = forall a. g a -> h a
class HFunctor f where
ffmap :: Functor g => (a -> b) -> f g a -> f g b
hfmap :: (Functor g, Functor h) => (g ~> h) -> (f g ~> f h)
instance HFunctor HExpoTree where
ffmap f HNoExpoTree = HNoExpoTree
ffmap f (HExpoTree x y) = HExpoTree (fmap f x) (fmap (fmap f) y)
hfmap f HNoExpoTree = HNoExpoTree
hfmap f (HExpoTree x y) = HExpoTree x (f (fmap f y))
type Alg f g = f g ~> g
newtype Mu f a = In { unIn :: f (Mu f) a }
instance HFunctor f => Functor (Mu f) where
fmap f (In r) = In (ffmap f r)
hfold :: (HFunctor f, Functor g) => Alg f g -> (Mu f ~> g)
hfold m (In u) = m (hfmap (hfold m) u)
An Alg ExpoTreeH g can be decomposed into a product of two natural transformations:
type ExpoTreeAlg g = forall a. Maybe a -> g (g a) -> g a
type NoExpoTreeAlg g = forall a. g a
{-# NOINLINE fold #-}
fold :: Functor g => ExpoTreeAlg g -> NoExpoTreeAlg g -> ExpoTree a -> g a
fold f z NoExpoTree = z
fold f z (ExpoTree x y) = f x (fold f z (fmap (fold f z) y))
The natural transformation here c ~> x is very interesting, and turns out to be quite necessary. Here's the build translation:
hbuild :: HFunctor f => (forall x. Alg f x -> (c ~> x)) -> (c ~> Mu f)
hbuild g = g In
newtype I :: (* -> *) where
I :: x -> I x
deriving (Show, Eq, Functor, Foldable, Traversable)
-- Needs to be a newtype, otherwise RULE firer gets bamboozled
newtype ExpoTreeBuilder c = ETP {runETP :: (forall x. Functor x
=> (forall a. Maybe a -> x (x a) -> x a)
-> (forall a. x a)
-> (forall a. c a -> x a)
)}
{-# NOINLINE build #-}
build :: ExpoTreeBuilder c -> forall a. c a -> ExpoTree a
build g = runETP g ExpoTree NoExpoTree
The newtype for the builder function is needed, because GHC 8.0 doesn't know how to fire the RULE without.
Now, the shortcut fusion rule:
{-# RULES "ExpoTree fold/build"
forall (g :: ExpoTreeBuilder c) c (f :: ExpoTreeAlg g) (n :: NoExpoTreeAlg g).
fold f n (build g c) = runETP g f n c #-}
Implementation of 'match' with 'build':
{-# INLINE match #-}
match :: Tree -> ExpoTree ()
match n = build (match_mk n) (I ())
where
match_mk :: Tree -> ExpoTreeBuilder I
match_mk Leaf = ETP $ \ f z (I c) -> f (Just c) z
match_mk (Node x y) = ETP $ \ f z c ->
-- NB: This fmap is bad for performance
f Nothing (fmap (const (runETP (match_mk y) f z c)) (runETP (match_mk x) f z c))
Implementation of 'appl' with 'fold' (we need to define a custom functor to define the return type.)
newtype PFunTree a = PFunTree { runPFunTree :: Tree -> Maybe a }
deriving (Functor)
{-# INLINE appl #-}
appl :: ExpoTree a -> PFunTree a
appl = fold appl_expoTree appl_noExpoTree
where
appl_expoTree :: ExpoTreeAlg PFunTree
appl_expoTree = \z f -> PFunTree $ \n ->
case n of Leaf -> z
Node n1 n2 -> runPFunTree f n1 >>= flip runPFunTree n2
appl_noExpoTree :: NoExpoTreeAlg PFunTree
appl_noExpoTree = PFunTree $ \n -> Nothing
Putting it all together:
applmatch :: Tree -> Tree -> Maybe ()
applmatch x = runPFunTree (appl (match x))
We can once again inspect the core with -ddump-simpl. Unfortunately, while we have successfully fused away the TrieMap data structure, we are left with suboptimal code due to the fmap in match. Eliminating this inefficiency is left to future work.

Answer 2

The paper appears to draw a parallel between ExpoNat a as a recursive Type and Tree as a recursive type constructor (Type -> Type).
newtype Fix f = Fix (f ( Fix f))
newtype HFix h a = HFix (h (HFix h) a)
Fix f represents the least fixed point of the endofunctor on the category of types and functions, f :: Type -> Type; HFix h represents the least fixed point of the endofunctor h on a category of functors and natural transformations, h :: (Type -> Type) -> (Type -> Type).
-- x ~ Fix (ExpoNatF a) ~ ExpoNat
data ExpoNatF a x = ExpoNatF (Maybe a) x | NoExpoNatF
fmap :: (x -> y) -> ExpoNatF a x -> ExpoNatF a y
fmap f (ExpoNatF u v) = ExpoNatF u (f v)
fmap _ NoExpoNatF = NoExpoNatF
-- f ~ HFix TreeMapH ~ TreeMap
data TreeMapH f a = TreeMapH (Maybe a) (f (f a)) | EmptyTreeMapH
hfmap :: (f ~> g) -> (TreeMapH f ~> TreeMapH g)
hfmap f (TreeMapH u v) = TreeMapH u ((fmap . fmap) f v)
hfmap _ EmptyTreeMapH = EmptyTreeMapH
-- (~>) is the type of natural transformations
type f ~> g = forall a. f a -> g a
Endofunctors give rise to algebras.
type Alg f a = f a -> a
type HAlg h f = h f ~> f
fold, or cata maps any algebra to a morphism (function|natural transformation).
cata :: Alg f a -> Fix f -> a
hcata :: HAlg h f -> (HFix h ~> h)
build constructs a value from its Church encoding.
type Church f = forall a. Alg f a -> a
type HChurch h = forall f. HAlg h f ~> f
build :: Church f -> Fix f
hbuild :: HChurch h -> HFix h a
-- The paper actually has a slightly different type for Church encodings, derived from the categorical view, but I'm pretty sure they're equivalent
build/fold fusion is summarized by one equation.
cata alg ( build f) = f alg
hcata alg (hbuild f) = f alg