module SemanticsState where

import           Control.Applicative
import           Data.IntMap.Strict (IntMap, (!))
import qualified Data.IntMap.Strict as IntMap
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

data Expr = Num Integer
          | Var String
          | Prim2 Op2 Expr Expr         -- Prim2 op operand operand
          | Let [(String, Expr)] Expr   -- Let [(name, rhs), ...] eval-me
          | Lambda String Expr          -- Lambda var body
          | App Expr Expr               -- App func arg
          | Assign String Expr          -- var := expr
          | Seq [Expr]                  -- sequential execution
    deriving (Eq, Show)

data Op2 = Plus | Minus | Mul  -- Just enough for demos.
    deriving (Eq, Show)

-- The type of possible answers from the interpreter.
data Value = VN Integer
           | VClosure (Map String Int) String Expr
           | VNone
    deriving (Eq, Show)


----------------------------------------------
-- The semantic monad and important primitives
----------------------------------------------

data MutM a = MkMutM (IntMap Value -> Either String (IntMap Value, a))

unMutM (MkMutM f) = f

instance Monad MutM where
    return a = MkMutM (\s -> Right (s, a))
    MkMutM f >>= k = MkMutM (\s -> case f s of
                                     Left msg -> Left msg
                                     Right (s', a) -> unMutM (k a) s')
instance Applicative MutM where
    pure = return
    mf <*> ma = mf >>= \f -> ma >>= \a -> return (f a)
instance Functor MutM where
    fmap f ma = ma >>= \a -> return (f a)

-- Return the whole memory. Used by new and load below.
getMem = MkMutM (\s -> Right (s, s))

-- Modify the memory by a function, i.e., change s to f s. Used by store below.
modMem f = MkMutM (\s -> Right (f s, ()))

-- The interpreter shouldn't directly use getMem and modMem. Instead, it just
-- needs to know how to new (allocate and init), load from an address, store to
-- an address.

-- Allocate and store initial value, return new address.  What's the new
-- address?  I simply use the largest used address plus one.
new :: Value -> MutM Int
new val =
    fmap IntMap.size getMem
    >>= \n -> store n val
    >> pure n

-- Read value at address.
load :: Int -> MutM Value
load a = fmap (! a) getMem

-- Write value to address.
store :: Int -> Value -> MutM ()
store a val = modMem (IntMap.insert a val)

-- And how to flag errors. Because no longer simply "Left".
raise :: String -> MutM a
raise msg = MkMutM (\_ -> Left msg)


------------------
-- The interpreter
------------------

mainInterp :: Expr -> Either String Value
mainInterp expr = fmap snd (unMutM (interp expr Map.empty) IntMap.empty)
-- This fmap is Either's fmap.

interp :: Expr -> Map String Int -> MutM Value

-- I first present those affected by the Env/Mem split.

interp (Var v) env =
    case Map.lookup v env of
      Just a -> load a
      Nothing -> raise "var not found"

interp (Assign v e) env =
    case Map.lookup v env of
      Nothing -> raise "var not found"
      Just addr ->
          interp e env
          >>= \val -> store addr val
          >> pure VNone

interp (Let eqns evalMe) env =
    extend env eqns
    >>= \env' -> interp evalMe env'
  where
    extend env [] = pure env
    extend env ((v,rhs) : eqns) =
        interp rhs env
        -- Now we need to allocate memory to store val
        >>= \val -> new val
        -- And the environment stores the address, not the value!
        >>= \addr -> let env' = Map.insert v addr env
        in extend env' eqns

interp (App f e) env =
    interp f env
    >>= \c -> case c of
      VClosure fEnv v body ->
          interp e env
          -- Now we need to allocate memory to store val.
          >>= \eVal -> new eVal
          -- And the environment stores the address, not the value!
          >>= \addr -> let bEnv = Map.insert v addr fEnv
          in interp body bEnv
      _ -> raise "type error: not lambda"

-- Seq is new but easy once you know the intention.

interp (Seq es) env = go es
  where
    -- The answer of the whole Seq is VNone if the list is empty, else the
    -- answer from the last expression.  But still run all expressions for the
    -- effects.
    go [] = pure VNone
    go [e] = interp e env
    go (e:es) = interp e env *> go es

-- The rest are like the functional language before.

interp (Num i) _ = pure (VN i)

interp (Lambda v b) env = pure (VClosure env v b)

interp (Prim2 op left right) env =
    fmap VN (liftA2 (discern op) (evalForInt left) (evalForInt right))
  where
    discern Plus = (+)
    discern Minus = (-)
    discern Mul = (*)
    evalForInt expr =
        interp expr env
        >>= \val -> case val of
          VN i -> pure i
          _ -> raise "type error: not integer"


-- let x=42; y=23
--  in x:=x+y; y:=x-y; x:=x-y;
--     x*10000 + y
swapDemo = mainInterp
    (Let [("x", Num 42), ("y", Num 23)]
         (Seq [ Assign "x" (Prim2 Plus (Var "x") (Var "y"))
              , Assign "y" (Prim2 Minus (Var "x") (Var "y"))
              , Assign "x" (Prim2 Minus (Var "x") (Var "y"))
              , Prim2 Plus (Prim2 Mul (Var "x") (Num 10000)) (Var "y")
              ]))


-- let f = (let x=0 in \y -> x:=x+1; x)
--  in f 0; f 0; f 0
counterDemo = mainInterp
    (Let [ ("f", Let [("x", Num 0)]
                     (Lambda "y" (Seq [ Assign "x" (Prim2 Plus (Var "x") (Num 1))
                                      , Var "x"
                                      ])))
         ]
         (Seq [ App (Var "f") (Num 0)
              , App (Var "f") (Num 0)
              , App (Var "f") (Num 0)
              ])
    )