MultiRec in combination with "Trees That Grow" approach

Bad.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Bad where

import Prelude
import Data.Kind ( Type )
import Generics.MultiRec.TH


-- Type family used like in the "trees that grow" paper.
-- Each phase of the compiler can add/remove information as needed.
type family Anno a


-- Expression type, containing annotations
data Expr phase
  = I (Anno phase) Int
  | Plus (Anno phase) (Expr phase) (Expr phase)
  | Mul (Anno phase) (Expr phase) (Expr phase)

-- Naive approach: use it like in multirec examples
data ASTF :: Type -> Type -> Type where
  Expr :: ASTF phase (Expr phase)
$(deriveAll ''ASTF)

{-
Fails with:

src/Bad.hs:1:1: error:
    Exception when trying to run compile-time code:
      unknown construct
-}

Good.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Good ( main ) where


import Prelude
import Data.Kind ( Type )
import Generics.MultiRec.TH
import Generics.MultiRec.Compos
import Generics.MultiRec.FoldAlgK


-- Type family used like in the "trees that grow" paper.
-- Each phase of the compiler can add/remove information as needed.
type family Anno (a :: Phase)

data Phase = Parsing | TypeChecking

type instance Anno 'Parsing = Span
type instance Anno 'TypeChecking = (Span, TypeInfo)

data Span = Span Int Int
  deriving (Eq, Show)

data TypeInfo = TInt
  deriving (Eq, Show)

-- GADT that establishes the link between value and type level.
data Tag (a :: Phase) where
  TagP :: Tag 'Parsing
  TagTC :: Tag 'TypeChecking

-- Important: the Ann data type now contains no type parameters
-- used in a type family! (This would break generated code from
-- multirec "deriveAll" helper.)
--
-- It hides the 'a' parameter internally (as an existential variable).
-- We can find out the actual value of the parameter by pattern matching
-- on the tag, making it possible to access the annotation data.
-- 
-- A con of this approach is that now it's possible to have annotations of
-- multiple phases in 1 AST, but this can be mitigated by writing helper 
-- functions that convert all annotations in the AST.
data Ann where
  Ann :: Tag a -> Anno a -> Ann


-- The Eq and Show instances are not needed for this example,
-- but can be useful during debugging..

instance Eq Ann where
  (Ann TagP ann) == (Ann TagP ann2) = ann == ann2
  (Ann TagTC ann) == (Ann TagTC ann2) = ann == ann2
  _ == _ = False

instance Show Ann where
  show (Ann TagP ann) = "(" ++ show ann ++ ")"
  show (Ann TagTC ann) = show ann


-- Actual AST / expression type, containing annotations
data Expr
  = I Ann Int
  | Plus Ann Expr Expr
  | Mul Ann Expr Expr
  deriving (Eq, Show)

-- Necessary multirec boilerplate for generating code.
data ASTF :: Type -> Type where
  Expr :: ASTF Expr
$(deriveAll ''ASTF)

-- Another expression type, without annotations
data Expr2 = I2 Int
           | Plus2 Expr2 Expr2
           | Mul2 Expr2 Expr2
           deriving (Eq, Show)


-- Some example transforms:

plus1 :: Expr -> Expr
plus1 = plus1' Expr
  where
    plus1' :: ASTF a -> a -> a
    plus1' Expr (I ann i) = I ann $ i + 1
    plus1' p x = compos plus1' p x

stripAnns :: Expr -> Expr2
stripAnns = fold algebra Expr
  where
    algebra :: Algebra ASTF Expr2
    algebra _ = const I2
              & const Plus2
              & const Mul2

-- A rather simple example, but this can contain arbritrarily
-- complex logic to update annotations
addTypeInfo :: Expr -> Expr
addTypeInfo = fold algebra Expr
  where
    algebra :: Algebra ASTF Expr
    algebra _ = convertI
              & convertPlus
              & convertMul
    addIntType :: Ann -> Ann  -- Necessary type annotation
    addIntType (Ann TagP spanInfo) = Ann TagTC (spanInfo, TInt)
    addIntType ann = error ("Unexpected tag: "<> show ann)
    convertI ann = I (addIntType ann)
    convertPlus ann = Plus (addIntType ann)
    convertMul ann = Mul (addIntType ann)

-- 1 + 2 * 3
ast :: Expr
ast =
  let mkSpan begin end = Ann TagP (Span begin end)
      one = I (mkSpan 0 1) 1
      two = I (mkSpan 4 5) 2
      three = I (mkSpan 8 9) 3
   in Mul (mkSpan 6 7) (Plus (mkSpan 2 3) one two) three

main :: IO ()
main = do
  print $ plus1 ast
  print $ stripAnns ast
  print $ addTypeInfo ast

Here are some related links to these snippets:

• multirec library
• multirec paper
• trees that grow paper