anf.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Anf where

import Data.IORef
import Control.Monad
import Data.Foldable (foldrM)

data BOp   = Add | Sub | Mul
    deriving Show
data Expr  = ELit Int | EBOp BOp Expr Expr | ECall String [Expr]
    deriving Show

-- ANF
type Var   = String
data Value = VVar Var | VInt Int
    deriving Show
data AExpr =  AHalt Value
    | ABOp Var BOp Value Value AExpr -- lhs, bop, v1, v2, remaining expr
    | ACall Var String [Value] AExpr -- lhs, name, args, remaining expr
    deriving Show

fresh :: IORef Int -> String -> IO String
fresh c prefix = do
    void $ modifyIORef' c (+ 1)
    c1 <- readIORef c
    pure $ prefix <> show c1

anf :: Expr -> IO AExpr
anf e = do
    c <- newIORef 0
    go e c (pure . AHalt)
  where
    go :: Expr -> IORef Int -> (Value -> IO AExpr) -> IO AExpr
    go (ELit i) _ k = k (VInt i)
    go (EBOp op el er) c k = go el c $ 
        \ael -> go er c $ 
        \aer -> do 
            v   <- fresh c "v"
            rem <- k $ VVar v
            pure $ ABOp v op ael aer rem
    go (ECall f es) c k = do
        let finalK vs = do
            v   <- fresh c "v"
            rem <- k $ VVar v
            pure $ ACall v f vs rem
 
        foldr (\e ctx vs -> go e c $ ctx . (vs ++) . pure) finalK es $ []

main :: IO ()
main = do
    r <- anf $ ECall "fn" [ELit 10, EBOp Add (ELit 20) (ELit 30)]
    print r