haskellabstract-syntax-treecomonadcatamorphismrecursion-schemes

How to work with AST with Cofree annotation?


I have this simple Expr AST and I can easily convert it to String.

import Prelude hiding (Foldable)
import qualified Prelude
import Data.Foldable as F
import Data.Functor.Foldable
import Data.Monoid
import Control.Comonad.Cofree

data ExprF r = Const Int
              | Add   r r
                deriving ( Show, Eq, Ord, Functor, Prelude.Foldable )

type Expr = Fix ExprF

testExpr = Fix $ Add (Fix (Const 1)) (Fix (Const 2))

convertToString :: Expr -> String
convertToString = cata $ \case
  e@(Const x) -> show x
  e@(Add x y) -> unwords [x, "+", y]

Now I want to add an additional data to it. So I am trying to use Cofree

type LineNumber = Int
type Expr2 = Cofree ExprF LineNumber

I can convert Expr to Expr2

addLineNumbers :: Expr -> Expr2
addLineNumbers = cata $ \case
  e@(Const _) -> 1 :< e
  e -> 2 :< e

But I cannot figure out how to convert Expr2 to String

convertToString2 :: Expr2 -> String
convertToString2 = cata $ \case
  e@(_ :< (Const x)) -> show x
  e@(_ :< (Add x y)) -> unwords [x, "+", y]

Also, is Cofree the best way to solve this annotation problem?


Solution

  • An alternative way of annotating your syntax tree is to compose the annotation into the base functor.

    -- constant functor
    newtype K c a = K c
        deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)
    
    -- functor product
    data (f :*: g) a = (:*:) { left :: f a, right :: g a }
        deriving (Eq, Ord, Show, Read, Functor, Foldable, Traversable)
    

    We're going to use the functor product to attach an annotation (inside a K) to each layer of the tree.

    type AnnExpr = Fix (K LineNumber :*: ExprF)
    

    If you can generate annotations while only inspecting a single layer of the tree (that is, your annotation-generating code can be expressed as a natural transformation) then you can use the following bit of machinery to modify the functor while keeping the fixpoint structure in place:

    hoistFix :: Functor f => (forall a. f a -> g a) -> Fix f -> Fix g
    hoistFix f = Fix . f . fmap (hoistFix f) . unFix
    

    This is of limited usefulness, though, as most interesting annotations such as type-checking require traversal of the syntax tree.

    You can reuse the code to tear down an Expr by simply ignoring the annotations. Given an algebra for ExprF...

    -- instructions for a stack machine
    data Inst = PUSH Int | ADD
    type Prog = [Inst]
    
    compile_ :: ExprF Prog -> Prog
    compile_ (Const x) = [PUSH x]
    compile_ (Add x y) = x ++ y ++ [ADD]
    

    ... you can use it to tear down either an Expr or an AnnExpr:

    compileE :: Expr -> Prog 
    compileE = cata compile_
    
    compileA :: AnnExpr -> Prog
    compileA = cata (compile_ . right)