A4.hs

module A4 where

{-
CSC324 — 2024F — Assignment 4

Tasks

1. State Monad
  (a) `get` and `put`
  (b) Instance of Functor
  (c) Instance of Applicative
2. Imperative Programming
  (a) Figure out a suitable `State` type
  (b) Define `lit` and `var`
  (c) Define `def`
  (d) Define `+=`, `-=`, and `*=`
  (e) Define `while`

Replace all "func = undefined" with a valid definition (or definitions) for "func".

IMPORTANT: You are NOT allowed to modify existing imports or add new imports.
-}
import Data.Maybe (Maybe (..), fromJust, fromMaybe, isJust)
import Prelude
  ( Applicative (..),
    Bool (..),
    Eq (..),
    Functor (..),
    Int,
    Monad (..),
    Ord (..),
    Show (..),
    String,
    fst,
    snd,
    last,
    lookup,
    length,
    undefined,
    ($),
    (*),
    (+),
    (++),
    (-),
    (<$>),
    (==),
  )

-- import Control.Monad.State.Lazy(State(..), get, put, evalState)

{-
-----------------------------------------------------------------------------
* Task 1: State Monad *
-----------------------------------------------------------------------------
-}

{-
Recall that during the lecture, we discussed the `WCounter` monad, which keeps track of a counter.

Here, we generalize the concept of a counter to a more general state.
We make the state type `s` a parameter of the `State` monad.
-}

newtype State s a = State {runState :: s -> (a, s)}

{-
Intuitively, `State s` is a wrapper over type `a` that carries a state of type `s`.

It represents a stateful computation that takes an initial state of type `s` and returns a _result_ of type `a` along with a final _state_ of type `s`. We can treat the state `s` as some state context that threads through the computation.

Given a `State` monad, we run it by providing an initial state and extracting the result and the final state. Sometimes, we only care about the result and ignore the final state. Let's define a helper function `evalState` that does this.
-}

evalState :: State s a -> s -> a
evalState (State f) s = fst (f s)

getNextState :: State s a -> s -> s
getNextState (State f) s = snd (f s)

{-

(a) `get` and `put`
Now, let's define two functions `get` and `put` that allow us to interact with the state.

- `get` returns the current state as the result.
- `put n` replaces the current state with new state `n`, and returns `()`.

Define the functions `get` and `put`. You can check the provided tests to better understand the expected behaviour.
-}

get :: State s s
get = State (\state -> (state, state))

put :: s -> State s ()
put newState = State (\state -> ((), newState))


{-
(b) Instance of Functor
Recall that a `Functor` is a type class on a type constructor `f` that allows you to apply a mapping to the value inside the type constructor. For `State s`, the mapping is applied to the result of the computation.

Implement `Functor` for `State s`.
-}

instance Functor (State s) where
  fmap f (State state_func) = State (\state ->
    let (value, curr_state) = state_func state
    in (f value, curr_state))

{-
(c) Instance of Applicative
It's intuitive that `State s` is a `Monad`. More interestingly, `State s` is also an instance of a weaker type class called `Applicative`.

Intuitively,
- `pure` (`return` in `Monad`): wraps a value in the type constructor.
- `<*>`: applies a mapping wrapped in the type constructor t to a value inside the type constructor. Notice how this is similar to `fmap`, but the mapping itself is also wrapped in the type constructor.

For `State s`,
- `pure` should wrap a value `a` in the `State` monad with the state unchanged.
- `<*>` should do something similar to `fmap`, but the state change of the mapping should be taken into account.

Also, an `Applicative` instance should satisfy the following laws:
- Identity: `pure id <*> v = v`
- Composition: `pure (.) <*> u <*> v <*> w = u <*> (v <*> w)`
- Homomorphism: `pure f <*> pure x = pure (f x)`
- Interchange: `u <*> pure y = pure ($ y) <*> u`

Implement `Applicative` for `State s`. Do NOT use the `Monad` instance to implement this, or you will receive 0 points.

Note: There are more than one correct implementation for this task. You can choose any implementation that satisfies the laws.
-}

instance Applicative (State s) where
  pure value = State (\state -> (value, state))
  wrapped_func <*> wrapped_value = State (\state ->
    let 
      (val_func, s1) = runState wrapped_func state 
      (val, s2) = runState wrapped_value s1
    in (val_func val, s2))

{-
-----------------------------------------------------------------------------
* Task 2: Imperative Programming *
-----------------------------------------------------------------------------
-}

