Renaming ast using cata

recursion-schemes-annotation.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TemplateHaskell #-}

import Data.Functor.Base
import Data.Functor.Foldable
import Data.Functor.Foldable.TH
import Data.Validation
import Control.Monad.Reader
import Data.Functor.Compose
import Data.Fix (Fix(..))
import Text.Show.Deriving
import Data.List (intersperse)

-- think much more constructors, around 40 in practice
data ExprF r = Const Int
            | Add r r
            | Ref String
            | Let String r
    deriving (Functor, Foldable, Traversable, Eq, Show)

$(deriveShow1 ''ExprF)

data AExprF r = Int :<< ExprF r
    deriving (Functor, Foldable, Traversable, Eq, Show)

$(deriveShow1 ''AExprF)

pattern x :< y = Fix (x :<< y)

type AExpr = Fix AExprF

type Env = [String]

namespace :: Env -> String
namespace = concat . intersperse "." . reverse

look :: String -> Env -> Maybe String
look x e = case dropWhile (/= x) e of
                    [] -> Nothing
                    xs -> Just $ namespace xs

type Errors = [String]
type Renamer = Compose (Reader Env) (Validation Errors)

renameAlg :: AExprF (Renamer AExpr) -> Renamer AExpr
renameAlg (p :<< Let name expr) = Compose $ do
                                    ns <- asks namespace
                                    let name' = if ns == [] then name else ns ++ "." ++ name
                                    expr' <- local (name:) $ getCompose expr
                                    return $ fmap (p :<) $ Let <$> pure name' <*> expr'
renameAlg (p :<< Ref name) = Compose $ do
                           mname <- asks $ look name
                           case mname of
                              Nothing -> pure $ Failure ["undeclared name " ++ name ++ " at " ++ show p]
                              Just name -> pure $ Success $ (p :< Ref name)
renameAlg xs = fmap Fix $ sequenceA xs -- this is the useful part

rename :: AExpr -> Validation Errors AExpr
rename x = runReader (getCompose $ cata renameAlg x) []
                                    
wrong :: AExpr
wrong = (0 :< Add (1 :< Let "x" (2 :< Add (3 :< Ref "y") (4 :< Ref "x")))
                    (5 :< Add (6 :< Let "z" (7 :< Add (8 :< Ref "h") (9 :< Add (10 :< Const 0) (11 :< Ref "z"))))
                               (12 :< Const 1)))

right :: AExpr
right = (0 :< Add (1 :< Let "x" (2 :< Let "y" (3 :< Ref "y")))
                  (4 :< Let "z" (5 :< Add (6 :< Const 0) (7 :< Let "k" (8 :< Ref "k")))))

main = do
    putStrLn "wrong"
    print wrong
    putStrLn "renamed wrong"
    print $ rename wrong
    putStrLn "right"
    print right
    putStrLn "renamed right"
    print $ rename right