Annotate de bruijn using compdata

Main.hs

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

import Data.Comp.Multi.Derive
import Data.Comp.Multi.Algebra
import Data.Comp.Multi.Term
import Data.Comp.Multi.HTraversable
import Data.Functor.Compose
import Data.Validation
import Control.Monad.Reader

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)

data (a :&: f) (g ::  * -> *) e = a e :&&: f g e

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

data Ann i where
    Index :: Int -> Ann Name
    None :: Ann Expr

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@(Const _) = PExpr x
proof x@(Add _ _) = PExpr x
proof x@(Lam _ _) = PExpr x
proof x@(Ref _) = PName x

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 NodeF (Renamer (Term ANodeF))
renameAlg (proof -> PName (Ref name)) = Renamer $ Compose $ do
                           mi <- asks $ look name
                           case mi of
                               Nothing -> pure $ Failure ["undeclared name " ++ name]
                               Just i -> pure $ Success $ (Index i :&: Ref name)
renameAlg (proof -> PExpr (Lam name expr)) = Renamer $ Compose $ do
                           expr' <- local (name:) $ getCompose $ getRenamer expr
                           return $ fmap (None :&:) $ Lam <$> pure name <*> expr'
renameAlg (proof -> PExpr xs) = Renamer $ fmap (None :&:) $ htraverse getRenamer xs -- Can use None!

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

main = pure ()