{-
  Name: Yuri Tatishchev
  Class: CS 252
  Assigment: HW2
  Date: 2026-03-06
  Description: Implements the big-step operational semantics for
    the WHILE language described in `while-semantics.pdf`
-}


module WhileInterp (
  Expression(..),
  Binop(..),
  Value(..),
  testProgram,
  run
) where

import Data.Map (Map)
import qualified Data.Map as Map

-- We represent variables as strings.
type Variable = String

-- The store is an associative map from variables to values.
-- (The store roughly corresponds with the heap in a language like Java).
type Store = Map Variable Value

data Expression =
    Var Variable                            -- x
  | Val Value                               -- v
  | Assign Variable Expression              -- x := e
  | Sequence Expression Expression          -- e1; e2
  | Op Binop Expression Expression
  | If Expression Expression Expression     -- if e1 then e2 else e3
  | While Expression Expression             -- while (e1) e2
  | BoolOp BoolBinop Expression Expression
  | Not Expression
  deriving (Show)

data Binop =
    Plus     -- +  :: Int  -> Int  -> Int
  | Minus    -- -  :: Int  -> Int  -> Int
  | Times    -- *  :: Int  -> Int  -> Int
  | Divide   -- /  :: Int  -> Int  -> Int
  | Gt       -- >  :: Int -> Int -> Bool
  | Ge       -- >= :: Int -> Int -> Bool
  | Lt       -- <  :: Int -> Int -> Bool
  | Le       -- <= :: Int -> Int -> Bool
  deriving (Show)

data BoolBinop = 
    And
  | Or
  deriving (Show)

data Value =
    IntVal Int
  | BoolVal Bool
  deriving (Show)


-- This function will be useful for defining binary operations.
-- The first case is done for you.
-- Be sure to explicitly check for a divide by 0 and throw an error.
applyOp :: Binop -> Value -> Value -> Value
applyOp Plus (IntVal i) (IntVal j) = IntVal $ i + j
applyOp Minus (IntVal i) (IntVal j) = IntVal $ i - j
applyOp Times (IntVal i) (IntVal j) = IntVal $ i * j
applyOp Divide (IntVal i) (IntVal j) = IntVal $ i `div` j
applyOp Gt (IntVal i) (IntVal j) = BoolVal $ i > j
applyOp Ge (IntVal i) (IntVal j) = BoolVal $ i >= j
applyOp Lt (IntVal i) (IntVal j) = BoolVal $ i < j
applyOp Le (IntVal i) (IntVal j) = BoolVal $ i <= j
applyOp _ _ _ = error "Not implemented for non-integer values"

applyBoolOp :: BoolBinop -> Value -> Value -> Value
applyBoolOp And (BoolVal b1) (BoolVal b2) = BoolVal $ b1 && b2
applyBoolOp Or (BoolVal b1) (BoolVal b2) = BoolVal $ b1 || b2
applyBoolOp _ _ _ = error "Not implemented for non-boolean values"

applyNot :: Value -> Value
applyNot (BoolVal b) = BoolVal $ not b
applyNot _ = error "Not implemented for non-boolean values"

-- Implement this function according to the specified semantics
evaluate :: Expression -> Store -> (Value, Store)
evaluate (Val v) s = (v, s)
evaluate (Var x) s = case Map.lookup x s of
  Just v -> (v, s)
  Nothing -> error "Variable not found"
evaluate (Assign x e) s =
  let (v,s') = evaluate e s
  in (v, Map.insert x v s')
evaluate (Sequence e1 e2) s =
  let (_,s1) = evaluate e1 s
      (v2,s') = evaluate e2 s1
  in (v2, s')
evaluate (Op o e1 e2) s =
  let (v1,s1) = evaluate e1 s
      (v2,s') = evaluate e2 s1
  in (applyOp o v1 v2, s')
evaluate (If e1 e2 e3) s =
  let (v1,s1) = evaluate e1 s
  in case v1 of
    BoolVal b -> if b then evaluate e2 s1 else evaluate e3 s1
    _ -> error "Not implemented for non-boolean values"
evaluate (While e1 e2) s =
  let (v1,s1) = evaluate e1 s
  in case v1 of
    BoolVal b -> if b then
                   let (_, s2) = evaluate e2 s1
                   in evaluate (While e1 e2) s2
                 else (BoolVal False, s1)
    _ -> error "Not implemented for non-boolean values"
evaluate (BoolOp o e1 e2) s =
  let (v1,s1) = evaluate e1 s
      (v2,s') = evaluate e2 s1
  in (applyBoolOp o v1 v2, s')
evaluate (Not e) s =
  let (v,s') = evaluate e s
  in (applyNot v, s')

-- Evaluates a program with an initially empty state
run :: Expression -> (Value, Store)
run prog = evaluate prog Map.empty

-- The same as run, but only returns the Store
testProgram :: Expression -> Store
testProgram prog = snd $ run prog