--
-- CBN interpreter with state to count number of computation steps already taken.
--
module InterpStateCBN where

data Term = Var String 
          | Con Int
          | Add Term Term
          | Lam String Term
          | App Term Term
          | Count
  deriving Show

data Value = Wrong
             | Num Int
             | Fun (M Value -> M Value)

instance Show Value where
    show Wrong = "<wrong>"
    show (Num  n) = show n
    show (Fun _) = "<function>"


newtype StateM a = StateM (Int -> (a, Int))
instance Show a => Show (StateM a) where
    show (StateM m) = let (v, n) = m 0 
             in "Value: "++(show v)++"\nSteps: "++(show n)++"\n"

unit v = StateM $ \s -> (v, s)
bad s = StateM $ \s -> (Wrong, s)
bind (StateM x) f = StateM $ \s0 ->
    let (v, s1) = x s0
        StateM next = f v
    in next s1

tick = StateM $ \s -> ((), s+1)
getSt = StateM $ \s -> (Num s, s)


type M a = StateM a

type Env = [(String, M Value)]

getVar :: Env -> String -> M Value
getVar ((x,v):xs) y = if x == y then v else getVar xs y
getVar [] y = bad $ "unbound variable "++y

add :: Value -> Value -> M Value
add (Num m) (Num n) = bind tick $ \_ -> unit $ Num $ m + n
add m n = bad $ "expected numbers, found "++(show m)++" and "++(show n)

apply :: Value -> M Value -> M Value
apply (Fun v1) v2 = bind tick $ \_ -> v1 v2
apply f _ = bad $ "expected function, found "++(show f)

interp :: Term -> Env -> M Value
interp (Var x) env = getVar env x
interp (Con n) _ = unit $ Num n
interp (Add e1 e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    add v1 v2
interp (Lam x e) env = unit $ Fun $ \v -> interp e ((x,v):env)
interp (App e1 e2) env = 
    bind (interp e1 env) $ \v1 ->
    apply v1 $ interp e2 env
interp Count _ = getSt

test e = interp e []

i = Lam "x" $ Var "x"
k = Lam "x" $ Lam "y" $ Var "x"
s = Lam "f" $ Lam "g" $ Lam "x" $ App (App (Var "f") (Var "x")) (App (Var "g") (Var "x"))
omega = App (Lam "x" $ App (Var "x") (Var "x")) (Lam "x" $ App (Var "x") (Var "x"))