instance Monad (State s) where
  return a = State (\s -> (a, s))

  State f >>= g = State (\s -> let (a, s') = f s in runState (g a) s')

{-
The essense of imperative programming is to have a sequence of statements that can modify the state of the program.
With the `State` monad and do notation, we can simulate imperative programming in Haskell.

Your goal is to define `def`, `var`, `lit`, `while`, `+=`, `-=` and `*=` functions that allow you to write an imperative style factorial function in Haskell.

You are free to define additional helper functions if you find them useful.
-}

factorial :: Int -> Int
factorial n = def $ do
  result <- var 1
  i <- var n
  while i (> 0) $ do
    result *= i
    i -= lit 1
  return result

{-
Don't worry if you have no clue how to implement these functions. Let's do it step by step.

First, notice that there are two different kinds of values that we need to represent in our imperative language: variables and literals.
We can define a `Value` type that represents either a variable or a literal of type `a`. A variable is represented by an integer index, and a literal is just a value of type `a`.
-}

data Value a = Var Int | Lit a
  deriving (Show, Eq)

{-
(a) Figure out a suitable `State` type
Recall how we defined the operational semantics of lambda calculus: we have an environment that maps variables to values. Similarly, here we need a state that maps variable indices to values.

Define a suitable State type `Ctx`. You can use either `data` or `type` to define `Ctx`.

- `getVar` that takes an index and returns the value of the variable at that index.
- `emptyCtx` that represents an initial state with no variables.

Hint: you need to track the next available variable index in the state, so that you can generate fresh variables when calling `var`.

For automatic testing, you should define a `emptyCtx` value that represents an initial state with no variables.
-}

-- Either an empty enviorment, or a part (mapping) from an index (int) to value (a) + recursive
data Ctx a = Empty | Context Int a (Ctx a) deriving Show

emptyCtx :: Ctx a
emptyCtx = Empty

getVar :: Ctx a -> Int -> a
getVar (Context searched_index literal rest) target_index 
  | (target_index == searched_index) = literal
  | True = getVar rest target_index
getVar Empty int = undefined

-- >>> a = Context 3 "1" (Context 2 "2" (Context 1 "3" Empty))
-- >>> a
-- Context 1 "1" (Context 2 "2" (Context 3 "3" Empty 1) 2) 3
-- >>> getVar a 2
-- "2"

{-
(b) Define `lit` and `var`
- `lit` creates a `Value` from a literal value.
- `var` creates a new variable in the state and returns a `Value` representing that variable.

Hint: don't forget to update the state to keep track of the new variable. Use do notation to make it easier to work with the `State` monad.
-}

lit :: a -> Value a
lit = Lit

-- >>> pure (lit 10)
-- Lit 10

largest_index :: Ctx a -> Int
largest_index Empty = 0
largest_index (Context index _ rest) = 
  let rest_largest = largest_index rest
  in if rest_largest > index then rest_largest else index

var :: a -> State (Ctx a) (Value a)
var newValue = State (\state -> 
  case state of
    Context old_index old_value old_rest -> (Var (old_index + 1), Context (largest_index state + 1) newValue state)
    Empty -> (Var 1, Context 1 newValue Empty)
  )

-- >>> s1 = var 111
-- >>> v1 = evalState s1 Empty
-- >>> v1
-- Var 1
-- >>> env2 = getNextState s1 Empty
-- >>> env2
-- Context 1 111 Empty
-- >>> s2 = var 222
-- >>> v2 = evalState s2 env2
-- >>> v2
-- Var 2
-- >>> env3 = getNextState s2 env2
-- >>> getVar env3 2
-- 222
-- >>> getVar env3 1
-- 111
--- >>> env3
-- Context 2 222 (Context 1 111 Empty)

{-
(c) Define `def`
`def` is a helper function that takes an imperative program, runs it with an initial state, and returns the final result.

Hint: You can use `undefined` as a value of any type you want, as long as you are sure that it will never be evaluated.
-}

def :: State (Ctx a) (Value a) -> a
def (State func) = 
  let (currVal, currEnv) = func Empty
  in case currVal of 
    Var index -> getVar currEnv index
    Lit n -> n

{-
(d) Define `+=`, `-=`, and `*=`
These functions should modify the value of a variable in the state, and return nothing. If the first argument is a literal, they do nothing.

For convenience, start with a helper function `modifyVar` that takes a mapping function `f` and applies it to a variable.

Note: for this task, you are free to remove `modifyVar` and implement `+=`, `-=`, and `*=` directly if you prefer.
-}

modifyVar :: (a -> a -> a) -> Value a -> Value a -> State (Ctx a) ()
modifyVar f (Var index_l) rhs = State (\state->
  let evalLhs = getVar state index_l
      evalRhs = case rhs of
        Lit n -> n
        Var index_r -> getVar state index_r
      new_val = f evalLhs evalRhs
      new_index = case state of 
        Empty -> index_l  -- make the new index index_l, so that it will be on top of the original index
        Context old_index _ _ -> index_l
  in ((), Context new_index new_val state)
  )
modifyVar _ (Lit n) _ = State (\state -> ((), state))

(+=) :: Value Int -> Value Int -> State (Ctx Int) ()
(+=) = modifyVar (+)

(-=) :: Value Int -> Value Int -> State (Ctx Int) ()
(-=) = modifyVar (-)

(*=) :: Value Int -> Value Int -> State (Ctx Int) ()
(*=) = modifyVar (*)

{-
(e) Define `while`
`while` takes a value, a predicate function, and a body program. It runs the body program repeatedly as long as the predicate function returns `True` on the value.

Hint:
Informally,
- while cond pred body = body; while cond pred body (if pred cond is True)
- while cond pred body = () (if pred cond is False)

Note: Completing this task does not guarantee any points. However, skipping this task will make the factorial function fail.
-}
while :: Value a -> (a -> Bool) -> State (Ctx a) () -> State (Ctx a) ()
while currVal predicate body = State (\state ->
  let evalCurrVal = 
        case currVal of
          Lit n -> n
          Var index -> getVar state index
      repeat_loop = predicate evalCurrVal
      new_env = getNextState (do body; while currVal predicate body) state
  in if repeat_loop
    then ((), new_env)
    else ((), state)
  )