module Machine where

import Data.Bits
import Data.Array.Base
import Data.Array.IO
import Data.Array.MArray

import qualified Data.IntMap as IM

import Control.Monad.State
import Control.Monad

import System.IO

import GHC.Base as B

import InstructionSet

-- the 8 general purpose registers of the machine
-- capable of holding one platter each
type MachRegisters = IOUArray Reg UInt32

-- an array of platters. A platter is a UInt32.
-- Note: array 0 stores the program
type MachPlatterArray = IOUArray UInt32 UInt32

-- Note: the program array ('0') is not put into the platter array table
-- The MachPlatterArray is a collection of platter arrays, keyed by a 32-bit indent
data MachState = MachState {
      mach_registers              :: !MachRegisters,

      -- a collection of platter arrays, keyed by a 32-bit ident
      mach_plattersTable          :: !(IM.IntMap MachPlatterArray),

      mach_program                :: !MachPlatterArray,
      mach_pc                     :: !UInt32,
      mach_nextFreePlatterArrayId :: !UInt32
    }

dumpState :: MachState -> IO String
dumpState s = do
  registers <- getElems (mach_registers s)
  return $ unlines $
    dumpRegisters registers ++
    dumpPC ++ dumpNextFreePlatterArrayId
  where dumpRegisters registers = map (\ (i,v) -> "reg[" ++ show i ++ "] = "
                                                  ++ show v)
                                      (zip [(0::UInt32) ..] registers)
        dumpPC = ["PC = " ++ show (mach_pc s)]
        dumpNextFreePlatterArrayId = ["Next Platter Array Id = " ++
                                      show (mach_nextFreePlatterArrayId s)]

-- the machine monad
type MA = StateT MachState IO

