module Hw4 where

data Term = 

    -- Lambda Calculus
    Var String         -- Var "_" is always unbound
  | String :\ Term     -- ("_" :\ e) is a thunk which throws away it's argument
  | Term :$: Term

    -- Integers
  | Con Integer
  | Term :+: Term
  | Term :-: Term

    -- Booleans
  | TT
  | FF
  | Not Term
  | Term :==: Term
  | If Term Term Term

    -- Pairs
  | Term :*: Term
  | Fst Term
  | Snd Term

    -- Lists
  | Term ::: Term
  | Nil 
  | LCase Term Term Term

    -- No-operation, has value Unit
  | Skip

    -- Sequencing
  | Term :& Term

    -- # of computation steps
  | Count

    -- Output
  | Out Term

    -- Nondeterminism
  | Amb Term Term

    -- Mutable references
  | MkRef Term
  | Read Term
  | Term := Term
  | ReadV String      -- Syntactic sugar, ReadV x  ==>  Read (Var x)
  | String :!= Term   -- Syntactic sugar, x :!= e  ==>  (Var x) := e

    -- Syntactic Sugar
  | Let [Decl] Term
  | LetRec [Decl] Term
  | While Term Term

  deriving Show

--
-- These declarations are for use in Let and LetRec terms.
--
data Decl = String :=: Term 
          | String :!=: Term   -- Syntactic sugar, x :!=: e  ==>  x :=: MkRef e
  deriving Show

--
-- fixity declarations for the infix Term and Decl constructors.
--
infix 1 :=:
infix 1 :!=:
infixr 2 :\
infixr 3 :&
infix 4 :!=
infix 4 :=
infixl 5 :$:
infixl 7 :+:
infixr 8 :::

--
-- Examples of well-formed terms can be found in the last part of this file.
--

data Value = Wrong
             | Fun (Value -> M Value)
             | Num Integer
             | Unit
             | List [Value]
             | Bl Bool
             | Ref Int
             | Pair (Value, Value)

instance Eq Value where
    Num x == Num y = x == y
    Unit == Unit = True
    List x == List y = x == y
    Bl x == Bl y = x == y
    Ref x == Ref y = x == y
    Pair x == Pair y = x == y
    _ == _ = False

instance Show Value where
    show Wrong = "<wrong>"
    show (Num  n) = show n
    show (Fun _) = "<function>"
    show Unit = "()"
    show (List x) = show x
    show (Bl x) = show x 
    show (Ref x) = "<ref "++(show x)++">"
    show (Pair x) = show x


data Err a = Err String | Succ a
instance Show a => Show (Err a) where
    show (Err s) = "Error: "++s
    show (Succ v) = "Success: "++(show v)


--
-- step is the number of computation steps.
-- refs is the values of the mutable reference cells (i.e. pointers). 
--   Cell names are Ints (indexing into refs), so (Ref 0)
--   is a reference to the first element of the refs field in MyState.
--
data MyState = MyState {step::Integer, refs::[Value]}
  deriving Show
initSt = MyState{step = 0, refs = []}


newtype All a = All (MyState -> [((String, Err a), MyState)])

instance Show a => Show (All a) where
    show (All x) = foldr go "" $ x initSt where
        go ((o,v),s) res = ("("++(show v)++" | Output: "++o++" | "++(show s)++")\n") ++ res
   

type M a = All a


type Env = [(String, Value)]


-- !!!!!!!!!!
-- Main interpreter combinators to be filled in.
--
unit :: a -> M a
unit v = 

bad :: String -> M a
bad err = 

bind :: M a -> (a -> M b) -> M b
bind (All x) f = 
--
-- !!!!!!!!!!

-- !!!!!!!!!!
-- Basic helper functions to be filled in
--
tick :: M ()   -- increments the number of computation steps.
tick = 

getStep :: M Value   -- retrieve current number of computation steps.
getStep = 

out :: Value -> M Value   -- send a value to the output stream.
out v = 

both :: M a -> M a -> M a   -- merge two results into one.
both (All x) (All y) = 
--
-- !!!!!!!!!!


pair proj (Pair v) = bind tick $ \_ -> unit $ proj v
pair _ v = bad $ "expecting pair, found "++(show v)

