{-# LANGUAGE ConstraintKinds, DeriveFunctor, GADTs, LambdaCase #-}
{-# LANGUAGE FunctionalDependencies, FlexibleContexts, FlexibleInstances #-}

module Processor where

import Control.Selective
import Control.Selective.Rigid.Free
import Data.Bool
import Data.Functor
import Data.Int (Int16)
import Data.Map.Strict (Map)
import Data.Word (Word8)
import Foreign.Marshal.Utils (fromBool)
import Prelude hiding (read, log)

import qualified Control.Monad.Trans.State as S
import qualified Data.Map.Strict as Map

-- See Section 5.3 of the paper: https://dl.acm.org/doi/10.1145/3341694.
-- Note that we have changed the naming.

-- | A standard @MonadState@ class extended with the 'Selective' interface.
class (Selective m, Monad m) => MonadState s m | m -> s where
    get   :: m s
    put   :: s -> m ()
    state :: (s -> (a, s)) -> m a

instance Monad m => MonadState s (S.StateT s m) where
    get   = S.get
    put   = S.put
    state = S.state

gets :: MonadState s m => (s -> a) -> m a
gets f = f <$> get

modify :: MonadState s m => (s -> s) -> m ()
modify f = state (\s -> ((), f s))

--------------------------------------------------------------------------------
-------- Types -----------------------------------------------------------------
--------------------------------------------------------------------------------

-- | All values are signed 16-bit words.
type Value = Int16

-- | The processor has four registers.
data Register = R0 | R1 | R2 | R3 deriving (Show, Eq, Ord)

-- | The register bank maps registers to values.
type RegisterBank = Map Register Value

-- | The address space is indexed by one byte.
type Address = Word8

-- | The memory maps addresses to signed 16-bit words.
type Memory = Map.Map Address Value

-- | The processor has two status flags.
data Flag = Zero     -- ^ tracks if the result of the last arithmetical operation was zero
          | Overflow -- ^ tracks integer overflow
    deriving (Show, Eq, Ord)

-- | A flag assignment.
type Flags = Map Flag Value

-- | Address in the program memory.
type InstructionAddress = Value

-- | A program execution log entry, recording either a read from a key and the
-- obtained value, or a write to a key, along with the written value.
data LogEntry k v where
    ReadEntry  :: k -> v -> LogEntry k v
    WriteEntry :: k -> v -> LogEntry k v

-- | A log is a sequence of log entries, in the execution order.
type Log k v = [LogEntry k v]

-- | The complete processor state.
data State = State { registers :: RegisterBank
                   , memory    :: Memory
                   , pc        :: InstructionAddress
                   , flags     :: Flags
                   , log       :: Log Key Value}

-- | Various elements of the processor state.
data Key = Reg Register | Cell Address | Flag Flag | PC deriving Eq

instance Show Key where
    show (Reg r)  = show r
    show (Cell a) = show a
    show (Flag f) = show f
    show PC       = "PC"

-- | The base functor for mutable processor state.
data RW a = Read  Key                 (Value -> a)
          | Write Key (Program Value) (Value -> a)
    deriving Functor

-- | A program is a free selective on top of the 'RW' base functor.
type Program a = Select RW a

instance Show (RW a) where
    show (Read  k          _) = "Read "  ++ show k
    show (Write k (Pure v) _) = "Write " ++ show k ++ " " ++ show v
    show (Write k _        _) = "Write " ++ show k

logEntry :: MonadState State m => LogEntry Key Value -> m ()
logEntry item = modify $ \s -> s { log = log s ++ [item] }

-- | Interpret the base functor in a 'MonadState'.
toState :: MonadState State m => RW a -> m a
toState = \case
    (Read k t) -> do
        v <- case k of Reg  r    -> gets ((Map.! r) . registers)
                       Cell addr -> gets ((Map.! addr) . memory)
                       Flag f    -> gets ((Map.! f) . flags)
                       PC        -> gets pc
        logEntry (ReadEntry k v)
        pure (t v)
    (Write k p t) -> do
        v <- runSelect toState p
        logEntry (WriteEntry k v)
        case k of
            Reg r     -> let regs' s = Map.insert r v (registers s)
                         in  state (\s -> (t v, s {registers = regs' s}))
            Cell addr -> let mem' s = Map.insert addr v (memory s)
                         in state (\s -> (t v, s {memory = mem' s}))
            Flag f    -> let flags' s = Map.insert f v (flags s)
                         in state (\s -> (t v, s {flags = flags' s}))
            PC        -> state (\s -> (t v, s {pc = v}))

-- | Interpret a program as a state transformer.
runProgramState :: Program a -> State -> (a, State)
runProgramState f = S.runState (runSelect toState f)

-- | Interpret the base functor in the selective functor 'Over'.
toOver :: RW a -> Over [RW ()] a
toOver (Read  k _   ) = Over [Read k (const ())]
toOver (Write k fv _) = runSelect toOver fv *> Over [Write k fv (const ())]

-- | Get all possible program effects.
getProgramEffects :: Program a -> [RW ()]
getProgramEffects = getOver . runSelect toOver

-- | A convenient alias for reading an element of the processor state.
read :: Key -> Program Value
read k = liftSelect (Read k id)

-- | A convenient alias for writing into an element of the processor state.
write :: Key -> Program Value -> Program Value
write k fv = liftSelect (Write k fv id)

--------------------------------------------------------------------------------
-------- Instructions ----------------------------------------------------------
--------------------------------------------------------------------------------

-- | The addition instruction, which reads the summands from a 'Register' and a
-- memory 'Address', adds them, writes the result back into the same register,
-- and also updates the state of the 'Zero' flag to indicate whether the
-- resulting 'Value' is zero.
add :: Register -> Address -> Program Value
add reg addr = let arg1   = read (Reg reg)
                   arg2   = read (Cell addr)
                   result = (+)   <$> arg1 <*> arg2
                   isZero = (==0) <$> write (Reg reg) result
               in write (Flag Zero) (bool 0 1 <$> isZero)

-- | A conditional branching instruction that performs a jump if the result of
-- the previous operation was zero.
jumpZero :: Value -> Program ()
jumpZero offset = let zeroSet  = (==1) <$> read (Flag Zero)
                      modifyPC = void $ write PC ((+offset) <$> read PC)
                  in whenS zeroSet modifyPC

-- | A simple block of instructions.
addAndJump :: Program ()
addAndJump = add R0 1 *> jumpZero 42

-----------------------------------
-- Add with overflow tracking -----
-----------------------------------

{-  The following example demonstrates how important it is to be aware of your
    effects.

    Problem: implement the semantics of the @add@ instruction which calculates
    the sum of two values and writes it to the specified destination, updates
    the 'Zero' flag if the result is zero, and also detects if integer overflow
    has occurred, updating the 'Overflow' flag accordingly.

    We'll take a look at two approaches that implement this semantics and see
    the crucial deference between them.
-}

-- | Add two values and detect integer overflow.
--
-- The interesting bit here is the call to the 'willOverflowPure' function.
-- Since the function is pure, the call @willOverflowPure <$> arg1 <*> arg2@
-- triggers only one 'read' of @arg1@ and @arg2@, even though we use the
-- variables many times in the 'willOverflowPure' implementation. Thus,
-- 'willOverflowPure' might be thought as an atomic processor microcommand.
addOverflow :: Key -> Key -> Key -> Program Value
addOverflow x y z =
    let arg1     = read x
        arg2     = read y
        result   = (+)   <$> arg1   <*> arg2
        isZero   = (==0) <$> write z result
        overflow = willOverflowPure <$> arg1 <*> arg2
    in write (Flag Zero)     (fromBool <$> isZero) *>
       write (Flag Overflow) (fromBool <$> overflow)

-- | A pure check for integer overflow during addition.
willOverflowPure :: Value -> Value -> Bool
willOverflowPure x y =
    let o1 = (>) y 0
        o2 = (>) x((-) maxBound y)
        o3 = (<) y 0
        o4 = (<) x((-) minBound y)
    in  (||) ((&&) o1 o2)
             ((&&) o3 o4)

-- | Add two values and detect integer overflow.
--
-- In this implementations we take a different approach and, when implementing
-- overflow detection, lift all the pure operations into 'Applicative'. This
-- forces the semantics to read the input variables more times than
-- 'addOverflow' does (@x@ is read 3x times, and @y@ is read 5x times).
addOverflowNaive :: Key -> Key -> Key -> Program Value
addOverflowNaive x y z =
    let arg1   = read x
        arg2   = read y
        result = (+)   <$> arg1 <*> arg2
        isZero = (==0) <$> write z result
        overflow = willOverflow arg1 arg2
    in write (Flag Zero)     (fromBool <$> isZero) *>
       write (Flag Overflow) (fromBool <$> overflow)

-- | An 'Applicative' check for integer overflow during addition.
willOverflow :: Program Value -> Program Value -> Program Bool
willOverflow arg1 arg2 =
    let o1 = (>) <$> arg2 <*> pure 0
        o2 = (>) <$> arg1 <*> ((-) maxBound <$> arg2)
        o3 = (<) <$> arg2 <*> pure 0
        o4 = (<) <$> arg1 <*> ((-) minBound <$> arg2)
    in  (||) <$> ((&&) <$> o1 <*> o2)
             <*> ((&&) <$> o3 <*> o4)

-----------------------------------
-- Example simulations ------------
-----------------------------------

renderState :: State -> String
renderState state =
  "Registers: " ++ show (registers state)          ++ "\n" ++
  "Flags: "     ++ show (Map.toList $ flags state) ++ "\n" ++
  "Log: "       ++ show (log state)

instance Show State where
    show = renderState

emptyRegisters :: RegisterBank
emptyRegisters = Map.fromList [(R0, 0), (R1, 0), (R2, 0), (R3, 0)]

emptyFlags :: Flags
emptyFlags = Map.fromList $ zip [Zero, Overflow] [0, 0..]

initialiseMemory :: [(Address, Value)] -> Memory
initialiseMemory m =
    let blankMemory = Map.fromList $ zip [0..maxBound] [0, 0..]
    in foldr (\(addr, value) acc -> Map.adjust (const value) addr acc) blankMemory m

boot :: Memory -> State
boot mem = State { registers = emptyRegisters
                 , pc = 0
                 , flags = emptyFlags
                 , memory = mem
                 , log   = [] }

twoAdds :: Program Value
twoAdds = add R0 0 *> add R0 0

addExample :: IO ()
addExample = do
    let initState = boot (initialiseMemory [(0, 2)])
    print . snd $ runProgramState twoAdds initState

---------------------------- Some boilerplate code -----------------------------

instance (Show k, Show v) => Show (LogEntry  k v) where
    show (ReadEntry k v)  = "Read (" ++ show k ++ ", " ++ show v ++ ")"
    show (WriteEntry k v) = "Write (" ++ show k ++ ", " ++ show v ++ ")"