Search code examples
haskellrecursion-schemes

Using a paramorphism inside of an apomorphism


I'm trying to use paramorphisms and apomorhisms (in haskell):

-- Fixed point of a Functor                                                           
newtype Fix f = In (f (Fix f))                                                        
                                                                                      
deriving instance (Eq (f (Fix f))) => Eq (Fix f)                                      
deriving instance (Ord (f (Fix f))) => Ord (Fix f)                                    
deriving instance (Show (f (Fix f))) => Show (Fix f)                                  
                                                                                      
out :: Fix f -> f (Fix f)                                                             
out (In f) = f  

type RAlgebra f a = f (Fix f, a) -> a                                                 
                                                                                      
para :: (Functor f) => RAlgebra f a -> Fix f -> a                                     
para rAlg = rAlg . fmap fanout . out                                                  
        where fanout t = (t, para rAlg t)                                             
                                                                                      
-- Apomorphism                                                                        
type RCoalgebra f a = a -> f (Either (Fix f) a)                                       
                                                                                      
apo :: Functor f => RCoalgebra f a -> a -> Fix f                                      
apo rCoalg = In . fmap fanin . rCoalg                                                 
        where fanin = either id (apo rCoalg) 

to define the following recursive function:

fun concat3 (v,E,r) = add(r,v)
   | concat3 (v,l,E) = add(l,v)
   | concat3 (v, l as T(v1,n1,l1,r1), r as T(v2,n2,l2,r2)) =
       if weight*n1 < n2 then T’(v2,concat3(v,l,l2),r2)
       else if weight*n2 < n1 then T’(v1,l1,concat3(v,r1,r))
       else N(v,l,r)

It takes two binary trees and an element that is greater than the values in the left tree and less than the values in the right tree and combines them into one binary tree :: value -> tree1 -> tree2 -> tree3

I have defined the add function (which inserts an element into a binary tree) as a paramorphism like so:

add :: Ord a => a -> RAlgebra (ATreeF a) (ATreeF' a)
add elem EmptyATreeF = In (NodeATreeF elem 1 (In EmptyATreeF) (In EmptyATreeF))
add elem (NodeATreeF cur _ (prevLeft, left) (prevRight, right))
        | elem < cur = bATreeConstruct cur left prevRight
        | elem > cur = bATreeConstruct cur prevLeft right
        | otherwise = nATreeConstruct cur prevLeft prevRight

When I try to write concat3 as an apomorphism:

concat3 :: Ord a => a -> RCoalgebra (ATreeF a) (ATreeF' a, ATreeF' a)                 
concat3 elem (In EmptyATreeF, In (NodeATreeF cur2 size2 left2 right2)) = 
     out para (insertATreeFSetPAlg elem) (In (NodeATreeF cur2 size2 (Left left2) (Left right2)))
     ...

Because the next level of the apomorphism has not been evaluated yet, I get a type error from the compiler.

Couldn't match type: Fix (ATreeF a)                                             
                     with: Either (Fix (ATreeF a)) (ATreeF' a, ATreeF' a)             
      Expected: ATreeF a (Either (Fix (ATreeF a)) (ATreeF' a, ATreeF' a))             
        Actual: ATreeF a (Fix (ATreeF a)) 

Is there another approach I can take?


Solution

  • Some missing context to explain the solution is that this is from an implementation of weight-balanced trees, specifically Adams's variant (which happens to be the data structure behind Data.Set and Data.Map.)

    A problem when writing concat3 as a coalgebra is that it is not corecursive, strictly speaking, because the recursive calls of concat3 are under a smart constructor T', i.e., a function (which does some non-trivial rebalancing).

    A solution is to introduce an intermediate representation which delays the evaluation of that smart constructor.

    -- | Tree with delayed rebalancing operations T', or Id when no rebalancing is needed
    data TreeF1 a x = E1 | T' a x x | Id (Tree a)
      deriving Functor
    

    So we can write a coalgebra of TreeF1:

    concatAlg :: Ord a => a -> RCoalgebra (TreeF1 a) (Tree a, Tree a)
    concatAlg v (In E, r) = Id (add r v)
    concatAlg v (l, In E) = Id (add l v)
    concatAlg v (l@(In (T v1 n1 l1 r1)), r@(In (T v2 n2 l2 r2))) =
      if balance * n1 < n2 then T' v2 (Right (l, l2)) (Left (In (Id r2)))
      else if balance * n2 < n1 then T' v1 (Left (In (Id l1))) (Right (r1, r))
      else Id (_N v1 l r)
    
    {- Reference implementation for comparison:
    fun concat3 (v,E,r) = add(r,v)
      | concat3 (v,l,E) = add(l,v)
      | concat3 (v, l as T(v1,n1,l1,r1), r as T(v2,n2,l2,r2)) =
           if weight*n1 < n2 then T’(v2,concat3(v,l,l2),r2)
           else if weight*n2 < n1 then T’(v1,l1,concat3(v,r1,r))
           else N(v,l,r)
    -}
    

    And we can convert a Fix (TreeF1 a) to Fix (Tree a) via a catamorphism, finally executing those delayed applications of rebalancing T'.

    _T :: a -> Tree a -> Tree a -> Tree a
    _T = error "todo: rebalance"
    
    type Algebra f a = f a -> a
    
    -- do the rebalancing on T' v l r nodes
    rebalanceAlg :: Algebra (TreeF1 a) (Tree a)
    rebalanceAlg E1 = In E
    rebalanceAlg (T' v l r) = _T v l r
    rebalanceAlg (Id t) = t
    

    So concat3 is a composition of cata and apo using the above algebras:

    concat3 :: Ord a => a -> Tree a -> Tree a -> Tree a
    concat3 v l r = (cata rebalanceAlg . apo (concatAlg v)) (l, r)
    

    You can fuse cata and apo so that, after some elementary compiler optimizations, the intermediate tree does not get allocated:

    -- fusion of (cata _ . apo _)
    cataApo :: Functor f => Algebra f b -> RCoalgebra f a -> a -> b
    cataApo alg coalg = go
      where
        go x = alg (either (cata alg) go <$> coalg x)
    
    concat3' :: Ord a => a -> Tree a -> Tree a -> Tree a
    concat3' v l r = cataApo rebalanceAlg (concatAlg v) (l, r)
    

    Full gist: https://gist.github.com/Lysxia/281010fbe40eac9be0b135d4733c3d5a