not' (Bl v) = unit $ Bl $ not v
not' x =  bad $ "expected bool, found "++(show x)

if' (Bl v) x y = bind tick $ \_ -> if v then x else y
if' v _ _ = bad $ "expected bool, found "++(show v)

cons x (List y) = bind tick $ \_ -> unit $ List $ x:y
cons x y = bad $ "expected list, found "++(show y)

lcase (List []) x y = bind tick $ \_ -> x
lcase (List (h:t)) x y = 
    bind tick $ \_ -> 
    bind y $ \v1 -> case v1 of 
        Fun f1 -> bind (f1 h) $ \v3 -> case v3 of
            Fun f2 -> f2 (List t)
            f2 -> bad $ "expecting function, found "++(show f2)
        f1 -> bad $ "expecting function, found "++(show f1)
lcase x _ _ = bad $ "expecting list, found "++(show x)


-- !!!!!!!!!!
-- three helper functions for interpreter refs to be filled in
--
getRef :: Value -> M Value   -- return the value of the given reference cell.
getRef (Ref x) = 
getRef x = bad $ "expecting ref, found "++(show x)

setRef :: Value -> Value -> M Value   -- set the given reference cell to the given value.
setRef (Ref x) v = 
setRef x _ = bad $ "expecting ref, found "++(show x)

newRef :: Value -> M Value   -- create, and return, a new reference cell.
newRef v = 
--
-- !!!!!!!!!!


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

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

