Ast Transformation - De bruijn Annotation along with Validation using Data.Comp.Multi

Main.hs

{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE QuantifiedConstraints #-}

import Data.Comp.Multi.Derive
import Data.Comp.Multi.Algebra
import Data.Comp.Multi.Term
import Data.Comp.Multi.HFunctor
import Data.Comp.Multi.HFoldable
import Data.Comp.Multi.HTraversable
import Data.Comp.Multi.Show
import Data.Functor.Compose
import Data.Validation
import Control.Monad.Reader
import Unsafe.Coerce

data Expr
data Name

data NodeF r i where
    Const :: Int -> NodeF r Expr
    Add :: r x -> r y -> NodeF r Expr
    Ref :: String -> NodeF r Name
    Lam :: String -> r i -> NodeF r Expr

$(makeHFunctor ''NodeF)
$(makeHFoldable ''NodeF)
$(makeHTraversable ''NodeF)
$(makeShowHF ''NodeF)

-- boilerplate due to some dubious choices of compdata

data (a :&: f) (r :: * -> *) i = a i :&: f r i

instance (ShowHF f, forall i. Show (a i)) => ShowHF (a :&: f) where
    showHF (x :&: y) = K $ case show x of
                            "" -> unK (showHF y)
                            s -> "(" ++ s ++ " :&: " ++ unK (showHF y) ++ ")"

pattern x :&&: y  = Term (x :&: y)

instance HFunctor f => HFunctor (a :&: f) where
    hfmap f (x :&: y) = (x :&: hfmap f y)
instance HFoldable f => HFoldable (a :&: f) where
    hfoldMap f (_ :&: y) = hfoldMap f y
instance HTraversable f => HTraversable (a :&: f) where
    htraverse f (x :&: y) = (:&:) <$> pure x <*> htraverse f y
    hmapM f (x :&: y) = (:&:) <$> pure x <*> hmapM f y

-- boilerplate ends here

data Located i where
    L :: Int -> Located i

instance Show (Located i) where
    show (L i) = show i

data Ann i where
    Z :: Int -> Int -> Ann Name
    N :: Int -> Ann Expr

instance Show (Ann i) where
    show (Z l i) = show l ++ " :&: index=" ++ show i
    show (N i) = show i

-- useless in this case, may be useful when a lot of irrelevant constructors

data Proof r i where
    PName :: NodeF r Name -> Proof r Name
    PExpr :: NodeF r Expr -> Proof r Expr

proof :: NodeF r i -> Proof r i
proof x@(Ref _) = PName x
proof x = unsafeCoerce $ PExpr $ unsafeCoerce x

-- end of proof

type LNodeF = Located :&: NodeF

type ANodeF = Ann :&: NodeF

type Env = [String]

look :: String -> Env -> Maybe Int
look x e = case length (takeWhile (/= x) e) of
                y | y == length e -> Nothing
                y | otherwise -> Just y

type Errors = [String]

newtype Renamer f i = Renamer { getRenamer :: Compose (Reader Env) (Validation Errors) (f i)}

renameAlg :: Alg LNodeF (Renamer (Term (ANodeF)))
renameAlg (L l :&: (proof -> PExpr (Lam name expr))) = Renamer $ Compose $ do
                                                expr' <- local (name:) $ getCompose $ getRenamer expr
                                                return $ fmap (N l :&&:) $ Lam <$> pure name <*> expr'
renameAlg (L l :&: (proof -> PName (Ref name))) = Renamer $ Compose $ do
                           mi <- asks $ look name
                           case mi of
                              Nothing -> pure $ Failure ["undeclared " ++ name ++ " at position " ++ show l]
                              Just i -> pure $ Success $ (Z l i :&&: Ref name)
renameAlg (L l :&: (proof -> PExpr xs)) = Renamer $ fmap (N l :&&:) $ htraverse getRenamer xs -- Can use N!!!!!

rename :: Term LNodeF i -> Validation Errors (Term ANodeF i)
rename x = runReader (getCompose $ getRenamer $ cata renameAlg x) []

right :: Term LNodeF Expr
right = L 0 :&&: Lam "x" (L 1 :&&: Lam "y" (L 2 :&&: Lam "z" (L 3 :&&: Add (L 4 :&&: Ref "z") (L 5 :&&: Ref "x"))))

wrong :: Term LNodeF Expr
wrong = L 0 :&&: Lam "x" (L 1 :&&: Lam "y" (L 2 :&&: Lam "z" (L 3 :&&: Add (L 4 :&&: Ref "k") (L 5 :&&: Ref "w"))))

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

Output:

right, original
(0 :&: (Lam "x" (1 :&: (Lam "y" (2 :&: (Lam "z" (3 :&: (Add (4 :&: (Ref "z")) (5 :&: (Ref "x"))))))))))
right, renamed
Success (0 :&: (Lam "x" (1 :&: (Lam "y" (2 :&: (Lam "z" (3 :&: (Add (4 :&: index=0 :&: (Ref "z")) (5 :&: index=2 :&: (Ref "x"))))))))))
wrong, original
(0 :&: (Lam "x" (1 :&: (Lam "y" (2 :&: (Lam "z" (3 :&: (Add (4 :&: (Ref "k")) (5 :&: (Ref "w"))))))))))
wrong, renamed
Failure ["undeclared k at position 4","undeclared w at position 5"]

One can actually also parameterize the NodeF datatype by a tag and use trees-that-grows approach in order to eg remove the name from the Lam case and use the bruijn index directly instead of a name on Ref. This is more easily done after actually doing the annotation , basically just a pure annotated pre-renamer tree -> post-rename tree transformation. Not actually changing the original tree is more flexible though as it has more information than the post-renamed tree.