module SemanticsState where

import Data.Map.Strict (Map, (!))
import Data.Map.Strict qualified as Map
import Prelude hiding ((>>=), pure, fmap)

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
          | AppByRef Expr Expr          -- AppByRef 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 model and important primitives
----------------------------------------------

data MutM a = MkMutM (Map Int Value -> Either ErrorType (Map Int Value, a))
data ErrorType = TypeError | VarNotFound
    deriving Show
unMutM (MkMutM f) = f

-- Give an answer.
pure :: a -> MutM a
pure a = MkMutM (\s -> Right (s, a))

-- Sequential composition that passes answer from 1st stage to 2nd.
(>>=) :: MutM a -> (a -> MutM b) -> MutM b
MkMutM f >>= k = MkMutM (\s -> case f s of
                                 Left msg -> Left msg
                                 Right (s', a) -> unMutM (k a) s')
infixl 1 >>=

-- Raise an error.
raise :: ErrorType -> MutM a
raise err = MkMutM (\_ -> Left err)

-- Apply a function to the answer.
fmap :: (a -> b) -> MutM a -> MutM b
fmap f ma = ma >>= \a -> pure (f a)

-- Sequential composition that combines answers by a function.
liftA2 :: (a -> b -> c) -> MutM a -> MutM b -> MutM c
liftA2 f ma mb = ma >>= \a -> mb >>= \b -> pure (f a b)

-- 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, use new,
-- load, store for individual addresses.

-- 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 Map.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 (Map.insert a val)


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

mainInterp :: Expr -> Either ErrorType Value
mainInterp expr = case unMutM (interp expr Map.empty) Map.empty of
  Left e -> Left e
  Right (_, v) -> Right v

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 VarNotFound

interp (Assign v e) env =
    case Map.lookup v env of
      Nothing -> raise VarNotFound
      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 TypeError

-- 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

-- Function application but pass by reference. The argument e is expected to be
-- the Var case.
interp (AppByRef f e) env =
    interp f env
    >>= \c -> case c of
      VClosure fEnv v body ->
        -- Require e to be Var
        case e of
          Var x ->
            -- Find address of x
            case Map.lookup x env of
              Nothing -> raise VarNotFound
              Just addr ->
                -- map v to address of x
                let bEnv = Map.insert v addr fEnv
                in interp body bEnv
          _ -> raise TypeError
      _ -> raise TypeError

-- 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 TypeError


-- 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)
              ])
    )

-- let x = 0
--  in (\y -> y:=2) x; x
-- but that function call is pass by reference. The function changes x.
passByRefDemo = mainInterp
    (Let [ ("x", Num 0) ]
         (Seq [ AppByRef (Lambda "y" (Assign "y" (Num 2))) (Var "x")
              , Var "x"
              ])
    )