apply :: Value -> 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 ("_" :\ e) env = unit $ Fun $ \_ -> interp e env
interp (x :\ e) env = unit $ Fun $ \v -> interp e ((x,v):env)
interp (e1 :$: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    apply v1 v2
interp (Con n) _ = unit $ Num n
interp (e1 :+: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    arith (+) v1 v2
interp (e1 :-: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    arith (-) v1 v2
interp Count _ = getStep
interp (Out e) env = bind (interp e env) out
interp (Amb e1 e2) env = both (interp e1 env) (interp e2 env)

interp (e1 :*: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    unit $ Pair (v1,v2)
interp (Fst e) env =
    bind (interp e env) (pair fst)
interp (Snd e) env =
    bind (interp e env) (pair snd)

interp TT _ = unit $ Bl True
interp FF _ = unit $ Bl False
interp (Not e) env = bind (interp e env) not'
interp (e1 :==: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    unit $ Bl $ v1 == v2
interp (If e1 e2 e3) env =
    bind (interp e1 env) $ \v1 -> 
    if' v1 (interp e2 env) (interp e3 env)

interp Skip _ = unit Unit
interp (e1 :& e2) env = 
    bind (interp e1 env) $ \_ -> interp e2 env 

interp Nil _  = unit $ List []
interp (e1 ::: e2) env = 
    bind (interp e1 env) $ \v1 ->
    bind (interp e2 env) $ \v2 ->
    cons v1 v2
interp (LCase e1 e2 e3) env =
    bind (interp e1 env) $ \v1 -> 
    lcase v1 (interp e2 env) (interp e3 env)

interp (MkRef e) env = bind (interp e env) newRef
interp (Read e) env = bind (interp e env) getRef
interp (ReadV x) env = bind (interp (Var x) env) getRef
interp (e1 := e2) env =
    bind (interp e1 env) $ \v1 -> 
    bind (interp e2 env) $ \v2 -> 
    setRef v1 v2
interp (x :!= e2) env =
    bind (interp (Var x) env) $ \v1 -> 
    bind (interp e2 env) $ \v2 -> 
    setRef v1 v2

-- !!!!!!!!!!
-- Three pieces of syntactic sugar to fill in
--
interp (Let decls e) env = 

interp (LetRec decls e) env = 

interp (While e1 e2) env = 
--
-- !!!!!!!!!!


test e = interp e []


------------------------------------------------------------------------
--
-- Example terms
--

--
-- identity function
--
i = "x" :\ Var "x"

--
-- k, or constant, combinator
--
k = "x" :\ "y" :\ Var "x"

--
-- s combinator
--
s = "f" :\ "g" :\ "x" :\ Var "f" :$: Var "x" :$: (Var "g" :$: Var "x")

--
-- omega should not terminate
--
omega = ("x" :\ Var "x" :$: Var "x") :$: ("x" :\ Var "x" :$: Var "x")


--
-- Pure Functional Fibonacci function
--
fib = LetRec ["f" :=: ("n" :\ If (Var "n" :==: Con 0) 
                                 (Con 0) 
                                 (If (Var "n" :==: Con 1) 
                                     (Con 1) 
                                     ((Var "f" :$: (Var "n" :-: Con 1)) :+: (Var "f" :$: (Var "n" :-: Con 2))) 
                                 )
                      )
             ]
             (Var "f")


--
-- Imperative Fibonacci function (both inputs are references, second is set to answer).
--
fibImp = "x" :\ "f" :\ 
    If (ReadV "x" :==: Con 0)
       ("f" :!= (Con 0))
       (Let ["k" :!=: Con 1, "g" :!=: Con 0] $ 
            "f" :!= Con 1 :& 
            While (Not $ ReadV "k" :==: ReadV "x")
                  (Let ["t" :=: ReadV "g"] $ 
                        "g" :!= ReadV "f" :&  
                        "f" :!= ReadV "f" :+: Var "t" :&
                        "k" :!= ReadV "k" :+: Con 1
                  ) :&
            Out (ReadV "f")
       )


--
-- Search function for lists
--
search = 
  LetRec ["srch" :=: "l" :\ "p" :\ "f" :\ "g" :\
             LCase (Var "l") 
                   (Var "g" :$: Skip) 
                   ("h" :\ "t" :\ 
                        If (Var "p" :$: Var "h") 
                           (Var "f" :$: Var "h") 
                           (Var "srch" :$: Var "t" :$: Var "p" :$: Var "f" :$: Var "g")
                   )
         ]
         (Var "srch")
          

--
-- Using state to memoize a function
--
memo = "f" :\ Let ["r" :!=: Nil] ("n" :\ 
           search :$: (ReadV "r") :$: 
                      ("x" :\ Fst (Var "x") :==: Var "n") :$: 
                      ("x" :\ Snd (Var "x")) :$:
                      ("_" :\ Let ["y" :=: Var "f" :$: Var "n"] 
                                  ("r" :!= (Var "n" :*: Var "y") ::: ReadV "r" :& Var "y")
                      )
           )


-- !!!!!!!!!!
-- Combinator to make memoized recursive functions, to be filled in.
--
memorec :: Term
memorec = 
--
-- !!!!!!!!!!


--
-- Recursively memoized Fibonacci function
--
fibm = memorec :$: ("f" :\ "n" :\ 
    If (Var "n" :==: Con 0) 
       (Con 0) 
       (If (Var "n" :==: Con 1) 
           (Con 1) 
           ((Var "f" :$: (Var "n" :-: Con 1)) :+: (Var "f" :$: (Var "n" :-: Con 2))) 
       )
   )


--
-- plus function which updates a counter before doing addition.
--
plus = "v" :\ "x" :\ "y" :\ "v" :!= ReadV "v" :+: Con 1 :& Var "x" :+: Var "y"


--
-- Pure Fibonacci function modified to record number of additions performed.
--
fib' = "x" :\ Let ["c" :!=: Con 0,] (
      LetRec ["f" :=: ("n" :\ If (Var "n" :==: Con 0) 
                                 (Con 0) 
                                 (If (Var "n" :==: Con 1) 
                                     (Con 1) 
                                     (plus :$: Var "c" :$: (Var "f" :$: (Var "n" :-: Con 1)) :$: (Var "f" :$: (Var "n" :-: Con 2)))
                                 )
                      )
             ]
             ((Var "f" :$: Var "x") :*: ReadV "c")
      )
 

--
-- Recursively memoized Fibonacci function modified to record the number of additions performed.
--
fibm' = "x" :\ Let ["c" :!=: Con 0] (
  (memorec :$: ("f" :\ "n" :\ 
                  If (Var "n" :==: Con 0) (Con 0) 
                         (If (Var "n" :==: Con 1) 
                             (Con 1) 
                             (plus :$: Var "c" :$: (Var "f" :$: (Var "n" :-: Con 1)) :$: (Var "f" :$: (Var "n" :-: Con 2))) 
                         )
               ) :$: 
               (Var "x")
   ) :*: ReadV "c"
 )