{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module State where

import Control.Monad.State

type V = String

data AST = Const Int
         | Var V
         | Plus AST AST
         | Let V AST AST
           deriving Show

ex1 = Plus (Const 1) (Const 2)
ex2 = Let "a" (Const 2) (Plus (Var "a") (Var "a"))
ex2res = Let "v0" (Const 2) (Plus (Var "v0") (Var "v0"))
ex3 = Var "apa"
ex3res = Var "apa"
ex4 = Let "a" (Const 1) (Let "b" (Const 2) (Plus (Var "a") (Var "b")))

-- Return all constants in the expression
-- renames all local variables to fresh names

-- CF Monad

data CFState =
  CFState { consts  :: [Int] 
          , subst   :: [(V,V)]
          , newName :: Int
          }

initState = CFState [] [] 0

newtype CF a = CF { unCF :: StateT CFState IO a }
  deriving (Functor, Monad)

addConstant :: Int -> CF ()
addConstant i = CF $ do st <- get
                        put (st { consts = i : consts st})

renameVar :: V -> CF V
renameVar v = CF $ do st <- get
                      case lookup v (subst st) of
                        Nothing -> return v
                        Just v' -> return v'

freshName :: CF V
freshName = CF $ do st <- get
                    let n = newName st
                    put (st { newName = n + 1 })
                    return ("v" ++ show n)

substitute :: V -> V -> CF a -> CF a
substitute v v' (CF m) = CF $
  do st <- get
     let s = subst st
     put (st { subst = (v,v') : s})
     a <- m
     st' <- get
     put (st { subst = s })
     return a

-- Implementation

constFresh :: AST -> IO (AST,[Int])
constFresh ast = do (a,st) <- runStateT (unCF (constFresh' ast)) initState
                    return (a,consts st)

constFresh' :: AST -> CF AST
constFresh' (Const i) = do addConstant i
                           return (Const i)
constFresh' (Var v) = do v' <- renameVar v
                         return (Var v')
constFresh' (Plus ast1 ast2) = do ast1' <- constFresh' ast1
                                  ast2' <- constFresh' ast2
                                  return (Plus ast1' ast2')
constFresh' (Let v ast1 ast2) = do
  ast1' <- constFresh' ast1
  v'    <- freshName
  ast2' <- substitute v v' (constFresh' ast2)
  return (Let v' ast1' ast2')