-- getReg r returns the value of register r
getReg :: Reg -> MA UInt32
getReg r = r `seq` do
  regs <- gets mach_registers
  liftIO $! regs `unsafeRead` (fromIntegral r)
{-# INLINE getReg #-}

-- setReg r i sets the value of register r to i
setReg :: Reg -> UInt32 -> MA ()
setReg r i = r `seq` i `seq` do
  regs <- gets mach_registers
  liftIO $! unsafeWrite regs (fromIntegral r) i
{-# INLINE setReg #-}

-- Index the platter array
getPlatterArray :: UInt32 -> MA MachPlatterArray
getPlatterArray i = do
  if i == 0
     then gets mach_program
     else do plattersTable <- gets mach_plattersTable
             case IM.lookup (fromIntegral i) plattersTable of
               Just pa -> return pa
               Nothing -> fail ("no platter array with id " ++ show i ++ " found")
{-# INLINE getPlatterArray #-}

addPlatterArray :: UInt32 -> MachPlatterArray -> MA ()
addPlatterArray i arr = do
  plattersTable <- gets mach_plattersTable
  let m = IM.insert (fromIntegral i) arr plattersTable
  modify $ \s -> s { mach_plattersTable = m }
{-# INLINE addPlatterArray #-}

-- remove a platter array
deletePlatterArray :: UInt32 -> MA ()
deletePlatterArray 0 = fail ("cannot delete platter array '0'")
deletePlatterArray i = do
  plattersTable <- gets mach_plattersTable
  let m = IM.delete (fromIntegral i) plattersTable
  modify $! \s -> s { mach_plattersTable = m }

-- getPlatter i j access the platter array with id i at position j
getPlatter :: UInt32 -> UInt32 -> MA UInt32
getPlatter i j = do
  arr <- getPlatterArray i
  liftIO $! arr `unsafeRead` (fromIntegral j)
{-# INLINE getPlatter #-}

-- setPlatter i j k sets the platter array with id i at position j to k
setPlatter :: UInt32 -> UInt32 -> UInt32 -> MA ()
setPlatter i j k = do
  arr <- getPlatterArray i
  liftIO $! unsafeWrite arr (fromIntegral j) k
{-# INLINE setPlatter #-}

-- The program counter is advanced to the next platter, if any.
incPC :: MA ()
incPC = modify $ \s -> let pc = 1 + mach_pc s in pc `seq` s { mach_pc = pc }
{-# INLINE incPC #-}

getInstr :: MA Instruction
getInstr = do
  pc        <- gets mach_pc
  prog      <- gets mach_program
  instrCode <- liftIO $ prog `unsafeRead` (fromIntegral pc)
  return $! decode instrCode
{-# INLINE getInstr #-}

-- initialState takes a program and constructs the initial machine state
--
-- The '0' array's contents are read from the "program" scroll.
--
-- all registers are initialised to platter value '0'
-- PC points to first platter of the '0' array, with offset '0'
--
initialState :: MachPlatterArray -> IO MachState
initialState prog = do
  regs <- newArray (minReg, maxReg) (toUInt32 (0::UInt32))
  return $ MachState {
      mach_registers              = regs,
      mach_plattersTable          = IM.empty,
      mach_program                = prog,
      mach_pc                     = toUInt32 (0 :: UInt32),
      mach_nextFreePlatterArrayId = toUInt32 (1 :: UInt32)
    }

-- runs the given program
run :: MachPlatterArray -> IO MachState
run prog = do
  state <- initialState prog
  execStateT execAll state

-- executes until the Halt instruction is reached
execAll :: MA ()
execAll = do
  instr <- getInstr
  unless (instr == Halt) $ do
    b <- exec instr
    when b incPC
    execAll

-- executes the next instruction and updates the machine state accordingly.
-- exec returns True if the caller of exec should increment the program
-- counter by one; False is returned if the caller should not touch the PC.
exec :: Instruction -> MA Bool
exec s | s `seq` False = undefined -- strict

-- Conditional move: A=B, unless C=0
exec (CMove regA regB regC) = {-# SCC "CMOVE" #-} do
  c <- getReg regC
  unless (c == 0) $ do
    b <- getReg regB
    setReg regA b
  return True

-- Array Index: A=array[B][C]
exec (ArrI regA regB regC) = {-# SCC "ARRI" #-} do
  b <- getReg regB
  c <- getReg regC
  i <- getPlatter b c
  setReg regA i
  return True

-- Array Amendment: array[A][B]=C
exec (ArrA regA regB regC) = {-# SCC "ARRA" #-} do
  a <- getReg regA
  b <- getReg regB
  c <- getReg regC
  setPlatter a b c
  return True

-- Addition: A=B + C, modulo 2^32
exec (Add regA regB regC) = {-# SCC "ADD" #-} do
  b <- getReg regB
  c <- getReg regC
  setReg regA (b + c)
  return True

-- Multiplication: A=B*C, mod 2^32
exec (Mult regA regB regC) = {-# SCC "MULT" #-} do
  b <- getReg regB
  c <- getReg regC
  setReg regA (b * c)
  return True

-- Division: unless (C==0) $ A=B `div` C
exec (Div regA regB regC) = {-# SCC "DIV" #-} do
  b <- getReg regB
  c <- getReg regC
  if c == 0
     then fail "divide by zero"
     else setReg regA (b `div` c)
  return True

-- Not-And: A= ....
exec (NAnd regA regB regC) = {-# SCC "NAND" #-} do
  b <- getReg regB
  c <- getReg regC
  setReg regA (complement $ b .&. c)
  return True

exec Halt = fail "exec should not be called for Halt"

-- allocate new array of 0s size C. B gets a unique id.
exec (Alloc regB regC) = {-# SCC "ALLOC" #-} do
  c    <- getReg regC
  arr  <- liftIO $ newArray (0, c - 1) 0
  idnt <- gets mach_nextFreePlatterArrayId
  if (idnt == 0)
     then fail "no more platter array identifiers available"
     else do modify (\s -> s { mach_nextFreePlatterArrayId = idnt + 1 })
             setReg regB idnt
             addPlatterArray idnt arr
  return True

-- Free an array
exec (Aband regC) = {-# SCC "ABAND" #-} do
  c <- getReg regC
  deletePlatterArray c
  return True

-- Display a value to the console
exec (Outp regC) = {-# SCC "OUTP" #-} do
  c <- getReg regC
  if c > toUInt32 (255 :: UInt32) then
     fail ("cannot print character " ++ show c)
     else let char = chr (fromIntegral c)
          in liftIO (putChar char >> hFlush stdout)
  return True

-- Block, waiting for input.
exec (Inp regC) = {-# SCC "INP" #-} do
  c <- getc
  setReg regC c
  return True

-- This is a frequent operation
-- There's an intentional delay to slow us down. The key is to avoid
-- copies where possible. In particular, array[0] = array[0] should be a
-- null op.
exec (Load regB regC) = {-# SCC "LOAD" #-} do
  b    <- getReg regB
  c    <- getReg regC
  if b == 0 -- array 0 copied to array 0 "must be handled with utmost velocity"
    then modify (\s -> s { mach_pc = c }) -- a branch
    else do arr  <- getPlatterArray b
            copy <- liftIO $ mapArray id arr
            modify (\s -> s { mach_program = copy, mach_pc = c })
  return False

-- Load immediate
exec (Ortho regA val) = {-# SCC "ORTHO" #-} do
  setReg regA val
  return True

getc :: MA UInt32
getc = do
  c <- liftIO $ (getChar >>= return . Just) `catch` (\_ -> return Nothing)
  return $ case c of
    Just c' -> toUInt32 (ord c')
    Nothing -> maxBound -- 1111111111