refactored evaluation of arithmetic sequences

module Evaluator where

import Prelude (class Show, top, class Semigroup, class Functor, class Monad, Unit, Ordering(..), (++), ($), (||), (&&), unit, return, (<*>), (<$>), pure, void, (==), otherwise, (>>=), (<), negate, (>), (>=), (<=), (/=), (-), (+), mod, div, (*), (<<<), compare, id, const, bind, show, map)
import Data.List (List(Nil, Cons), singleton, concatMap, intersect, zipWith, zipWithA, length, (:), replicate, drop, updateAt, (!!),concat)
import Data.StrMap (StrMap)
import Data.StrMap as Map
import Data.Tuple (Tuple(Tuple))
import Data.Maybe (Maybe(Nothing, Just), fromMaybe)
import Data.Foldable (foldl, foldr, foldMap, product)
import Data.Traversable (traverse)
import Data.Identity (Identity, runIdentity)
import Data.Either (Either(..), either)
import Data.Monoid (class Monoid)
import Data.String (fromChar, toChar)
import Data.Enum (fromEnum, toEnum)
import Control.Bind (join)
import Control.Apply ((*>))
import Control.Alt ((<|>))
import Control.Monad.State.Trans (StateT, modify, execStateT)
import Control.Monad.Except.Trans (ExceptT, throwError, runExceptT)

import AST (Atom(..), Binding(..), Definition(Def), Expr(..), Op(..),Path(..))

data EvalError =
    PathError Path Expr
  | IndexError Int Int
  | DivByZero
  | EvalError Expr
  | BinaryOpError Op Expr Expr
  | UnaryOpError Op Expr
  | NameCaptureError (List String)
  | UnknownFunction String
  | NoMatchingFunction String (List MatchingError)
  | CannotEvaluate Expr
  | NoError
  | MoreErrors EvalError EvalError

instance semigroupEvalError :: Semigroup EvalError where
  append e1 e2 = MoreErrors e1 e2

instance monoidEvalError :: Monoid EvalError where
  mempty = NoError

instance showEvalError :: Show EvalError where
  show (PathError p e)            = "PathError " ++ show p ++ " " ++ show e
  show (IndexError i l)           = "IndexError " ++ show i ++ " " ++ show l
  show DivByZero                  = "DivByZero"
  show (EvalError e)              = "EvalError " ++ show e
  show (BinaryOpError o e1 e2)    = "BinaryOpError " ++ show o ++ " (" ++ show e1 ++ ") (" ++ show e2 ++ ")"
  show (UnaryOpError o e)         = "UnaryOpError " ++ show o ++ " " ++ show e
  show (NameCaptureError ns)      = "NameCaptureError " ++ show ns
  show (UnknownFunction n)        = "UnknownFunction " ++ show n
  show (NoMatchingFunction n mes) = "NoMatchingFunction " ++ show n ++ " " ++ show mes
  show (CannotEvaluate e)         = "CannotEvaluate " ++ show e
  show NoError                    = "NoError"
  show (MoreErrors e1 e2)         = "(MoreErrors " ++ show e1 ++ " " ++ show e2 ++ ")"

data MatchingError =
    MatchingError Binding Expr
  | StrictnessError Binding Expr
  | TooFewArguments (List Binding) (List Expr)

instance showMatchingError :: Show MatchingError where
  show (MatchingError b e)     = "MatchingError " ++ show b ++ " " ++ show e
  show (StrictnessError b e)   = "StrictnessError " ++ show b ++ " " ++ show e
  show (TooFewArguments bs es) = "TooFewArguments " ++ show bs ++ " " ++ show es

type Evaluator = ExceptT EvalError Identity

runEvalM :: forall a. Evaluator a -> Either EvalError a
runEvalM = runIdentity <<< runExceptT

type Matcher = ExceptT MatchingError Identity

runMatcherM :: forall a. Matcher a -> Either MatchingError a
runMatcherM = runIdentity <<< runExceptT

mapWithPath :: Path -> (Expr -> Evaluator Expr) -> Expr -> Evaluator Expr
mapWithPath p f = go p
  where
  go End e     = f e
  go (Fst p) e = case e of
    Binary op e1 e2 -> Binary op <$> go p e1 <*> pure e2
    Unary op e      -> Unary op  <$> go p e
    SectL e op      -> SectL     <$> go p e <*> pure op
    IfExpr ce te ee -> IfExpr    <$> go p ce <*> pure te <*> pure ee
    ArithmSeq s b e -> ArithmSeq <$> go p s <*> pure b <*> pure e
    Lambda bs body  -> Lambda bs <$> go p body
    App e es        -> App       <$> go p e <*> pure es
    _               -> throwError $ PathError (Fst p) e
  go (Snd p) e = case e of
    Binary op e1 e2 -> Binary op e1 <$> go p e2
    SectR op e      -> SectR op     <$> go p e
    IfExpr ce te ee -> IfExpr <$> pure ce <*> go p te <*> pure ee
    ArithmSeq s b e -> ArithmSeq <$> pure s <*> (mapM' (go p) b) <*> pure e
    _               -> throwError $ PathError (Snd p) e
  go (Thrd p) e = case e of
    IfExpr ce te ee -> IfExpr <$> pure ce <*> pure te <*> go p ee
    ArithmSeq s b (Just e) -> ArithmSeq <$> pure s <*> pure b <*> (Just <$> (go p e))
    _               -> throwError $ PathError (Thrd p) e
  go (Nth n p) e = case e of
    List es  -> List  <$> mapIndex n (go p) es
    NTuple es -> NTuple <$> mapIndex n (go p) es
    App e es -> App e <$> mapIndex n (go p) es
    _        -> throwError $ PathError (Nth n p) e

mapIndex :: forall a. Int -> (a -> Evaluator a) -> (List a) -> Evaluator (List a)
mapIndex i f as = do
  case as !! i of
    Nothing -> throwError $ IndexError i (length as)
    Just a  -> do
      a' <- f a
      case updateAt i a' as of
        Nothing  -> throwError $ IndexError i (length as)
        Just as' -> return as'

evalPath1 :: Env -> Path -> Expr -> Either EvalError Expr
evalPath1 env path expr = runEvalM $ mapWithPath path (eval1 env) expr

evalPathAll :: Env -> Path -> Expr -> Either EvalError Expr
evalPathAll env path expr = runEvalM $ mapWithPath path (return <<< eval env) expr

type Env = StrMap (List (Tuple (List Binding) Expr))

defsToEnv :: (List Definition) -> Env
defsToEnv = foldl insertDef Map.empty

envToDefs :: Env -> (List Definition)
envToDefs env = concat $ map tupleToDef $ Data.StrMap.toList env
  where
    tupleToDef (Tuple name defs) = map
                                    (\(Tuple bin expr) -> Def name bin expr)
                                    defs

insertDef :: Env -> Definition -> Env
insertDef env (Def name bindings body) = case Map.lookup name env of
  Nothing   -> Map.insert name (singleton $ Tuple bindings body) env
  Just defs -> Map.insert name (defs ++ (singleton $ Tuple bindings body)) env

eval1 :: Env -> Expr -> Evaluator Expr
eval1 env expr = case expr of
  (Binary op e1 e2)                  -> binary env op e1 e2
  (Unary op e)                       -> unary env op e
  (Atom (Name name))                 -> apply env name Nil
  (IfExpr (Atom (Bool true)) te _)   -> return te
  (IfExpr (Atom (Bool false)) _ ee)  -> return ee
  (ArithmSeq start step end)         -> evalArithmSeq start step end
--  (List (e:es))                      -> return $ Binary Cons e (List es)
  (App (Binary Composition f g) (Cons e Nil)) -> return $ App f (singleton $ App g (Cons e Nil))
  (App (Lambda binds body) args)     -> tryAll env (singleton $ Tuple binds body) args "lambda" Nil
  (App (SectL e1 op) (Cons e2 Nil))           -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (SectR op e2) (Cons e1 Nil))           -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (PrefixOp op) (Cons e1 (Cons e2 Nil)))         -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (Atom (Name name)) args)      -> apply env name args
  (App (App func es) es')            -> return $ App func (es ++ es')
  _                                  -> throwError $ CannotEvaluate expr

eval :: Env -> Expr -> Expr
eval env expr = evalToBinding env expr (Lit (Name "_|_"))

evalToBinding :: Env -> Expr -> Binding -> Expr
evalToBinding env expr bind = case bind of
  (Lit (Name "_|_")) -> recurse env expr bind
  (Lit (Name _))     -> expr
  (Lit _)            -> case expr of
    (Atom _) -> expr
    _        -> recurse env expr bind

  (ConsLit b bs)     -> case expr of
    (Binary Colon e es) -> Binary Colon (evalToBinding env e b) (evalToBinding env es bs)
    (List (Cons e es))  -> evalToBinding env (Binary Colon e (List es)) bind
    _                   -> recurse env expr bind

  (ListLit bs)       -> case expr of
    (List es) -> if (length es == length bs) then List (zipWith (evalToBinding env) es bs) else expr
    _         -> recurse env expr bind

  (NTupleLit bs)     -> case expr of
    (NTuple es)  -> NTuple (zipWith (evalToBinding env) es bs)
    _            -> recurse env expr bind

recurse :: Env -> Expr -> Binding -> Expr
recurse env expr bind = if expr == eval1d then expr else evalToBinding env eval1d bind
  where
    eval1d :: Expr
    eval1d = either (const expr') id $ runEvalM $ eval1 env expr'
    expr' :: Expr
    expr' = case expr of
      (Binary op e1 e2)  ->
        Binary op (evalToBinding env e1 bind) (evalToBinding env e2 bind)
      (Unary op e)       ->
        Unary op (evalToBinding env e bind)
      (List es)          ->
        List ((\e -> evalToBinding env e bind) <$> es)
      (NTuple es)        ->
        NTuple ((\e -> evalToBinding env e bind) <$> es)
      (IfExpr c t e)     ->
        IfExpr (evalToBinding env c bind) t e
      (ArithmSeq c t e)     ->
        ArithmSeq (evalToBinding env c bind) ((\x -> evalToBinding env x bind) <$> t) ((\x -> evalToBinding env x bind) <$> e)
      (App f args)       -> do
        App (evalToBinding env f bind) args
      _                  ->
        expr

wrapLambda :: (List Binding) -> (List Expr) -> Expr -> Evaluator Expr
wrapLambda binds args body =
  case compare (length binds) (length args) of
    EQ -> return body
    GT -> return $ Lambda (drop (length args) binds) body
    LT -> return $ App body (drop (length binds) args)

------------------------------------------------------------------------------------------
-- Arithmetic Sequences
------------------------------------------------------------------------------------------

{-
packs the evaluation result for arithmetic sequences (AS)
example:
Trip (Just x) (Just y) (Just z) will be displayed as x : [y, z ..]
if no end of (AS) was given or x : [y, z .. end] if end was given
-}
data Trip a = Trip a a a

instance functorTrip :: Functor Trip where
  map f (Trip x y z) = Trip (f x) (f y) (f z)

intFromStepTo :: Int -> Maybe Int -> Maybe Int -> Trip (Maybe Int)
intFromStepTo start Nothing Nothing     = Trip (Just start) (Just (start + 1)) Nothing
intFromStepTo start (Just step) Nothing = Trip (Just start) (Just step) (Just (step + step - start))
intFromStepTo start Nothing (Just end)  = case start > end of
  true  -> Trip Nothing Nothing Nothing
  false -> case start == end of
    true  -> Trip (Just end) Nothing Nothing
    false -> Trip (Just start) (Just (start + 1)) Nothing
intFromStepTo start (Just step) (Just end) = case (start <= step && start > end) || (start > step && start < end) of
  true  -> Trip Nothing Nothing Nothing
  false -> case start == end of
    true  -> Trip (Just end) Nothing Nothing
    false -> Trip (Just (start)) (Just (step)) (Just (step + step - start))

--detect whether top or bottom for Boolean was reached
intToBool :: Trip (Maybe Int) -> Trip (Maybe Boolean)
intToBool trip = case temp of
  Trip (Just x) (Just y) z -> if x == y then Trip (Just x) Nothing Nothing else temp
  _ -> temp
  where
    temp :: Trip (Maybe Boolean)
    temp = (\x -> intToBool' <$> x) <$> trip

    intToBool' :: Int -> Boolean
    intToBool' i = if i <= 0 then false else true

exprFromStepTo :: Expr -> Maybe Expr -> Maybe Expr -> Trip (Maybe Expr)
exprFromStepTo start step end = case start of
  Atom (AInt i) -> (\x -> (Atom <<< AInt) <$> x) <$> intTrip
  Atom (Bool b) -> (\x -> (Atom <<< Bool) <$> x) <$> (intToBool intTrip)
  Atom (Char c) -> (\x -> (Atom <<< Char <<< fromChar) <$> x) <$> (\x -> join (toEnum <$> x)) <$> intTrip
  _             -> Trip Nothing Nothing Nothing
    where
      intTrip = intFromStepTo (unsafeExprToInt start) (unsafeExprToInt <$> step) (unsafeExprToInt <$> end)

unsafeExprToInt :: Expr -> Int
unsafeExprToInt (Atom (AInt i))   = i
unsafeExprToInt (Atom (Bool b))   = fromEnum b
unsafeExprToInt (Atom (Char str)) = fromEnum $ fromMaybe 'E' (toChar str)
unsafeExprToInt _ = top

evalArithmSeq :: Expr -> Maybe Expr -> Maybe Expr -> Evaluator Expr
evalArithmSeq start step end = case foldr (&&) true (isValid <$> [Just start, step, end]) of
  false -> throwError $ CannotEvaluate $ ArithmSeq start step end
  true  -> evalArithmSeq'
  where
    isValid :: Maybe Expr -> Boolean
    isValid Nothing = true
    isValid (Just (Atom (Name _))) = false
    isValid (Just (Atom _)) = true
    isValid _ = false

    evalArithmSeq' :: Evaluator Expr
    evalArithmSeq' = case (exprFromStepTo start step end) of
      Trip (Just a) (Just b) nextStep -> return $ Binary Colon a (ArithmSeq b nextStep end)
      Trip (Just a) Nothing _         -> return $ List (singleton a)
      Trip Nothing _ _                -> return $ List Nil

------------------------------------------------------------------------------------------

binary :: Env -> Op -> Expr -> Expr -> Evaluator Expr
binary env op = case op of
  Power  -> aint Power (\i j -> product $ replicate j i)
  Mul    -> aint Mul (*)
  Add    -> aint Add (+)
  Sub    -> aint Sub (-)
  Colon  -> \e es -> case es of
    (List ls) -> return $ List $ e:ls
    _         -> throwError $ BinaryOpError Colon e es
  Append -> \es1 es2 -> case (Tuple es1 es2) of
    (Tuple (List ls1) (List ls2)) -> return $ List $ ls1 ++ ls2
    _                             -> throwError $ BinaryOpError Append es1 es2
  Equ    -> ord Equ (==) (==) (==)
  Neq    -> ord Neq (/=) (/=) (/=)
  Leq    -> ord Leq (<=) (<=) (<=)
  Lt     -> ord Lt  (<)  (<)  (<)
  Geq    -> ord Geq (>=) (>=) (>=)
  Gt     -> ord Gt  (>)  (>)  (>)
  And    -> \e1 e2 -> case (Tuple e1 e2) of
    (Tuple (Atom (Bool false)) _                  ) -> return $ Atom $ Bool false
    (Tuple _                   (Atom (Bool false))) -> return $ Atom $ Bool false
    (Tuple (Atom (Bool true))  (Atom (Bool true)) ) -> return $ Atom $ Bool true
    (Tuple _                   _                  ) -> throwError $ BinaryOpError And e1 e2
  Or     -> \e1 e2 -> case (Tuple e1 e2) of
    (Tuple (Atom (Bool true))  _                   ) -> return $ Atom $ Bool true
    (Tuple _                   (Atom (Bool true))  ) -> return $ Atom $ Bool true
    (Tuple (Atom (Bool false))  (Atom (Bool false))) -> return $ Atom $ Bool false
    (Tuple _                   _                   ) -> throwError $ BinaryOpError And e1 e2
  Dollar -> (\f e -> return $ App f (singleton e))
  Composition -> \e1 e2 -> throwError $ BinaryOpError And e1 e2
  InfixFunc name -> \e1 e2 -> apply env name (e1 : e2 : Nil)
  where
    aint :: Op -> (Int -> Int -> Int) -> Expr -> Expr -> Evaluator Expr
    aint _   f (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ f i j
    aint op  _ e1              e2              = throwError $ BinaryOpError op e1 e2

    ord :: Op -> (Int -> Int -> Boolean) -> (String -> String -> Boolean) -> (Boolean -> Boolean -> Boolean)-> Expr -> Expr -> Evaluator Expr
    ord _  fi _  _  (Atom (AInt i))  (Atom (AInt j))  = return $ Atom $ Bool $ fi i j
    ord _  _  fs _  (Atom (Char s1)) (Atom (Char s2)) = return $ Atom $ Bool $ fs s1 s2
    ord _  _  _  fb (Atom (Bool b1)) (Atom (Bool b2)) = return $ Atom $ Bool $ fb b1 b2
    ord op _  _  _  e1               e2               = throwError $ BinaryOpError op e1 e2

unary :: Env -> Op -> Expr -> Evaluator Expr
unary env Sub (Atom (AInt i)) = return $ Atom $ AInt (-i)
unary env op e = throwError $ UnaryOpError op e

apply :: Env -> String -> (List Expr) -> Evaluator Expr
apply env "div" (Cons e1 (Cons e2 _)) = division e1 e2
apply env "mod" (Cons e1 (Cons e2 _)) = modulo e1 e2
apply env name args = case Map.lookup name env of
  Nothing    -> throwError $ UnknownFunction name
  Just cases -> tryAll env cases args name Nil

--builtin ipo
division :: Expr -> Expr -> Evaluator Expr
division (Atom (AInt i)) (Atom (AInt 0)) = throwError DivByZero
division (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ div i j
division e1 e2 = throwError $ BinaryOpError (InfixFunc "div") e1 e2

--builtin mod
modulo :: Expr -> Expr -> Evaluator Expr
modulo (Atom (AInt i)) (Atom (AInt 0)) = throwError DivByZero
modulo (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ mod i j
modulo e1 e2 = throwError $ BinaryOpError (InfixFunc "mod") e1 e2

tryAll :: Env -> (List (Tuple (List Binding) Expr)) -> (List Expr) -> String -> (List MatchingError) -> Evaluator Expr
tryAll _   Nil                        _    name errs = throwError $ NoMatchingFunction name errs
tryAll env (Cons (Tuple binds body) cases) args name errs | (length args) < (length binds) = tryAll env cases args name (errs ++ (singleton $ TooFewArguments binds args))
tryAll env (Cons (Tuple binds body) cases) args name errs = case runMatcherM $ matchls' env binds args of
  Right replMap                 -> replace' replMap body >>= wrapLambda binds args
  Left se@(StrictnessError _ _) -> throwError $ NoMatchingFunction name (errs ++ singleton se)
  Left err                      -> tryAll env cases args name (errs ++ singleton err)

whnf :: Expr -> Boolean
whnf (Atom (Name _)) = false
whnf (Atom _)        = true
whnf (List _)        = true
whnf (NTuple _)      = true
whnf _               = false

checkStrictness :: Binding -> Expr -> MatchingError
checkStrictness bs es | whnf es   = MatchingError bs es
                      | otherwise = StrictnessError bs es

matchls' :: Env -> (List Binding) -> (List Expr) -> Matcher (StrMap Expr)
matchls' env bs es = execStateT (zipWithA (\b e -> match' b (evalToBinding env e b)) bs es) Map.empty

match' :: Binding -> Expr -> StateT (StrMap Expr) Matcher Unit
match' (Lit (Name name)) e                   = modify (Map.insert name e)
match' (Lit ba)          (Atom ea)           = case ba == ea of
                                                 true  -> return unit
                                                 false -> throwError $ MatchingError (Lit ba) (Atom ea)
match' (Lit b)           e                   = throwError $ checkStrictness (Lit b) e

match' (ConsLit b bs)    (Binary Colon e es)  = match' b e *> match' bs es
match' (ConsLit b bs)    (List (Cons e es))  = match' (ConsLit b bs) (Binary Colon e (List es))
match' (ConsLit b bs)    (List Nil)           = throwError $ MatchingError (ConsLit b bs) (List Nil)
match' (ConsLit b bs)    e                   = throwError $ checkStrictness (ConsLit b bs) e

match' (ListLit (Cons b bs))  (Binary Colon e es)  = match' b e *> match' (ListLit bs) es
match' (ListLit bs)      (List es)           = case length bs == length es of
                                                 true  -> void $ zipWithA match' bs es
                                                 false -> throwError $ MatchingError (ListLit bs) (List es)
match' (ListLit bs)      e                   = throwError $ checkStrictness (ListLit bs) e

match' (NTupleLit bs)    (NTuple es)         = case length bs == length es of
                                                 true  -> void $ zipWithA match' bs es
                                                 false -> throwError $ MatchingError (NTupleLit bs) (NTuple es)
match' (NTupleLit bs)    e                   = throwError $ checkStrictness (NTupleLit bs) e

mapM' :: forall a b m. (Monad m) => (a -> m b) -> Maybe a -> m (Maybe b)
mapM' f Nothing  = pure Nothing
mapM' f (Just x) = Just <$> (f x)

replace' :: StrMap Expr -> Expr -> Evaluator Expr
replace' subs = go
  where
  go expr = case expr of
    a@(Atom (Name name)) -> case Map.lookup name subs of
      Just subExpr -> return subExpr
      Nothing      -> return a
    (List exprs)         -> List <$> (traverse go exprs)
    (NTuple exprs)       -> NTuple <$> (traverse go exprs)
    (Binary op e1 e2)    -> Binary <$> pure op <*> go e1 <*> go e2
    (Unary op e)         -> Unary <$> pure op <*> go e
    (SectL e op)         -> SectL <$> go e <*> pure op
    (SectR op e)         -> SectR <$> pure op <*> go e
    (IfExpr ce te ee)    -> IfExpr <$> go ce <*> go te <*> go ee
    (ArithmSeq ce te ee) -> ArithmSeq <$> go ce <*> (mapM' go te) <*> (mapM' go ee)
    (Lambda binds body)  -> (avoidCapture subs binds) *> (Lambda <$> pure binds <*> replace' (foldr Map.delete subs (boundNames' binds)) body)
    (App func exprs)     -> App <$> go func <*> traverse go exprs
    e                    -> return e

avoidCapture :: StrMap Expr -> (List Binding) -> Evaluator Unit
avoidCapture subs binds = case intersect (concatMap freeVariables $ Map.values subs) (boundNames' binds) of
  Nil       -> return unit
  captures -> throwError $ NameCaptureError captures

freeVariables :: Expr -> (List String)
freeVariables _ = Nil
-- freeVariables = nub <<< foldExpr
--   (\a -> case a of
--     Name name -> singleton $ name
--     _         -> Nil)
--   concat
--   concat
--   (\_ f1 f2 -> f1 ++ f2)
--   (\_ f -> f)
--   (\f _ -> f)
--   (\_ f -> f)
--   (\_ -> Nil)
--   (\f1 f2 f3 -> f1 ++ f2 ++ f3)
--   (\bs f -> nub f \\ boundNames' bs)
--   (\f fs -> f ++ concat fs)

boundNames' :: (List Binding) -> (List String)
boundNames' = concatMap boundNames

boundNames :: Binding -> (List String)
boundNames = go
  where
  go (Lit (Name name)) = singleton name
  go (ConsLit b1 b2)   = go b1 ++ go b2
  go (ListLit bs)      = foldMap go bs
  go (NTupleLit bs)    = foldMap go bs
  go _                 = Nil

module Evaluator where

import Prelude (class Show, top, bottom, class Semigroup, class Monad, Unit, Ordering(..), (++), ($), unit, return, (<*>), (<$>), pure, void, (==), otherwise, (>>=), (<), negate, (>), (>=), (<=), (/=), (-), (+), mod, div, (*), (<<<), compare, id, const, bind, show, map)
import Data.List (List(Nil, Cons), singleton, concatMap, intersect, zipWith, zipWithA, length, (:), replicate, drop, updateAt, (!!),concat)
import Data.StrMap (StrMap)
import Data.StrMap as Map
import Data.Tuple (Tuple(Tuple), fst, snd)
import Data.Maybe (Maybe(Nothing, Just), fromMaybe)
import Data.Foldable (foldl, foldr, foldMap, product)
import Data.Traversable (traverse)
import Data.Identity (Identity, runIdentity)
import Data.Either (Either(..), either)
import Data.Monoid (class Monoid)
import Data.String (fromChar, toChar)
import Data.Enum (class Enum, fromEnum, toEnum, succ)
import Control.Bind (join)
import Control.Apply ((*>))
import Control.Alt ((<|>))
import Control.Monad.State.Trans (StateT, modify, execStateT)
import Control.Monad.Except.Trans (ExceptT, throwError, runExceptT)

import AST (Atom(..), Binding(..), Definition(Def), Expr(..), Op(..),Path(..))

data EvalError =
    PathError Path Expr
  | IndexError Int Int
  | DivByZero
  | EvalError Expr
  | BinaryOpError Op Expr Expr
  | UnaryOpError Op Expr
  | NameCaptureError (List String)
  | UnknownFunction String
  | NoMatchingFunction String (List MatchingError)
  | CannotEvaluate Expr
  | NoError
  | MoreErrors EvalError EvalError

instance semigroupEvalError :: Semigroup EvalError where
  append e1 e2 = MoreErrors e1 e2

instance monoidEvalError :: Monoid EvalError where
  mempty = NoError

instance showEvalError :: Show EvalError where
  show (PathError p e)            = "PathError " ++ show p ++ " " ++ show e
  show (IndexError i l)           = "IndexError " ++ show i ++ " " ++ show l
  show DivByZero                  = "DivByZero"
  show (EvalError e)              = "EvalError " ++ show e
  show (BinaryOpError o e1 e2)    = "BinaryOpError " ++ show o ++ " (" ++ show e1 ++ ") (" ++ show e2 ++ ")"
  show (UnaryOpError o e)         = "UnaryOpError " ++ show o ++ " " ++ show e
  show (NameCaptureError ns)      = "NameCaptureError " ++ show ns
  show (UnknownFunction n)        = "UnknownFunction " ++ show n
  show (NoMatchingFunction n mes) = "NoMatchingFunction " ++ show n ++ " " ++ show mes
  show (CannotEvaluate e)         = "CannotEvaluate " ++ show e
  show NoError                    = "NoError"
  show (MoreErrors e1 e2)         = "(MoreErrors " ++ show e1 ++ " " ++ show e2 ++ ")"

data MatchingError =
    MatchingError Binding Expr
  | StrictnessError Binding Expr
  | TooFewArguments (List Binding) (List Expr)

instance showMatchingError :: Show MatchingError where
  show (MatchingError b e)     = "MatchingError " ++ show b ++ " " ++ show e
  show (StrictnessError b e)   = "StrictnessError " ++ show b ++ " " ++ show e
  show (TooFewArguments bs es) = "TooFewArguments " ++ show bs ++ " " ++ show es

type Evaluator = ExceptT EvalError Identity

runEvalM :: forall a. Evaluator a -> Either EvalError a
runEvalM = runIdentity <<< runExceptT

type Matcher = ExceptT MatchingError Identity

runMatcherM :: forall a. Matcher a -> Either MatchingError a
runMatcherM = runIdentity <<< runExceptT

mapWithPath :: Path -> (Expr -> Evaluator Expr) -> Expr -> Evaluator Expr
mapWithPath p f = go p
  where
  go End e     = f e
  go (Fst p) e = case e of
    Binary op e1 e2 -> Binary op <$> go p e1 <*> pure e2
    Unary op e      -> Unary op  <$> go p e
    SectL e op      -> SectL     <$> go p e <*> pure op
    IfExpr ce te ee -> IfExpr    <$> go p ce <*> pure te <*> pure ee
    ArithmSeq s b e -> ArithmSeq <$> go p s <*> pure b <*> pure e
    Lambda bs body  -> Lambda bs <$> go p body
    App e es        -> App       <$> go p e <*> pure es
    _               -> throwError $ PathError (Fst p) e
  go (Snd p) e = case e of
    Binary op e1 e2 -> Binary op e1 <$> go p e2
    SectR op e      -> SectR op     <$> go p e
    IfExpr ce te ee -> IfExpr <$> pure ce <*> go p te <*> pure ee
    ArithmSeq s b e -> ArithmSeq <$> pure s <*> (mapM' (go p) b) <*> pure e
    _               -> throwError $ PathError (Snd p) e
  go (Thrd p) e = case e of
    IfExpr ce te ee -> IfExpr <$> pure ce <*> pure te <*> go p ee
    ArithmSeq s b (Just e) -> ArithmSeq <$> pure s <*> pure b <*> (Just <$> (go p e))
    _               -> throwError $ PathError (Thrd p) e
  go (Nth n p) e = case e of
    List es  -> List  <$> mapIndex n (go p) es
    NTuple es -> NTuple <$> mapIndex n (go p) es
    App e es -> App e <$> mapIndex n (go p) es
    _        -> throwError $ PathError (Nth n p) e

mapIndex :: forall a. Int -> (a -> Evaluator a) -> (List a) -> Evaluator (List a)
mapIndex i f as = do
  case as !! i of
    Nothing -> throwError $ IndexError i (length as)
    Just a  -> do
      a' <- f a
      case updateAt i a' as of
        Nothing  -> throwError $ IndexError i (length as)
        Just as' -> return as'

evalPath1 :: Env -> Path -> Expr -> Either EvalError Expr
evalPath1 env path expr = runEvalM $ mapWithPath path (eval1 env) expr

evalPathAll :: Env -> Path -> Expr -> Either EvalError Expr
evalPathAll env path expr = runEvalM $ mapWithPath path (return <<< eval env) expr

type Env = StrMap (List (Tuple (List Binding) Expr))

defsToEnv :: (List Definition) -> Env
defsToEnv = foldl insertDef Map.empty

envToDefs :: Env -> (List Definition)
envToDefs env = concat $ map tupleToDef $ Data.StrMap.toList env
  where
    tupleToDef (Tuple name defs) = map
                                    (\(Tuple bin expr) -> Def name bin expr)
                                    defs

insertDef :: Env -> Definition -> Env
insertDef env (Def name bindings body) = case Map.lookup name env of
  Nothing   -> Map.insert name (singleton $ Tuple bindings body) env
  Just defs -> Map.insert name (defs ++ (singleton $ Tuple bindings body)) env

eval1 :: Env -> Expr -> Evaluator Expr
eval1 env expr = case expr of
  (Binary op e1 e2)                  -> binary env op e1 e2
  (Unary op e)                       -> unary env op e
  (Atom (Name name))                 -> apply env name Nil
  (IfExpr (Atom (Bool true)) te _)   -> return te
  (IfExpr (Atom (Bool false)) _ ee)  -> return ee
  (ArithmSeq start step end)         -> evalArithmSeq $ ArithmSeq start step end
--  (List (e:es))                      -> return $ Binary Cons e (List es)
  (App (Binary Composition f g) (Cons e Nil)) -> return $ App f (singleton $ App g (Cons e Nil))
  (App (Lambda binds body) args)     -> tryAll env (singleton $ Tuple binds body) args "lambda" Nil
  (App (SectL e1 op) (Cons e2 Nil))           -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (SectR op e2) (Cons e1 Nil))           -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (PrefixOp op) (Cons e1 (Cons e2 Nil)))         -> binary env op e1 e2 <|> (return $ Binary op e1 e2)
  (App (Atom (Name name)) args)      -> apply env name args
  (App (App func es) es')            -> return $ App func (es ++ es')
  _                                  -> throwError $ CannotEvaluate expr

eval :: Env -> Expr -> Expr
eval env expr = evalToBinding env expr (Lit (Name "_|_"))

evalToBinding :: Env -> Expr -> Binding -> Expr
evalToBinding env expr bind = case bind of
  (Lit (Name "_|_")) -> recurse env expr bind
  (Lit (Name _))     -> expr
  (Lit _)            -> case expr of
    (Atom _) -> expr
    _        -> recurse env expr bind

  (ConsLit b bs)     -> case expr of
    (Binary Colon e es) -> Binary Colon (evalToBinding env e b) (evalToBinding env es bs)
    (List (Cons e es))  -> evalToBinding env (Binary Colon e (List es)) bind
    _                   -> recurse env expr bind

  (ListLit bs)       -> case expr of
    (List es) -> if (length es == length bs) then List (zipWith (evalToBinding env) es bs) else expr
    _         -> recurse env expr bind

  (NTupleLit bs)     -> case expr of
    (NTuple es)  -> NTuple (zipWith (evalToBinding env) es bs)
    _            -> recurse env expr bind

recurse :: Env -> Expr -> Binding -> Expr
recurse env expr bind = if expr == eval1d then expr else evalToBinding env eval1d bind
  where
    eval1d :: Expr
    eval1d = either (const expr') id $ runEvalM $ eval1 env expr'
    expr' :: Expr
    expr' = case expr of
      (Binary op e1 e2)  ->
        Binary op (evalToBinding env e1 bind) (evalToBinding env e2 bind)
      (Unary op e)       ->
        Unary op (evalToBinding env e bind)
      (List es)          ->
        List ((\e -> evalToBinding env e bind) <$> es)
      (NTuple es)        ->
        NTuple ((\e -> evalToBinding env e bind) <$> es)
      (IfExpr c t e)     ->
        IfExpr (evalToBinding env c bind) t e
      (ArithmSeq c t e)     ->
        ArithmSeq (evalToBinding env c bind) ((\x -> evalToBinding env x bind) <$> t) ((\x -> evalToBinding env x bind) <$> e)
      (App f args)       -> do
        App (evalToBinding env f bind) args
      _                  ->
        expr

wrapLambda :: (List Binding) -> (List Expr) -> Expr -> Evaluator Expr
wrapLambda binds args body =
  case compare (length binds) (length args) of
    EQ -> return body
    GT -> return $ Lambda (drop (length args) binds) body
    LT -> return $ App body (drop (length binds) args)

------------------------------------------------------------------------------------------
-- Arithmetic Sequences
------------------------------------------------------------------------------------------

atomNext :: Atom -> Tuple (Maybe Atom) (Maybe Atom)
atomNext (AInt i) = atomNextTo (AInt i) (AInt top)
atomNext (Bool b) = atomNextTo (Bool b) (Bool top)
atomNext (Char x) = atomNextTo (Char x) (Char (fromChar top))
atomNext _ = Tuple Nothing Nothing

--allways rising sequence
atomNextTo :: Atom -> Atom -> Tuple (Maybe Atom) (Maybe Atom)
atomNextTo (AInt x) (AInt y) = if x == y then Tuple (Just (AInt y)) Nothing else
  if x < y then Tuple (Just (AInt x)) (Just (AInt (x + 1))) else Tuple Nothing Nothing
atomNextTo (Bool x) (Bool y) = if x == y then Tuple (Just (Bool y)) Nothing else
  if x < y then Tuple (Just (Bool x)) (Bool <$> (succ x)) else Tuple Nothing Nothing
atomNextTo (Char x) (Char y) = case toChar x of
  Just a -> case toChar y of
    Just b  -> if a == b then Tuple (Just (Char y)) Nothing else
                 if a < b then Tuple (Just (Char x)) ((Char <<< fromChar) <$> (succ a))
                          else Tuple Nothing Nothing
    Nothing -> Tuple Nothing Nothing
  Nothing -> Tuple Nothing Nothing
atomNextTo _ _ = Tuple Nothing Nothing

intNextStepTo :: Int -> Int -> Int -> Tuple (Maybe Int) (Maybe Int)
intNextStepTo x y z = case x == z of
  true  -> Tuple (Just z) Nothing
  false -> case x <= y of
    true  -> if x > z then Tuple Nothing Nothing else Tuple (Just x) (Just (y + y - x))
    false -> if x < z then Tuple Nothing Nothing else Tuple (Just x) (Just (y + y - x))

eNextStepTo :: forall e. (Enum e) => e -> e -> e -> Tuple (Maybe e) (Maybe e)
eNextStepTo x y z = case (fromEnum x) <= (fromEnum y) of
  true  -> Tuple (join (toEnum <$> (fst tup))) ((\q -> fromMaybe top (toEnum q)) <$> (snd tup))
  false -> Tuple (join (toEnum <$> (fst tup))) ((\q -> fromMaybe bottom (toEnum q)) <$> (snd tup))
  where tup = intNextStepTo (fromEnum x) (fromEnum y) (fromEnum z)

atomNextStepTo :: Atom -> Atom -> Atom -> Tuple (Maybe Atom) (Maybe Atom)
atomNextStepTo (AInt x) (AInt y) (AInt z) = let tup = intNextStepTo x y z in
  Tuple (AInt <$> (fst tup)) (AInt <$> (snd tup))
atomNextStepTo (Bool x) (Bool y) (Bool z) = let tup = eNextStepTo x y z in
  Tuple (Bool <$> (fst) tup) (Bool <$> (snd tup))
atomNextStepTo (Char x) (Char y) (Char z) = case toChar x of
  Nothing -> Tuple Nothing Nothing
  Just a  -> case toChar y of
    Nothing -> Tuple Nothing Nothing
    Just b  -> case toChar z of
      Nothing -> Tuple Nothing Nothing
      Just c  -> let tup = eNextStepTo a b c in
        Tuple ((Char <<< fromChar) <$> (fst tup)) ((Char <<< fromChar) <$> (snd tup))
atomNextStepTo _ _ _ = Tuple Nothing Nothing

atomNextStep :: Atom -> Atom -> Tuple (Maybe Atom) (Maybe Atom)
atomNextStep (AInt x) (AInt y) = case x < y of
  true  -> atomNextStepTo (AInt x) (AInt y) (AInt top)
  false -> atomNextStepTo (AInt x) (AInt y) (AInt bottom)
atomNextStep (Bool x) (Bool y) = case x < y of
  true  -> atomNextStepTo (Bool x) (Bool y) (Bool top)
  false -> atomNextStepTo (Bool x) (Bool y) (Bool bottom)
atomNextStep (Char x) (Char y) = case toChar x of
  Just a  -> case toChar y of
    Just b  -> if a < b then atomNextStepTo (Char x) (Char y) (Char (fromChar top))
                        else atomNextStepTo (Char x) (Char y) (Char (fromChar bottom))
    Nothing -> Tuple Nothing Nothing
  Nothing -> Tuple Nothing Nothing
atomNextStep _ _ = Tuple Nothing Nothing

evalArithmSeq :: Expr -> Evaluator Expr
evalArithmSeq as = case as of
  ArithmSeq (Atom (AInt _)) _ _ -> evalArithmSeq' as
  ArithmSeq (Atom (Bool _)) _ _ -> evalArithmSeq' as
  ArithmSeq (Atom (Char _)) _ _ -> evalArithmSeq' as
  _                             -> throwError $ CannotEvaluate as
  where
    evalArithmSeq' :: Expr -> Evaluator Expr
    evalArithmSeq' as = case as of
      ArithmSeq (Atom x) Nothing Nothing -> case atomNext x of
        Tuple Nothing _         -> return $ List Nil
        Tuple (Just a) Nothing  -> return $ List (singleton (Atom a))
        Tuple (Just a) (Just b) -> return $ Binary Colon (Atom a) (ArithmSeq (Atom b) Nothing Nothing)

      ArithmSeq (Atom x) (Just (Atom y)) Nothing -> case atomNextStep x y of
        Tuple Nothing _                -> return $ List Nil
        Tuple (Just a) Nothing         -> return $ List (singleton (Atom a))
        Tuple (Just a) (Just nextStep) -> return $ Binary Colon (Atom a) (ArithmSeq (Atom y) (Just (Atom nextStep)) Nothing)

      ArithmSeq (Atom x) Nothing end@(Just (Atom z)) -> case atomNextTo x z of
        Tuple Nothing _         -> return $ List Nil
        Tuple (Just a) Nothing  -> return $ List (singleton (Atom a))
        Tuple (Just a) (Just b) -> return $ Binary Colon (Atom a) (ArithmSeq (Atom b) Nothing Nothing)

      ArithmSeq (Atom x) (Just (Atom y)) end@(Just (Atom z)) -> case atomNextStepTo x y z of
        Tuple Nothing _                -> return $ List Nil
        Tuple (Just a) Nothing         -> return $ List (singleton (Atom a))
        Tuple (Just a) (Just nextStep) -> return $ Binary Colon (Atom a) (ArithmSeq (Atom y) (Just (Atom nextStep)) end)

      _ -> throwError $ CannotEvaluate as

------------------------------------------------------------------------------------------
-- Arithmetic Sequences End
------------------------------------------------------------------------------------------

binary :: Env -> Op -> Expr -> Expr -> Evaluator Expr
binary env op = case op of
  Power  -> aint Power (\i j -> product $ replicate j i)
  Mul    -> aint Mul (*)
  Add    -> aint Add (+)
  Sub    -> aint Sub (-)
  Colon  -> \e es -> case es of
    (List ls) -> return $ List $ e:ls
    _         -> throwError $ BinaryOpError Colon e es
  Append -> \es1 es2 -> case (Tuple es1 es2) of
    (Tuple (List ls1) (List ls2)) -> return $ List $ ls1 ++ ls2
    _                             -> throwError $ BinaryOpError Append es1 es2
  Equ    -> ord Equ (==) (==) (==)
  Neq    -> ord Neq (/=) (/=) (/=)
  Leq    -> ord Leq (<=) (<=) (<=)
  Lt     -> ord Lt  (<)  (<)  (<)
  Geq    -> ord Geq (>=) (>=) (>=)
  Gt     -> ord Gt  (>)  (>)  (>)
  And    -> \e1 e2 -> case (Tuple e1 e2) of
    (Tuple (Atom (Bool false)) _                  ) -> return $ Atom $ Bool false
    (Tuple _                   (Atom (Bool false))) -> return $ Atom $ Bool false
    (Tuple (Atom (Bool true))  (Atom (Bool true)) ) -> return $ Atom $ Bool true
    (Tuple _                   _                  ) -> throwError $ BinaryOpError And e1 e2
  Or     -> \e1 e2 -> case (Tuple e1 e2) of
    (Tuple (Atom (Bool true))  _                   ) -> return $ Atom $ Bool true
    (Tuple _                   (Atom (Bool true))  ) -> return $ Atom $ Bool true
    (Tuple (Atom (Bool false))  (Atom (Bool false))) -> return $ Atom $ Bool false
    (Tuple _                   _                   ) -> throwError $ BinaryOpError And e1 e2
  Dollar -> (\f e -> return $ App f (singleton e))
  Composition -> \e1 e2 -> throwError $ BinaryOpError And e1 e2
  InfixFunc name -> \e1 e2 -> apply env name (e1 : e2 : Nil)
  where
    aint :: Op -> (Int -> Int -> Int) -> Expr -> Expr -> Evaluator Expr
    aint _   f (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ f i j
    aint op  _ e1              e2              = throwError $ BinaryOpError op e1 e2

    ord :: Op -> (Int -> Int -> Boolean) -> (String -> String -> Boolean) -> (Boolean -> Boolean -> Boolean)-> Expr -> Expr -> Evaluator Expr
    ord _  fi _  _  (Atom (AInt i))  (Atom (AInt j))  = return $ Atom $ Bool $ fi i j
    ord _  _  fs _  (Atom (Char s1)) (Atom (Char s2)) = return $ Atom $ Bool $ fs s1 s2
    ord _  _  _  fb (Atom (Bool b1)) (Atom (Bool b2)) = return $ Atom $ Bool $ fb b1 b2
    ord op _  _  _  e1               e2               = throwError $ BinaryOpError op e1 e2

unary :: Env -> Op -> Expr -> Evaluator Expr
unary env Sub (Atom (AInt i)) = return $ Atom $ AInt (-i)
unary env op e = throwError $ UnaryOpError op e

apply :: Env -> String -> (List Expr) -> Evaluator Expr
apply env "div" (Cons e1 (Cons e2 _)) = division e1 e2
apply env "mod" (Cons e1 (Cons e2 _)) = modulo e1 e2
apply env name args = case Map.lookup name env of
  Nothing    -> throwError $ UnknownFunction name
  Just cases -> tryAll env cases args name Nil

--builtin ipo
division :: Expr -> Expr -> Evaluator Expr
division (Atom (AInt i)) (Atom (AInt 0)) = throwError DivByZero
division (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ div i j
division e1 e2 = throwError $ BinaryOpError (InfixFunc "div") e1 e2

--builtin mod
modulo :: Expr -> Expr -> Evaluator Expr
modulo (Atom (AInt i)) (Atom (AInt 0)) = throwError DivByZero
modulo (Atom (AInt i)) (Atom (AInt j)) = return $ Atom $ AInt $ mod i j
modulo e1 e2 = throwError $ BinaryOpError (InfixFunc "mod") e1 e2

tryAll :: Env -> (List (Tuple (List Binding) Expr)) -> (List Expr) -> String -> (List MatchingError) -> Evaluator Expr
tryAll _   Nil                        _    name errs = throwError $ NoMatchingFunction name errs
tryAll env (Cons (Tuple binds body) cases) args name errs | (length args) < (length binds) = tryAll env cases args name (errs ++ (singleton $ TooFewArguments binds args))
tryAll env (Cons (Tuple binds body) cases) args name errs = case runMatcherM $ matchls' env binds args of
  Right replMap                 -> replace' replMap body >>= wrapLambda binds args
  Left se@(StrictnessError _ _) -> throwError $ NoMatchingFunction name (errs ++ singleton se)
  Left err                      -> tryAll env cases args name (errs ++ singleton err)

whnf :: Expr -> Boolean
whnf (Atom (Name _)) = false
whnf (Atom _)        = true
whnf (List _)        = true
whnf (NTuple _)      = true
whnf _               = false

checkStrictness :: Binding -> Expr -> MatchingError
checkStrictness bs es | whnf es   = MatchingError bs es
                      | otherwise = StrictnessError bs es

matchls' :: Env -> (List Binding) -> (List Expr) -> Matcher (StrMap Expr)
matchls' env bs es = execStateT (zipWithA (\b e -> match' b (evalToBinding env e b)) bs es) Map.empty

match' :: Binding -> Expr -> StateT (StrMap Expr) Matcher Unit
match' (Lit (Name name)) e                   = modify (Map.insert name e)
match' (Lit ba)          (Atom ea)           = case ba == ea of
                                                 true  -> return unit
                                                 false -> throwError $ MatchingError (Lit ba) (Atom ea)
match' (Lit b)           e                   = throwError $ checkStrictness (Lit b) e

match' (ConsLit b bs)    (Binary Colon e es)  = match' b e *> match' bs es
match' (ConsLit b bs)    (List (Cons e es))  = match' (ConsLit b bs) (Binary Colon e (List es))
match' (ConsLit b bs)    (List Nil)           = throwError $ MatchingError (ConsLit b bs) (List Nil)
match' (ConsLit b bs)    e                   = throwError $ checkStrictness (ConsLit b bs) e

match' (ListLit (Cons b bs))  (Binary Colon e es)  = match' b e *> match' (ListLit bs) es
match' (ListLit bs)      (List es)           = case length bs == length es of
                                                 true  -> void $ zipWithA match' bs es
                                                 false -> throwError $ MatchingError (ListLit bs) (List es)
match' (ListLit bs)      e                   = throwError $ checkStrictness (ListLit bs) e

match' (NTupleLit bs)    (NTuple es)         = case length bs == length es of
                                                 true  -> void $ zipWithA match' bs es
                                                 false -> throwError $ MatchingError (NTupleLit bs) (NTuple es)
match' (NTupleLit bs)    e                   = throwError $ checkStrictness (NTupleLit bs) e

mapM' :: forall a b m. (Monad m) => (a -> m b) -> Maybe a -> m (Maybe b)
mapM' f Nothing  = pure Nothing
mapM' f (Just x) = Just <$> (f x)

replace' :: StrMap Expr -> Expr -> Evaluator Expr
replace' subs = go
  where
  go expr = case expr of
    a@(Atom (Name name)) -> case Map.lookup name subs of
      Just subExpr -> return subExpr
      Nothing      -> return a
    (List exprs)         -> List <$> (traverse go exprs)
    (NTuple exprs)       -> NTuple <$> (traverse go exprs)
    (Binary op e1 e2)    -> Binary <$> pure op <*> go e1 <*> go e2
    (Unary op e)         -> Unary <$> pure op <*> go e
    (SectL e op)         -> SectL <$> go e <*> pure op
    (SectR op e)         -> SectR <$> pure op <*> go e
    (IfExpr ce te ee)    -> IfExpr <$> go ce <*> go te <*> go ee
    (ArithmSeq ce te ee) -> ArithmSeq <$> go ce <*> (mapM' go te) <*> (mapM' go ee)
    (Lambda binds body)  -> (avoidCapture subs binds) *> (Lambda <$> pure binds <*> replace' (foldr Map.delete subs (boundNames' binds)) body)
    (App func exprs)     -> App <$> go func <*> traverse go exprs
    e                    -> return e

avoidCapture :: StrMap Expr -> (List Binding) -> Evaluator Unit
avoidCapture subs binds = case intersect (concatMap freeVariables $ Map.values subs) (boundNames' binds) of
  Nil       -> return unit
  captures -> throwError $ NameCaptureError captures

freeVariables :: Expr -> (List String)
freeVariables _ = Nil
-- freeVariables = nub <<< foldExpr
--   (\a -> case a of
--     Name name -> singleton $ name
--     _         -> Nil)
--   concat
--   concat
--   (\_ f1 f2 -> f1 ++ f2)
--   (\_ f -> f)
--   (\f _ -> f)
--   (\_ f -> f)
--   (\_ -> Nil)
--   (\f1 f2 f3 -> f1 ++ f2 ++ f3)
--   (\bs f -> nub f \\ boundNames' bs)
--   (\f fs -> f ++ concat fs)

boundNames' :: (List Binding) -> (List String)
boundNames' = concatMap boundNames

boundNames :: Binding -> (List String)
boundNames = go
  where
  go (Lit (Name name)) = singleton name
  go (ConsLit b1 b2)   = go b1 ++ go b2
  go (ListLit bs)      = foldMap go bs
  go (NTupleLit bs)    = foldMap go bs
  go _                 = Nil


Bug fix. mixed the address and value fields in a call to wr16. This is what you get when you abstract away types using type and not newtype!

module Ops where

import Prelude
import Data.Int.Bits
import Data.Maybe
import Data.Array ((!!)) as A

import Types
import MainMem
import Utils

import Debug

--NOTE: consider moving the timings to 2D arrays.
--separate from the opcode's logic

-- Additions
-- =========

--ADD A,R
addRegToA :: GetReg -> Regs -> Regs
addRegToA getReg regs = addXToA (getReg regs) regs

--ADD A,(HL)
addHLMemToA :: Mem -> Regs
addHLMemToA { mainMem, regs } = (addXToA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--ADD A,n
addImmToA :: Mem -> Regs
addImmToA { mainMem, regs } =
  (addXToA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

addXToA :: I8 -> Regs -> Regs
addXToA x regs =
  regs {a = sum.res, f = sum.flags, m = 1}
 where sum = addI8s regs.a x

--ADC A,R
addRegCarryToA :: GetReg -> Regs -> Regs
addRegCarryToA getReg regs = addXCarryToA (getReg regs) regs

--ADC A,(HL)
addHLMemCarryToA :: Mem -> Regs
addHLMemCarryToA { mainMem, regs } = (addXCarryToA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--ADC A,n
addImmCarryToA :: Mem -> Regs
addImmCarryToA { mainMem, regs } =
  (addXCarryToA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Test edge case of carry flag of 1 that causes a half-carry,
--Is the half-carry flag set correctly?
addXCarryToA :: I8 -> Regs -> Regs
addXCarryToA x regs =
  regs {a = sum.res, f = f', m = 1}
 where
  --Overriding addI8's half-carry flag because it used regs.a+carry
  --Which is essential for the other flags, but not for this one.
  f' = setHalfCarryFlag x regs.a sum.res sum.flags
  sum = addI8s (regs.a + carry) x
  carry = if regs.f .&. carryFlag /= 0 then 1 else 0

--ADD SP,|n|
addImmToSP :: Mem -> Regs
addImmToSP { regs, mainMem } =
  regs {sp = sp', pc = regs.pc + 2, m = 4}
 where
  sp' = (absI8 imm + regs.sp) .&. 0xFFFF
  imm = rd8 (regs.pc+1) mainMem

addI8s :: I8 -> I8 -> { res :: I8, flags :: I8 }
addI8s x1 x2 = { res , flags }
 where
  flags =  testZeroFlag res
       .|. testCarryFlag8 sum
       .|. testHalfCarryFlag x1 x2 res
  res = sum .&. 255
  sum = x1 + x2

--ADD HL,RR
add2RegsToHL :: GetReg -> GetReg -> Regs -> Regs
add2RegsToHL msByteReg lsByteReg regs =
  regs { h = split.ms, l = split.ls, f = sum.flags, m = 3 }
 where
  split = splitI16 sum.res
  sum = addI16s hl joinedRegs regs.f
  joinedRegs = joinRegs msByteReg lsByteReg regs
  hl = joinRegs h l regs

--ADD HL,SP
addSPToHL :: Regs -> Regs
addSPToHL regs =
  regs { h = split.ms, l = split.ls, f = sum.flags, m = 3 }
 where
  split = splitI16 sum.res
  sum = addI16s hl regs.sp regs.f
  hl = joinRegs h l regs

addI16s :: I8 -> I16 -> I16 -> { res :: I16, flags :: I8 }
addI16s oldFlags x1 x2 = { res, flags }
 where
  flags =  setCarryFlag16 sum
        $  oldFlags
       .|. testZeroFlag res
  res = sum .&. 0xFFFF
  sum = x1 + x2

-- Subtractions
-- ============

--SUB A,R
subRegFromA :: GetReg -> Regs -> Regs
subRegFromA getReg regs = subXFromA (getReg regs) regs

--SUB A,(HL)
subHLMemFromA :: Mem -> Regs
subHLMemFromA { mainMem, regs } = (subXFromA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--SUB A,n
subImmFromA :: Mem -> Regs
subImmFromA { mainMem, regs } = (subXFromA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Subtractions and additions are relatively similar. Consider
--Using same functions for both, if it doesn't obfuscate much.
subXFromA :: I8 -> Regs -> Regs
subXFromA x regs =
  regs {a = diff.res, f = diff.flags, m = 1}
 where diff = subI8s regs.a x

--SBC A,R
subRegCarryFromA :: GetReg -> Regs -> Regs
subRegCarryFromA getReg regs = subXCarryFromA (getReg regs) regs

--SBC A,(HL)
subHLMemCarryFromA :: Mem -> Regs
subHLMemCarryFromA { mainMem, regs } = (subXCarryFromA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--SBC A,n
subImmCarryToA :: Mem -> Regs
subImmCarryToA { mainMem, regs } = (subXCarryFromA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Test edge case of carry flag of 1 that causes a half-carry,
--Is the half-carry flag set correctly?
subXCarryFromA :: I8 -> Regs -> Regs
subXCarryFromA x regs =
  regs {a = diff.res, f = f', m = 1}
 where
  f' = setHalfCarryFlag x regs.a diff.res diff.flags
  diff = subI8s (regs.a - carry) x
  carry = if regs.f .&. carryFlag /= 0 then 1 else 0

subI8s :: I8 -> I8 -> { res :: I8, flags :: I8 }
subI8s x1 x2 = { res , flags }
 where
  flags =  testZeroFlag res
       .|. testNegCarryFlag diff
       .|. subtractionFlag
       .|. testHalfCarryFlag x1 x2 res
  res = diff .&. 255
  diff = x1 - x2

-- Boolean operations
-- ==================

--AND A,R
andOpRegIntoA :: GetReg -> Regs -> Regs
andOpRegIntoA = boolOpRegIntoA (.&.)

--OR A,R
orOpRegIntoA :: GetReg -> Regs -> Regs
orOpRegIntoA = boolOpRegIntoA (.|.)

--XOR A,R
xorOpRegIntoA :: GetReg -> Regs -> Regs
xorOpRegIntoA = boolOpRegIntoA (.^.)

boolOpRegIntoA :: (I8 -> I8 -> I8) -> GetReg -> Regs -> Regs
boolOpRegIntoA op getReg regs = boolOpXIntoA op (getReg regs) regs

--AND A,(HL)
andOpHLMemIntoA :: Mem -> Regs
andOpHLMemIntoA = boolOpHlMemIntoA (.&.)

--OR A,(HL)
orOpHLMemIntoA :: Mem -> Regs
orOpHLMemIntoA = boolOpHlMemIntoA (.|.)

--XOR A,(HL)
xorOpHLMemIntoA :: Mem -> Regs
xorOpHLMemIntoA = boolOpHlMemIntoA (.^.)

boolOpHlMemIntoA  :: (I8 -> I8 -> I8) -> Mem -> Regs
boolOpHlMemIntoA op { mainMem, regs } =
  (boolOpXIntoA op hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--AND A,Imm
andOpImmIntoA :: Mem -> Regs
andOpImmIntoA  = boolOpImmIntoA (.&.)

--OR A,Imm
orOpImmIntoA :: Mem -> Regs
orOpImmIntoA = boolOpImmIntoA (.|.)

--XOR A,Imm
xorOpImmIntoA :: Mem -> Regs
xorOpImmIntoA = boolOpImmIntoA (.^.)

boolOpImmIntoA :: (I8 -> I8 -> I8) -> Mem -> Regs
boolOpImmIntoA op { mainMem, regs } =
  (boolOpXIntoA op imm regs) {pc = regs.pc + 2, m = 2}
 where imm = rd8 (regs.pc+1) mainMem

boolOpXIntoA :: (I8 -> I8 -> I8) -> I8 -> Regs -> Regs
boolOpXIntoA op x regs =
  regs { a = a', f = testZeroFlag a', m = 1 }
 where a' = (regs.a `op` x) .&. 255

-- Bit operations
-- ==============

--CPL
cmplA :: Regs -> Regs
cmplA regs = regs { a = a', f = testZeroFlag a' ,m = 1 }
 where a' = 255 .^. regs.a

--BIT N,R
testBitNOfReg :: I8 -> GetReg -> Regs -> Regs
testBitNOfReg n getReg regs = testBitNOfX n (getReg regs) regs

--BIT N,(HL)
testBitNOfHLMem :: I8 -> Mem -> Regs
testBitNOfHLMem n { mainMem, regs } =
  (testBitNOfX n hlMem regs) { m = 3 }
 where
  hlMem = rd8 (joinRegs h l regs) mainMem

testBitNOfX :: Int -> I8 -> Regs -> Regs
testBitNOfX n x regs = regs { f = f', m = 2 }
 where
  f' =  setZeroFlag bitTest
     $  halfCarryFlag
    .|. (regs.f .&. carryFlag) -- NOTE: should it be 0x1F instead?
  bitTest = x .&. (1 `shl` n)

--SET N,R
--RES N,R
setBitNOfReg :: Boolean -> Int -> SetReg -> GetReg
               -> Regs -> Regs
setBitNOfReg setVal n setReg getReg regs = setReg reg' regs { m = 2 }
 where reg' = setBitNOfX setVal n (getReg regs)

--SET N,(HL)
--RES N,(HL)
setBitNOfHLMem :: Boolean -> Int -> Mem -> Mem
setBitNOfHLMem setVal n mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { m = 4 } }
 where
  mainMem' = wr8 hl' addr mainMem
  hl' = setBitNOfX setVal n $ rd8 addr mainMem
  addr = joinRegs h l regs

setBitNOfX :: Boolean -> Int -> I8 -> I8
setBitNOfX setVal n x = if setVal
  then x .|. bitShift
  else x .&. complement bitShift
 where bitShift = 1 `shl` n

-- Rotations
-- =========

data Dir = Left | Right
derive instance eqDir :: Eq Dir

-- RL[C] A, base-opcode
-- RR[C] A, base-opcode
rotA :: Dir -> Boolean -> Regs -> Regs
rotA dir isCarryRot regs = regs { a = rotated.res, f = rotated.flags, m = 1 }
 where rotated = rotX {dir,isCarryRot,isCB:false} regs.a regs.f

-- RL[C] R, CB-prefix instruction
-- RR[C] R, CB-prefix instruction
rotReg :: Dir -> Boolean -> SetReg -> GetReg -> Regs -> Regs
rotReg dir isCarryRot setReg getReg regs =
  setReg rotated.res regs { f = rotated.flags, m = 2 }
 where rotated = rotX {dir,isCarryRot,isCB:true} (getReg regs) regs.f

-- RL[C] (HL)
-- RR[C] (HL)
rotHLMem :: Dir -> Boolean -> Mem -> Mem
rotHLMem dir isCarryRot mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { f = rotated.flags, m = 4 } }
 where
  mainMem' = wr8 rotated.res addr mainMem
  addr = joinRegs h l regs
  rotated = rotX {dir,isCarryRot,isCB:true} (rd8 addr mainMem) regs.f

rotX :: {dir::Dir,isCarryRot::Boolean,isCB::Boolean}
      -> I8 -> I8 -> { res :: I8, flags :: I8 }
rotX { dir, isCarryRot, isCB } x oldFlags = { res, flags }
 where
  res = 255 .&. (currCarry + (x `shiftFunc` 1))
  -- RotA basic, use old flags, RotR CB, start fresh with zeroFlag
  flags = setNewCarry
    $ if isCB then testZeroFlag res else oldFlags
  currCarry = if   (not isCarryRot) && (oldFlags .&. carryFlag /= 0)
                || (isCarryRot      && isNewCarry)
    then case dir of
      Left -> 1
      Right -> 0x80
    else 0
  setNewCarry = if isNewCarry
    then (carryFlag .|. _)
    else (cmplCarryFlag .&. _)
  isNewCarry = x .&. edgeBit /= 0
  edgeBit = case dir of
    Left -> 0x80
    Right -> 1
  shiftFunc = case dir of
    Left -> shl
    Right -> zshr

-- Shifts
-- ======

--SLA R
--SRA R preserve sign
--SRL R
shiftReg :: Dir -> Boolean -> SetReg -> GetReg
         -> Regs -> Regs
shiftReg dir sign setReg getReg regs =
  setReg shifted.res regs { f = shifted.flags, m = 2 }
 where
  shifted = shiftX dir sign $ getReg regs

shiftMemHL :: Dir -> Boolean -> Mem -> Mem
shiftMemHL dir sign mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { f = shifted.flags, m = 4 } }
 where
  mainMem' = wr8 shifted.res addr mainMem
  shifted = shiftX dir sign $ rd8 addr mainMem
  addr = joinRegs h l regs

shiftX :: Dir -> Boolean -> I8
       -> { res::I8, flags::I8 }
shiftX dir sign x = { res, flags }
 where
  res =  (255 .&. _)
     <<< (if sign && dir == Right then (currCarry + _) else id)
      $  x `shiftFunc` 1
  shiftFunc = case dir of
    Left -> shl
    Right -> zshr
  flags = newCarry .|. testZeroFlag res
  newCarry = if x .&. edgeBit /= 0 then carryFlag else 0
  currCarry = x .&. 0x80 --sign is only relevant for right shifts
  edgeBit = case dir of
    Left -> 0x80
    Right -> 1

-- Increments / Decrements
-- =======================

--INC R
incReg :: GetReg -> SetReg -> Regs -> Regs
incReg = incDecReg incI8

--DEC R
decReg :: GetReg -> SetReg -> Regs -> Regs
decReg getReg setReg regs = regs' { f = setFlag subtractionFlag regs'.f }
  where regs' = (incDecReg decI8 getReg setReg regs)

incDecReg :: (I8 -> { res :: I8, carry :: Boolean })
          -> GetReg -> SetReg -> Regs -> Regs
incDecReg op getReg setReg regs =
  setReg inced .res regs {f = testZeroFlag inced.res, m = 1}
 where
  inced  = op $ getReg regs

--INC RR
incRegWithCarry :: GetReg -> GetReg -> SetReg -> SetReg
                -> Regs -> Regs
incRegWithCarry = incDecRegWithCarry incI8

--DEC RR
decRegWithCarry :: GetReg -> GetReg -> SetReg -> SetReg
                -> Regs -> Regs
decRegWithCarry = incDecRegWithCarry decI8

--NOTE: consider refactoring to something along the lines of
--"split2Regs . incI16 . joinRegs", if you have leeway, performance-wise.
incDecRegWithCarry :: (I8 -> { res :: I8, carry :: Boolean })
          -> GetReg -> GetReg -> SetReg -> SetReg
          -> Regs -> Regs
incDecRegWithCarry op getCarryReg getMainReg setCarryReg setMainReg regs =
  setCarryReg carryReg' <<< setMainReg mainReg' $ regs{m = 1}
 where
  carryReg' = (if mainRegInc.carry then _.res <<< op else id)
    $ getCarryReg regs
  mainReg' = mainRegInc.res
  mainRegInc = op $ getMainReg regs

--INC SP
incSP :: Regs -> Regs
incSP = incDecSP $ \x -> 0xFFFF .&. (x + 1)

--DEC SP
decSP :: Regs -> Regs
decSP = incDecSP $ \x -> 0xFFFF .&. (x - 1)

incDecSP :: (I16 -> I16) -> Regs -> Regs
incDecSP op regs = regs {sp = op $ regs.sp, m = 1}

--INC (HL)
incHLMem :: Mem -> Mem
incHLMem = incDecHLMem incI8

--DEC (HL)
decHLMem :: Mem -> Mem
decHLMem = incDecHLMem decI8

incDecHLMem :: (I8 -> { res :: I8, carry :: Boolean })
            -> Mem -> Mem
incDecHLMem op mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs {f = testZeroFlag inced.res, m = 3} }
 where
  mainMem' = wr8 inced.res addr mainMem
  inced = op $ rd8 addr mainMem
  addr = joinRegs h l regs

incI8 :: I8 -> { res :: I8, carry :: Boolean }
incI8 i8 = { res : i8', carry : i8' == 0 }
 where
  i8' = (i8 + 1) .&. 255

decI8 :: I8 -> { res :: I8, carry :: Boolean }
decI8 i8 = { res : i8', carry : i8' == 255 }
 where
  i8' = (i8 - 1) .&. 255

-- Loads
-- =====

--LD R,R
ldRegFromReg :: SetReg -> GetReg
      -> Regs -> Regs
ldRegFromReg setDestReg getSrcReg regs = setDestReg (getSrcReg regs)
  regs { m = 1 }

--LD (nn),SP
ldMemImmFromSP :: Mem -> Mem
ldMemImmFromSP mem@{mainMem, regs} =
  mem { mainMem = mainMem', regs = regs { pc = regs.pc + 3, m = 5 } }
 where
  mainMem' = wr16 regs.sp addr mainMem
  addr = rd16 (regs.pc+1) mainMem

--LD SP,HL
ldSPFromHL :: Regs -> Regs
ldSPFromHL regs  =
  regs { sp = sp', m = 2 }
 where sp' = joinRegs h l regs

ldSPFromImm :: Mem -> Regs
ldSPFromImm { mainMem, regs } =
  regs { sp = sp', pc = regs.pc + 3, m = 3 }
 where sp' = rd16 (regs.pc+1) mainMem

--LD RR,nn
ldTwoRegsFromImm :: SetReg -> SetReg
               -> Mem -> Regs
ldTwoRegsFromImm setReg1 setReg2 mem = regs'' { m = 3 }
 where
  regs'  = ldRegFromImm setReg2 mem
  regs'' = ldRegFromImm setReg1 mem { regs = regs' }

--LD R,(IOC)
ldRegFromFF00CMem :: SetReg -> Mem -> Regs
ldRegFromFF00CMem setReg mem@{regs, mainMem} = regs' { m = 2 }
 where regs' = ldRegFromFF00PlusX regs.c setReg mem

--LD R,(IOn)
ldRegFromFF00ImmMem :: SetReg -> Mem -> Regs
ldRegFromFF00ImmMem setReg mem@{regs, mainMem} = regs' { pc = regs.pc + 2, m = 3 }
 where regs' = ldRegFromFF00PlusX (rd8 (regs.pc+1) mainMem) setReg mem

ldRegFromFF00PlusX :: I8 -> SetReg -> Mem -> Regs
ldRegFromFF00PlusX  x setReg mem@{regs} = ldRegFromMem setReg addr mem
 where addr = 0xFF00 + x

--LD R,n
ldRegFromImm :: SetReg -> Mem -> Regs
ldRegFromImm setReg mem@{regs} = regs' { pc = regs.pc + 2 }
  where regs' = ldRegFromMem setReg (regs.pc+1) mem

--LD R,(RR)
ldRegFromMem2R :: SetReg -> GetReg -> GetReg
               -> Mem -> Regs
ldRegFromMem2R setReg msByteReg lsByteReg mem@{regs} =
  ldRegFromMem setReg ((msByteReg regs `shl` 8) + lsByteReg regs) mem

--LD A,(nn)
ldRegAFromMemImm :: Mem -> Regs
ldRegAFromMemImm mem@{ mainMem, regs } = regs' { pc = regs.pc + 3, m = 4 }
 where
  regs' = ldRegFromMem setA addr mem
  addr = rd16 (regs.pc+1) mainMem

--LDD A,(HL)
ldRegFromMemHLDec :: SetReg -> Mem -> Regs
ldRegFromMemHLDec = ldRegFromMemHLIncDec decRegWithCarry

--LDI A,(HL)
ldRegFromMemHLInc :: SetReg -> Mem -> Regs
ldRegFromMemHLInc = ldRegFromMemHLIncDec incRegWithCarry

ldRegFromMemHLIncDec :: (GetReg -> GetReg -> SetReg -> SetReg -> Regs -> Regs)
                   -- Inc/Dec reg operation's type signature
                   -> SetReg -> Mem -> Regs
ldRegFromMemHLIncDec op setReg mem@{ regs } = regs' {m = 2}
 where
  regs' = op h l setH setL
        $ ldRegFromMem setA addr mem
  addr = joinRegs h l regs

ldRegFromMem :: SetReg -> I16 -> Mem -> Regs
ldRegFromMem setReg addr { mainMem, regs } =
  setReg (rd8 addr mainMem) regs{m = 2}

--LD HL,SP|n|
ldHLFromSPImm :: Mem -> Regs
ldHLFromSPImm  mem@{ regs, mainMem } =
  regs {h = split.ms, l = split.ls, pc = regs.pc + 2, m = 3}
 where
  split = splitI16 x
  x = (absI8 imm + regs.sp) .&. 0xFFFF
  imm = rd8 (regs.pc+1) mainMem

--LDI (HL),A
ldMemHLFromRegInc :: GetReg -> Mem -> Mem
ldMemHLFromRegInc = ldMemHLFromRegIncDec incRegWithCarry

--LDD (HL),A
ldMemHLFromRegDec :: GetReg -> Mem -> Mem
ldMemHLFromRegDec  = ldMemHLFromRegIncDec decRegWithCarry

ldMemHLFromRegIncDec :: (GetReg -> GetReg -> SetReg -> SetReg -> Regs -> Regs)
                   -- Inc/Dec reg operation's type signature
                   -> GetReg
                   -> Mem -> Mem
ldMemHLFromRegIncDec op getReg mem@{regs} =
  mem { mainMem = mainMem', regs = regs' {m = 2} }
 where
  regs' = op h l setH setL regs
  mainMem' = _.mainMem $ ldMemFromReg addr a mem
  addr = joinRegs h l regs

--LD (IOC),R
ldFF00CMemFromReg :: GetReg -> Mem -> Mem
ldFF00CMemFromReg getReg mem@{regs, mainMem} =
  mem' { regs = regs{m = 2} }
 where mem' = ldFF00PlusXFromReg  regs.c getReg mem

--LD (IOn),R
ldFF00ImmMemFromReg :: GetReg -> Mem -> Mem
ldFF00ImmMemFromReg getReg mem@{regs, mainMem} =
  mem' { regs = regs{pc = regs.pc + 2, m = 3} }
 where
  mem' = ldFF00PlusXFromReg imm getReg mem
  imm = rd8 (regs.pc+1) mainMem

ldFF00PlusXFromReg :: I8 -> GetReg -> Mem -> Mem
ldFF00PlusXFromReg   x getReg mem@{regs, mainMem} = mem'
 where
  mem' = ldMemFromReg addr getReg mem
  addr = 0xFF00 + x

--LD (nn),A
ldMemImmFromRegA :: Mem -> Mem
ldMemImmFromRegA  mem@{ mainMem, regs } =
 mem' { regs = regs{m = 4, pc = regs.pc + 3} }
 where
  mem' = ldMemFromReg addr a mem
  addr = rd16 (regs.pc+1) mainMem

--LD (HL),n
ldMemHLFromImm :: Mem -> Mem
ldMemHLFromImm  mem@{mainMem,regs} =
  mem { mainMem = mainMem' , regs = regs{m = 3, pc = regs.pc + 2} }
 where
  mainMem' = wr8 imm addr mainMem
  addr =  joinRegs h l regs
  imm = rd8 (regs.pc+1) mainMem

--LD (RR),R
ldMem2RFromReg :: GetReg -> GetReg -> GetReg
                        -> Mem -> Mem
ldMem2RFromReg msbAddrReg lsbAddrReg getReg mem@{ regs } =
  ldMemFromReg (joinRegs msbAddrReg lsbAddrReg regs) getReg mem

ldMemFromReg :: I16 -> GetReg
           -> Mem -> Mem
ldMemFromReg addr getReg mem@{ regs } = ldMemFromX addr (getReg regs) mem

ldMemFromX :: I16 -> I8
           -> Mem -> Mem
ldMemFromX addr x mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs{m = 2} }
 where mainMem' = wr8 x addr mainMem

-- Compares
-- ========

--CP A,R
compAToReg :: GetReg -> Regs -> Regs
compAToReg getReg regs = compAToX (getReg regs) regs

--CP A,(HL)
compAToMemHL :: Mem -> Regs
compAToMemHL { mainMem, regs } = (compAToX hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--CP A,n
compAToImm :: Mem -> Regs
compAToImm { mainMem, regs } = (compAToX imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

compAToX :: I8 -> Regs -> Regs
compAToX x regs = regs { f = diff.flags, m = 1 }
 where
  diff = subI8s regs.a x

-- Stack operations
-- ================

--NOTE: On the one hand, there's an advantage of letting the compiler catch
--more mistakes by using row polymorphism in this case,
--on the other hand it obfuscates the signature, is there a better way?
--forall r s. { regs :: {sp::I8 | s}, mem :: MainMem | r }
--         -> { regs :: {sp::I8 | s}, mem :: MainMem | r }
--NOTE: Only 4 versions of PUSH so consider hard-coding them for performance.
--PUSH RR
pushReg :: GetReg -> GetReg
        -> Mem -> Mem
pushReg msByteReg lsByteReg mem@{regs,mainMem} =
  mem { mainMem = mainMem', regs = regs { sp = regs.sp - 2, m = 3 } }
 where
  mainMem' =  wr8 (lsByteReg regs) (regs.sp - 2)
          <<< wr8 (msByteReg regs) (regs.sp - 1)
           $  mainMem

--POP RR
popReg :: SetReg -> SetReg
       -> Mem -> Regs
popReg setMsByteReg setLsByteReg { mainMem, regs = regs@{sp} }
    = setMsByteReg (rd8 (sp + 1) mainMem)
   <<< setLsByteReg (rd8 sp mainMem)
    $  regs { sp = sp + 2, m = 3 }

-- Jumps
-- =====

-- JR NZ,n
-- JR Z,n
-- JR NC,n
-- JR C,n
jumpRelImmFlag :: Boolean -> I8 -> Mem -> Regs
jumpRelImmFlag inverse flag mem@{mainMem, regs} =
  if inverse `xor` isSetFlag flag regs.f
    then jumpRelImm mem
    else regs { pc = regs.pc + 2, m = 2 }

{----JR n--}
{--jumpRelImm :: Mem -> Regs--}
{--jumpRelImm { mainMem, regs } =--}
  {--regs { pc = pc', m = 3 }--}
 {--where--}
  {----NOTE: is it necessary to abs the immediate?--}
  {--pc' = regs.pc + 1 + absI8 imm--}
  {--imm = rd8 regs.pc mainMem--}

--JR n
jumpRelImm :: Mem -> Regs
jumpRelImm { mainMem, regs } =
  regs { pc = pc', m = 3 }
 where
  --NOTE: is it necessary to abs the immediate?
  pc' = regs.pc + 2 + absI8 imm
  imm = rd8 (regs.pc+1) mainMem

-- JP NZ,nn
-- JP Z,nn
-- JP NC,nn
-- JP C,nn
jumpImmFlag :: Boolean -> I8 -> Mem -> Regs
jumpImmFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then (jumpImm mem) { m = 4 }
    else regs { pc = regs.pc + 3, m = 3 }

-- JP nn
jumpImm :: Mem -> Regs
jumpImm { mainMem, regs } =
  regs { pc = pc', m = 3 }
 where pc' = rd16 (regs.pc+1) mainMem

-- JP (HL)
jumpHL :: Regs -> Regs
jumpHL regs = regs { pc = pc', m = 1 }
 where pc' = joinRegs h l regs

{--jumpX :: I8 -> Regs--}
{--jumpX { mainMem, regs } =--}
  {--regs { pc = pc', m = 3 }--}
 {--where pc' = rd16 regs.pc mainMem--}

-- Call Z,nn
-- Call NZ,nn
-- Call C,nn
-- Call NC,nn
callImmFlag :: Boolean -> I8 -> Mem -> Mem
callImmFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then callImm mem
    else mem { regs = regs {pc = regs.pc + 3, m = 3} }

-- Call nn
callImm :: Mem -> Mem
callImm mem@{mainMem,regs} =
  mem { mainMem =  mainMem', regs = regs { pc = imm, sp = regs.sp - 2, m = 5 } }
 where
  mainMem' = wr16 (regs.pc + 3) (regs.sp-2) $ mainMem
  imm = rd16 (regs.pc+1) mainMem

-- RET Z
-- RET NZ
-- RET C
-- RET NC
retFlag :: Boolean -> I8 -> Mem -> Regs
retFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then ret mem
    else regs { m = 1 }

-- RETI
retEnableInterrupt :: Mem -> Regs
retEnableInterrupt { mainMem, regs, svdRegs } =
  restoreRegs svdRegs regs { ime = true, pc = pc', sp = regs.sp + 2, m = 3 }
 where pc' = rd16 regs.sp mainMem

-- RET
ret :: Mem -> Regs
ret { mainMem, regs } =
  regs { pc = pc', sp = regs.sp + 2, m = 3 }
 where pc' = rd16 regs.sp mainMem

--RST XX
callRoutine :: I16 -> Mem -> Mem
callRoutine addr { mainMem, regs, svdRegs } =
  { mainMem : mainMem', regs : regs', svdRegs : saveRegs regs svdRegs }
 where
  regs' = regs { sp = regs.sp - 2, pc = addr, m = 3 }
  mainMem' = wr16 regs.pc (regs.sp - 2) mainMem

-- Misc.
-- =====

--SWAP R
swapReg :: SetReg -> GetReg -> Regs -> Regs
swapReg setReg getReg regs = setReg reg' $ regs { m = 1 }
 where
  reg' = ((0x0F.&.reg) `shl` 4) .|. ((0xF0.&.reg) `zshr` 4)
  reg = getReg regs

--SWAP (HL)
swapMemHL :: Mem -> Mem
swapMemHL mem@{mainMem, regs} =
  mem { mainMem = mainMem', regs = regs { m = 4 } }
 where
  mainMem' = wr8 hlMem' addr mainMem
  hlMem' = ((0x0F.&.hlMem) `shl` 4) .|. ((0xF0.&.hlMem) `zshr` 4)
  hlMem = rd8 addr mainMem
  addr = joinRegs h l regs

--Why a command called 'clear' actually flips the tag
--and not unset it, is a question that shouldn't be directed at me!
--CCF
clearCarryFlag :: Regs -> Regs
clearCarryFlag = changeCarryFlag false

--SCF
setCarryFlag :: Regs -> Regs
setCarryFlag = changeCarryFlag true

changeCarryFlag :: Boolean -> Regs -> Regs
changeCarryFlag isSet regs = regs { f = f', m = 1 }
 where f' =  unsetFlag subtractionFlag
         <<< unsetFlag halfCarryFlag
         <<< (if isSet then setFlag else flipFlag) carryFlag
          $  regs.f

--DAA
adjAForBCDAdd :: Regs -> Regs
adjAForBCDAdd regs = regs { f = adjF regs.f, a = a', m = 1 }
 where
  a' = adjHighNyb <<< adjLowNyb $ regs.a
  adjLowNyb = if (regs.f .&. halfCarryFlag /= 0) || (regs.a .&. 15 > 9)
    then (6 + _) else id
  adjHighNyb = if shouldAdjHighNyb then (0x60 + _) else id
  adjF = if shouldAdjHighNyb then (carryFlag .|. _) else (cmplCarryFlag .&. _)
  shouldAdjHighNyb = (regs.f .&. halfCarryFlag /= 0) || (regs.a > 0x99)

--NOP
nop :: Regs -> Regs
nop = _ { m = 1}

--Invalid opCode
invalidOpCode :: Z80State -> Z80State
--NOTE: LOG THIS! for debugging purposes
invalidOpCode state = state { stop = true }

--STOP
--NOTE: this is just a guess
stop :: Z80State -> Z80State
stop state@{mem=mem@{regs}} =
  state { stop = true
        , mem = mem { regs = regs { m = 1 } }
        }

--HALT
halt :: Z80State -> Z80State
halt state@{mem=mem@{regs}} =
  state { halt = true
        , mem = mem { regs = regs { m = 1 } }
        }

--DI
--EI
setInterrupts :: Boolean -> Regs -> Regs
setInterrupts enable = _ { ime = enable, m = 1 }

--Extended Ops
execExtOps :: Array (Z80State -> Z80State)
           -> Z80State -> Z80State
execExtOps opsMap state@{ mem = mem@{mainMem,regs} } = extOp state'
 where
  --NOTE replace toMaybe, with something to log errors
  extOp = getOpcode opCode opsMap
  state' = state { mem = mem { regs = regs { pc = 65535 .&. (regs.pc + 1) } } }
  opCode = rd8 regs.pc mainMem

-- Helpers
-- =======

saveRegs :: Regs -> SavedRegs -> SavedRegs
saveRegs regs svdRegs =
  svdRegs { a = regs.a, b = regs.b, c = regs.c, d = regs.d
          , e = regs.e, f = regs.f, h = regs.h, l = regs.l
          }

restoreRegs :: SavedRegs -> Regs -> Regs
restoreRegs svdRegs regs =
  regs { a = svdRegs.a, b = svdRegs.b, c = svdRegs.c, d = svdRegs.d
       , e = svdRegs.e, f = svdRegs.f, h = svdRegs.h, l = svdRegs.l
       }

joinRegs :: GetReg -> GetReg -> Regs -> I16
joinRegs msByteReg lsByteReg regs = (msByteReg regs `shl` 8) + lsByteReg regs

absI8 :: I8 -> I8
absI8 i8 = if i8 > 127 then negI8 i8 else i8

negI8 :: I8 -> I8
negI8 i8 = (complement i8 + 1) .&. 255

splitI16 :: I16 -> { ms :: I8, ls :: I8 }
splitI16 x = { ms : (x `zshr` 8) .&. 255
             , ls : x .&. 255
             }

isSetFlag :: I8 -> I8 -> Boolean
isSetFlag f fs = f .&. fs /= 0

setFlag :: I8 -> I8 -> I8
setFlag  f = (f .|. _)

unsetFlag :: I8 -> I8 -> I8
unsetFlag f = ((complement f) .&. _)

flipFlag :: I8 -> I8 -> I8
flipFlag f = (f .^. _)

testZeroFlag :: Int -> Int
testZeroFlag x = if x == 0 then 0x80 else 0

setZeroFlag :: Int -> I8 -> I8
setZeroFlag x = if x == 0
  then (_ .|. zeroFlag)
  else (_ .&. cmplZeroFlag)

testCarryFlag8 :: Int -> I8
testCarryFlag8 x = if x > 255 then carryFlag else 0

setCarryFlag16 :: Int -> I8 -> I8
setCarryFlag16 x = if x > 0xFFFF
  then (_ .|. carryFlag)
  else (_ .&. cmplCarryFlag)

testNegCarryFlag :: Int -> I8
testNegCarryFlag reg = if reg < 0 then carryFlag else 0

testHalfCarryFlag :: I8 -> I8 -> I8 -> I8
testHalfCarryFlag reg1 reg2 sum = if res /= 0 then halfCarryFlag else 0
 where res = 0x10 .&. (reg1 .^. reg2 .^. sum)

setHalfCarryFlag :: I8 -> I8 -> I8 -> I8 -> I8
setHalfCarryFlag reg1 reg2 sum =
  if res /= 0 then (_ .|. halfCarryFlag) else (_ .&. cmplHalfCarryFlag)
 where res = 0x10 .&. (reg1 .^. reg2 .^. sum)

--NOTE: replace fromMaybe with an error report mechanism to trace bad cases
getOpcode :: Int -> Array (Z80State -> Z80State) -> Z80State -> Z80State
getOpcode ix table = fromMaybe id $ table A.!! ix

xor :: Boolean -> Boolean -> Boolean
xor true x = not x
xor false x = x

zeroFlag :: I8
zeroFlag = 0x80
cmplZeroFlag :: I8
cmplZeroFlag = 0x7F
carryFlag :: I8
carryFlag = 0x10
cmplCarryFlag :: I8
cmplCarryFlag = 0xEF
halfCarryFlag :: I8
halfCarryFlag = 0x20
cmplHalfCarryFlag :: I8
cmplHalfCarryFlag = 0xD0
subtractionFlag :: I8
subtractionFlag = 0x40

type GetReg = Regs -> I8
type SetReg = I8 -> Regs -> Regs

module Ops where

import Prelude
import Data.Int.Bits
import Data.Maybe
import Data.Array ((!!)) as A

import Types
import MainMem
import Utils

import Debug

--NOTE: consider moving the timings to 2D arrays.
--separate from the opcode's logic

-- Additions
-- =========

--ADD A,R
addRegToA :: GetReg -> Regs -> Regs
addRegToA getReg regs = addXToA (getReg regs) regs

--ADD A,(HL)
addHLMemToA :: Mem -> Regs
addHLMemToA { mainMem, regs } = (addXToA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--ADD A,n
addImmToA :: Mem -> Regs
addImmToA { mainMem, regs } =
  (addXToA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

addXToA :: I8 -> Regs -> Regs
addXToA x regs =
  regs {a = sum.res, f = sum.flags, m = 1}
 where sum = addI8s regs.a x

--ADC A,R
addRegCarryToA :: GetReg -> Regs -> Regs
addRegCarryToA getReg regs = addXCarryToA (getReg regs) regs

--ADC A,(HL)
addHLMemCarryToA :: Mem -> Regs
addHLMemCarryToA { mainMem, regs } = (addXCarryToA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--ADC A,n
addImmCarryToA :: Mem -> Regs
addImmCarryToA { mainMem, regs } =
  (addXCarryToA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Test edge case of carry flag of 1 that causes a half-carry,
--Is the half-carry flag set correctly?
addXCarryToA :: I8 -> Regs -> Regs
addXCarryToA x regs =
  regs {a = sum.res, f = f', m = 1}
 where
  --Overriding addI8's half-carry flag because it used regs.a+carry
  --Which is essential for the other flags, but not for this one.
  f' = setHalfCarryFlag x regs.a sum.res sum.flags
  sum = addI8s (regs.a + carry) x
  carry = if regs.f .&. carryFlag /= 0 then 1 else 0

--ADD SP,|n|
addImmToSP :: Mem -> Regs
addImmToSP { regs, mainMem } =
  regs {sp = sp', pc = regs.pc + 2, m = 4}
 where
  sp' = (absI8 imm + regs.sp) .&. 0xFFFF
  imm = rd8 (regs.pc+1) mainMem

addI8s :: I8 -> I8 -> { res :: I8, flags :: I8 }
addI8s x1 x2 = { res , flags }
 where
  flags =  testZeroFlag res
       .|. testCarryFlag8 sum
       .|. testHalfCarryFlag x1 x2 res
  res = sum .&. 255
  sum = x1 + x2

--ADD HL,RR
add2RegsToHL :: GetReg -> GetReg -> Regs -> Regs
add2RegsToHL msByteReg lsByteReg regs =
  regs { h = split.ms, l = split.ls, f = sum.flags, m = 3 }
 where
  split = splitI16 sum.res
  sum = addI16s hl joinedRegs regs.f
  joinedRegs = joinRegs msByteReg lsByteReg regs
  hl = joinRegs h l regs

--ADD HL,SP
addSPToHL :: Regs -> Regs
addSPToHL regs =
  regs { h = split.ms, l = split.ls, f = sum.flags, m = 3 }
 where
  split = splitI16 sum.res
  sum = addI16s hl regs.sp regs.f
  hl = joinRegs h l regs

addI16s :: I8 -> I16 -> I16 -> { res :: I16, flags :: I8 }
addI16s oldFlags x1 x2 = { res, flags }
 where
  flags =  setCarryFlag16 sum
        $  oldFlags
       .|. testZeroFlag res
  res = sum .&. 0xFFFF
  sum = x1 + x2

-- Subtractions
-- ============

--SUB A,R
subRegFromA :: GetReg -> Regs -> Regs
subRegFromA getReg regs = subXFromA (getReg regs) regs

--SUB A,(HL)
subHLMemFromA :: Mem -> Regs
subHLMemFromA { mainMem, regs } = (subXFromA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--SUB A,n
subImmFromA :: Mem -> Regs
subImmFromA { mainMem, regs } = (subXFromA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Subtractions and additions are relatively similar. Consider
--Using same functions for both, if it doesn't obfuscate much.
subXFromA :: I8 -> Regs -> Regs
subXFromA x regs =
  regs {a = diff.res, f = diff.flags, m = 1}
 where diff = subI8s regs.a x

--SBC A,R
subRegCarryFromA :: GetReg -> Regs -> Regs
subRegCarryFromA getReg regs = subXCarryFromA (getReg regs) regs

--SBC A,(HL)
subHLMemCarryFromA :: Mem -> Regs
subHLMemCarryFromA { mainMem, regs } = (subXCarryFromA hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--SBC A,n
subImmCarryToA :: Mem -> Regs
subImmCarryToA { mainMem, regs } = (subXCarryFromA imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

--NOTE: Test edge case of carry flag of 1 that causes a half-carry,
--Is the half-carry flag set correctly?
subXCarryFromA :: I8 -> Regs -> Regs
subXCarryFromA x regs =
  regs {a = diff.res, f = f', m = 1}
 where
  f' = setHalfCarryFlag x regs.a diff.res diff.flags
  diff = subI8s (regs.a - carry) x
  carry = if regs.f .&. carryFlag /= 0 then 1 else 0

subI8s :: I8 -> I8 -> { res :: I8, flags :: I8 }
subI8s x1 x2 = { res , flags }
 where
  flags =  testZeroFlag res
       .|. testNegCarryFlag diff
       .|. subtractionFlag
       .|. testHalfCarryFlag x1 x2 res
  res = diff .&. 255
  diff = x1 - x2

-- Boolean operations
-- ==================

--AND A,R
andOpRegIntoA :: GetReg -> Regs -> Regs
andOpRegIntoA = boolOpRegIntoA (.&.)

--OR A,R
orOpRegIntoA :: GetReg -> Regs -> Regs
orOpRegIntoA = boolOpRegIntoA (.|.)

--XOR A,R
xorOpRegIntoA :: GetReg -> Regs -> Regs
xorOpRegIntoA = boolOpRegIntoA (.^.)

boolOpRegIntoA :: (I8 -> I8 -> I8) -> GetReg -> Regs -> Regs
boolOpRegIntoA op getReg regs = boolOpXIntoA op (getReg regs) regs

--AND A,(HL)
andOpHLMemIntoA :: Mem -> Regs
andOpHLMemIntoA = boolOpHlMemIntoA (.&.)

--OR A,(HL)
orOpHLMemIntoA :: Mem -> Regs
orOpHLMemIntoA = boolOpHlMemIntoA (.|.)

--XOR A,(HL)
xorOpHLMemIntoA :: Mem -> Regs
xorOpHLMemIntoA = boolOpHlMemIntoA (.^.)

boolOpHlMemIntoA  :: (I8 -> I8 -> I8) -> Mem -> Regs
boolOpHlMemIntoA op { mainMem, regs } =
  (boolOpXIntoA op hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--AND A,Imm
andOpImmIntoA :: Mem -> Regs
andOpImmIntoA  = boolOpImmIntoA (.&.)

--OR A,Imm
orOpImmIntoA :: Mem -> Regs
orOpImmIntoA = boolOpImmIntoA (.|.)

--XOR A,Imm
xorOpImmIntoA :: Mem -> Regs
xorOpImmIntoA = boolOpImmIntoA (.^.)

boolOpImmIntoA :: (I8 -> I8 -> I8) -> Mem -> Regs
boolOpImmIntoA op { mainMem, regs } =
  (boolOpXIntoA op imm regs) {pc = regs.pc + 2, m = 2}
 where imm = rd8 (regs.pc+1) mainMem

boolOpXIntoA :: (I8 -> I8 -> I8) -> I8 -> Regs -> Regs
boolOpXIntoA op x regs =
  regs { a = a', f = testZeroFlag a', m = 1 }
 where a' = (regs.a `op` x) .&. 255

-- Bit operations
-- ==============

--CPL
cmplA :: Regs -> Regs
cmplA regs = regs { a = a', f = testZeroFlag a' ,m = 1 }
 where a' = 255 .^. regs.a

--BIT N,R
testBitNOfReg :: I8 -> GetReg -> Regs -> Regs
testBitNOfReg n getReg regs = testBitNOfX n (getReg regs) regs

--BIT N,(HL)
testBitNOfHLMem :: I8 -> Mem -> Regs
testBitNOfHLMem n { mainMem, regs } =
  (testBitNOfX n hlMem regs) { m = 3 }
 where
  hlMem = rd8 (joinRegs h l regs) mainMem

testBitNOfX :: Int -> I8 -> Regs -> Regs
testBitNOfX n x regs = regs { f = f', m = 2 }
 where
  f' =  setZeroFlag bitTest
     $  halfCarryFlag
    .|. (regs.f .&. carryFlag) -- NOTE: should it be 0x1F instead?
  bitTest = x .&. (1 `shl` n)

--SET N,R
--RES N,R
setBitNOfReg :: Boolean -> Int -> SetReg -> GetReg
               -> Regs -> Regs
setBitNOfReg setVal n setReg getReg regs = setReg reg' regs { m = 2 }
 where reg' = setBitNOfX setVal n (getReg regs)

--SET N,(HL)
--RES N,(HL)
setBitNOfHLMem :: Boolean -> Int -> Mem -> Mem
setBitNOfHLMem setVal n mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { m = 4 } }
 where
  mainMem' = wr8 hl' addr mainMem
  hl' = setBitNOfX setVal n $ rd8 addr mainMem
  addr = joinRegs h l regs

setBitNOfX :: Boolean -> Int -> I8 -> I8
setBitNOfX setVal n x = if setVal
  then x .|. bitShift
  else x .&. complement bitShift
 where bitShift = 1 `shl` n

-- Rotations
-- =========

data Dir = Left | Right
derive instance eqDir :: Eq Dir

-- RL[C] A, base-opcode
-- RR[C] A, base-opcode
rotA :: Dir -> Boolean -> Regs -> Regs
rotA dir isCarryRot regs = regs { a = rotated.res, f = rotated.flags, m = 1 }
 where rotated = rotX {dir,isCarryRot,isCB:false} regs.a regs.f

-- RL[C] R, CB-prefix instruction
-- RR[C] R, CB-prefix instruction
rotReg :: Dir -> Boolean -> SetReg -> GetReg -> Regs -> Regs
rotReg dir isCarryRot setReg getReg regs =
  setReg rotated.res regs { f = rotated.flags, m = 2 }
 where rotated = rotX {dir,isCarryRot,isCB:true} (getReg regs) regs.f

-- RL[C] (HL)
-- RR[C] (HL)
rotHLMem :: Dir -> Boolean -> Mem -> Mem
rotHLMem dir isCarryRot mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { f = rotated.flags, m = 4 } }
 where
  mainMem' = wr8 rotated.res addr mainMem
  addr = joinRegs h l regs
  rotated = rotX {dir,isCarryRot,isCB:true} (rd8 addr mainMem) regs.f

rotX :: {dir::Dir,isCarryRot::Boolean,isCB::Boolean}
      -> I8 -> I8 -> { res :: I8, flags :: I8 }
rotX { dir, isCarryRot, isCB } x oldFlags = { res, flags }
 where
  res = 255 .&. (currCarry + (x `shiftFunc` 1))
  -- RotA basic, use old flags, RotR CB, start fresh with zeroFlag
  flags = setNewCarry
    $ if isCB then testZeroFlag res else oldFlags
  currCarry = if   (not isCarryRot) && (oldFlags .&. carryFlag /= 0)
                || (isCarryRot      && isNewCarry)
    then case dir of
      Left -> 1
      Right -> 0x80
    else 0
  setNewCarry = if isNewCarry
    then (carryFlag .|. _)
    else (cmplCarryFlag .&. _)
  isNewCarry = x .&. edgeBit /= 0
  edgeBit = case dir of
    Left -> 0x80
    Right -> 1
  shiftFunc = case dir of
    Left -> shl
    Right -> zshr

-- Shifts
-- ======

--SLA R
--SRA R preserve sign
--SRL R
shiftReg :: Dir -> Boolean -> SetReg -> GetReg
         -> Regs -> Regs
shiftReg dir sign setReg getReg regs =
  setReg shifted.res regs { f = shifted.flags, m = 2 }
 where
  shifted = shiftX dir sign $ getReg regs

shiftMemHL :: Dir -> Boolean -> Mem -> Mem
shiftMemHL dir sign mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs { f = shifted.flags, m = 4 } }
 where
  mainMem' = wr8 shifted.res addr mainMem
  shifted = shiftX dir sign $ rd8 addr mainMem
  addr = joinRegs h l regs

shiftX :: Dir -> Boolean -> I8
       -> { res::I8, flags::I8 }
shiftX dir sign x = { res, flags }
 where
  res =  (255 .&. _)
     <<< (if sign && dir == Right then (currCarry + _) else id)
      $  x `shiftFunc` 1
  shiftFunc = case dir of
    Left -> shl
    Right -> zshr
  flags = newCarry .|. testZeroFlag res
  newCarry = if x .&. edgeBit /= 0 then carryFlag else 0
  currCarry = x .&. 0x80 --sign is only relevant for right shifts
  edgeBit = case dir of
    Left -> 0x80
    Right -> 1

-- Increments / Decrements
-- =======================

--INC R
incReg :: GetReg -> SetReg -> Regs -> Regs
incReg = incDecReg incI8

--DEC R
decReg :: GetReg -> SetReg -> Regs -> Regs
decReg getReg setReg regs = regs' { f = setFlag subtractionFlag regs'.f }
  where regs' = (incDecReg decI8 getReg setReg regs)

incDecReg :: (I8 -> { res :: I8, carry :: Boolean })
          -> GetReg -> SetReg -> Regs -> Regs
incDecReg op getReg setReg regs =
  setReg inced .res regs {f = testZeroFlag inced.res, m = 1}
 where
  inced  = op $ getReg regs

--INC RR
incRegWithCarry :: GetReg -> GetReg -> SetReg -> SetReg
                -> Regs -> Regs
incRegWithCarry = incDecRegWithCarry incI8

--DEC RR
decRegWithCarry :: GetReg -> GetReg -> SetReg -> SetReg
                -> Regs -> Regs
decRegWithCarry = incDecRegWithCarry decI8

--NOTE: consider refactoring to something along the lines of
--"split2Regs . incI16 . joinRegs", if you have leeway, performance-wise.
incDecRegWithCarry :: (I8 -> { res :: I8, carry :: Boolean })
          -> GetReg -> GetReg -> SetReg -> SetReg
          -> Regs -> Regs
incDecRegWithCarry op getCarryReg getMainReg setCarryReg setMainReg regs =
  setCarryReg carryReg' <<< setMainReg mainReg' $ regs{m = 1}
 where
  carryReg' = (if mainRegInc.carry then _.res <<< op else id)
    $ getCarryReg regs
  mainReg' = mainRegInc.res
  mainRegInc = op $ getMainReg regs

--INC SP
incSP :: Regs -> Regs
incSP = incDecSP $ \x -> 0xFFFF .&. (x + 1)

--DEC SP
decSP :: Regs -> Regs
decSP = incDecSP $ \x -> 0xFFFF .&. (x - 1)

incDecSP :: (I16 -> I16) -> Regs -> Regs
incDecSP op regs = regs {sp = op $ regs.sp, m = 1}

--INC (HL)
incHLMem :: Mem -> Mem
incHLMem = incDecHLMem incI8

--DEC (HL)
decHLMem :: Mem -> Mem
decHLMem = incDecHLMem decI8

incDecHLMem :: (I8 -> { res :: I8, carry :: Boolean })
            -> Mem -> Mem
incDecHLMem op mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs {f = testZeroFlag inced.res, m = 3} }
 where
  mainMem' = wr8 inced.res addr mainMem
  inced = op $ rd8 addr mainMem
  addr = joinRegs h l regs

incI8 :: I8 -> { res :: I8, carry :: Boolean }
incI8 i8 = { res : i8', carry : i8' == 0 }
 where
  i8' = (i8 + 1) .&. 255

decI8 :: I8 -> { res :: I8, carry :: Boolean }
decI8 i8 = { res : i8', carry : i8' == 255 }
 where
  i8' = (i8 - 1) .&. 255

-- Loads
-- =====

--LD R,R
ldRegFromReg :: SetReg -> GetReg
      -> Regs -> Regs
ldRegFromReg setDestReg getSrcReg regs = setDestReg (getSrcReg regs)
  regs { m = 1 }

--LD (nn),SP
ldMemImmFromSP :: Mem -> Mem
ldMemImmFromSP mem@{mainMem, regs} =
  mem { mainMem = mainMem', regs = regs { pc = regs.pc + 3, m = 5 } }
 where
  mainMem' = wr16 regs.sp addr mainMem
  addr = rd16 (regs.pc+1) mainMem

--LD SP,HL
ldSPFromHL :: Regs -> Regs
ldSPFromHL regs  =
  regs { sp = sp', m = 2 }
 where sp' = joinRegs h l regs

ldSPFromImm :: Mem -> Regs
ldSPFromImm { mainMem, regs } =
  regs { sp = sp', pc = regs.pc + 3, m = 3 }
 where sp' = rd16 (regs.pc+1) mainMem

--LD RR,nn
ldTwoRegsFromImm :: SetReg -> SetReg
               -> Mem -> Regs
ldTwoRegsFromImm setReg1 setReg2 mem = regs'' { m = 3 }
 where
  regs'  = ldRegFromImm setReg2 mem
  regs'' = ldRegFromImm setReg1 mem { regs = regs' }

--LD R,(IOC)
ldRegFromFF00CMem :: SetReg -> Mem -> Regs
ldRegFromFF00CMem setReg mem@{regs, mainMem} = regs' { m = 2 }
 where regs' = ldRegFromFF00PlusX regs.c setReg mem

--LD R,(IOn)
ldRegFromFF00ImmMem :: SetReg -> Mem -> Regs
ldRegFromFF00ImmMem setReg mem@{regs, mainMem} = regs' { pc = regs.pc + 2, m = 3 }
 where regs' = ldRegFromFF00PlusX (rd8 (regs.pc+1) mainMem) setReg mem

ldRegFromFF00PlusX :: I8 -> SetReg -> Mem -> Regs
ldRegFromFF00PlusX  x setReg mem@{regs} = ldRegFromMem setReg addr mem
 where addr = 0xFF00 + x

--LD R,n
ldRegFromImm :: SetReg -> Mem -> Regs
ldRegFromImm setReg mem@{regs} = regs' { pc = regs.pc + 2 }
  where regs' = ldRegFromMem setReg (regs.pc+1) mem

--LD R,(RR)
ldRegFromMem2R :: SetReg -> GetReg -> GetReg
               -> Mem -> Regs
ldRegFromMem2R setReg msByteReg lsByteReg mem@{regs} =
  ldRegFromMem setReg ((msByteReg regs `shl` 8) + lsByteReg regs) mem

--LD A,(nn)
ldRegAFromMemImm :: Mem -> Regs
ldRegAFromMemImm mem@{ mainMem, regs } = regs' { pc = regs.pc + 3, m = 4 }
 where
  regs' = ldRegFromMem setA addr mem
  addr = rd16 (regs.pc+1) mainMem

--LDD A,(HL)
ldRegFromMemHLDec :: SetReg -> Mem -> Regs
ldRegFromMemHLDec = ldRegFromMemHLIncDec decRegWithCarry

--LDI A,(HL)
ldRegFromMemHLInc :: SetReg -> Mem -> Regs
ldRegFromMemHLInc = ldRegFromMemHLIncDec incRegWithCarry

ldRegFromMemHLIncDec :: (GetReg -> GetReg -> SetReg -> SetReg -> Regs -> Regs)
                   -- Inc/Dec reg operation's type signature
                   -> SetReg -> Mem -> Regs
ldRegFromMemHLIncDec op setReg mem@{ regs } = regs' {m = 2}
 where
  regs' = op h l setH setL
        $ ldRegFromMem setA addr mem
  addr = joinRegs h l regs

ldRegFromMem :: SetReg -> I16 -> Mem -> Regs
ldRegFromMem setReg addr { mainMem, regs } =
  setReg (rd8 addr mainMem) regs{m = 2}

--LD HL,SP|n|
ldHLFromSPImm :: Mem -> Regs
ldHLFromSPImm  mem@{ regs, mainMem } =
  regs {h = split.ms, l = split.ls, pc = regs.pc + 2, m = 3}
 where
  split = splitI16 x
  x = (absI8 imm + regs.sp) .&. 0xFFFF
  imm = rd8 (regs.pc+1) mainMem

--LDI (HL),A
ldMemHLFromRegInc :: GetReg -> Mem -> Mem
ldMemHLFromRegInc = ldMemHLFromRegIncDec incRegWithCarry

--LDD (HL),A
ldMemHLFromRegDec :: GetReg -> Mem -> Mem
ldMemHLFromRegDec  = ldMemHLFromRegIncDec decRegWithCarry

ldMemHLFromRegIncDec :: (GetReg -> GetReg -> SetReg -> SetReg -> Regs -> Regs)
                   -- Inc/Dec reg operation's type signature
                   -> GetReg
                   -> Mem -> Mem
ldMemHLFromRegIncDec op getReg mem@{regs} =
  mem { mainMem = mainMem', regs = regs' {m = 2} }
 where
  regs' = op h l setH setL regs
  mainMem' = _.mainMem $ ldMemFromReg addr a mem
  addr = joinRegs h l regs

--LD (IOC),R
ldFF00CMemFromReg :: GetReg -> Mem -> Mem
ldFF00CMemFromReg getReg mem@{regs, mainMem} =
  mem' { regs = regs{m = 2} }
 where mem' = ldFF00PlusXFromReg  regs.c getReg mem

--LD (IOn),R
ldFF00ImmMemFromReg :: GetReg -> Mem -> Mem
ldFF00ImmMemFromReg getReg mem@{regs, mainMem} =
  mem' { regs = regs{pc = regs.pc + 2, m = 3} }
 where
  mem' = ldFF00PlusXFromReg imm getReg mem
  imm = rd8 (regs.pc+1) mainMem

ldFF00PlusXFromReg :: I8 -> GetReg -> Mem -> Mem
ldFF00PlusXFromReg   x getReg mem@{regs, mainMem} = mem'
 where
  mem' = ldMemFromReg addr getReg mem
  addr = 0xFF00 + x

--LD (nn),A
ldMemImmFromRegA :: Mem -> Mem
ldMemImmFromRegA  mem@{ mainMem, regs } =
 mem' { regs = regs{m = 4, pc = regs.pc + 3} }
 where
  mem' = ldMemFromReg addr a mem
  addr = rd16 (regs.pc+1) mainMem

--LD (HL),n
ldMemHLFromImm :: Mem -> Mem
ldMemHLFromImm  mem@{mainMem,regs} =
  mem { mainMem = mainMem' , regs = regs{m = 3, pc = regs.pc + 2} }
 where
  mainMem' = wr8 imm addr mainMem
  addr =  joinRegs h l regs
  imm = rd8 (regs.pc+1) mainMem

--LD (RR),R
ldMem2RFromReg :: GetReg -> GetReg -> GetReg
                        -> Mem -> Mem
ldMem2RFromReg msbAddrReg lsbAddrReg getReg mem@{ regs } =
  ldMemFromReg (joinRegs msbAddrReg lsbAddrReg regs) getReg mem

ldMemFromReg :: I16 -> GetReg
           -> Mem -> Mem
ldMemFromReg addr getReg mem@{ regs } = ldMemFromX addr (getReg regs) mem

ldMemFromX :: I16 -> I8
           -> Mem -> Mem
ldMemFromX addr x mem@{mainMem,regs} =
  mem { mainMem = mainMem', regs = regs{m = 2} }
 where mainMem' = wr8 x addr mainMem

-- Compares
-- ========

--CP A,R
compAToReg :: GetReg -> Regs -> Regs
compAToReg getReg regs = compAToX (getReg regs) regs

--CP A,(HL)
compAToMemHL :: Mem -> Regs
compAToMemHL { mainMem, regs } = (compAToX hlMem regs) { m = 2 }
 where hlMem = rd8 (joinRegs h l regs) mainMem

--CP A,n
compAToImm :: Mem -> Regs
compAToImm { mainMem, regs } = (compAToX imm regs) { pc = regs.pc + 2, m = 2 }
 where imm = rd8 (regs.pc+1) mainMem

compAToX :: I8 -> Regs -> Regs
compAToX x regs = regs { f = diff.flags, m = 1 }
 where
  diff = subI8s regs.a x

-- Stack operations
-- ================

--NOTE: On the one hand, there's an advantage of letting the compiler catch
--more mistakes by using row polymorphism in this case,
--on the other hand it obfuscates the signature, is there a better way?
--forall r s. { regs :: {sp::I8 | s}, mem :: MainMem | r }
--         -> { regs :: {sp::I8 | s}, mem :: MainMem | r }
--NOTE: Only 4 versions of PUSH so consider hard-coding them for performance.
--PUSH RR
pushReg :: GetReg -> GetReg
        -> Mem -> Mem
pushReg msByteReg lsByteReg mem@{regs,mainMem} =
  mem { mainMem = mainMem', regs = regs { sp = regs.sp - 2, m = 3 } }
 where
  mainMem' =  wr8 (lsByteReg regs) (regs.sp - 2)
          <<< wr8 (msByteReg regs) (regs.sp - 1)
           $  mainMem

--POP RR
popReg :: SetReg -> SetReg
       -> Mem -> Regs
popReg setMsByteReg setLsByteReg { mainMem, regs = regs@{sp} }
    = setMsByteReg (rd8 (sp + 1) mainMem)
   <<< setLsByteReg (rd8 sp mainMem)
    $  regs { sp = sp + 2, m = 3 }

-- Jumps
-- =====

-- JR NZ,n
-- JR Z,n
-- JR NC,n
-- JR C,n
jumpRelImmFlag :: Boolean -> I8 -> Mem -> Regs
jumpRelImmFlag inverse flag mem@{mainMem, regs} =
  if inverse `xor` isSetFlag flag regs.f
    then jumpRelImm mem
    else regs { pc = regs.pc + 2, m = 2 }

{----JR n--}
{--jumpRelImm :: Mem -> Regs--}
{--jumpRelImm { mainMem, regs } =--}
  {--regs { pc = pc', m = 3 }--}
 {--where--}
  {----NOTE: is it necessary to abs the immediate?--}
  {--pc' = regs.pc + 1 + absI8 imm--}
  {--imm = rd8 regs.pc mainMem--}

--JR n
jumpRelImm :: Mem -> Regs
jumpRelImm { mainMem, regs } =
  regs { pc = pc', m = 3 }
 where
  --NOTE: is it necessary to abs the immediate?
  pc' = regs.pc + 2 + absI8 imm
  imm = rd8 (regs.pc+1) mainMem

-- JP NZ,nn
-- JP Z,nn
-- JP NC,nn
-- JP C,nn
jumpImmFlag :: Boolean -> I8 -> Mem -> Regs
jumpImmFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then (jumpImm mem) { m = 4 }
    else regs { pc = regs.pc + 3, m = 3 }

-- JP nn
jumpImm :: Mem -> Regs
jumpImm { mainMem, regs } =
  regs { pc = pc', m = 3 }
 where pc' = rd16 (regs.pc+1) mainMem

-- JP (HL)
jumpHL :: Regs -> Regs
jumpHL regs = regs { pc = pc', m = 1 }
 where pc' = joinRegs h l regs

{--jumpX :: I8 -> Regs--}
{--jumpX { mainMem, regs } =--}
  {--regs { pc = pc', m = 3 }--}
 {--where pc' = rd16 regs.pc mainMem--}

-- Call Z,nn
-- Call NZ,nn
-- Call C,nn
-- Call NC,nn
callImmFlag :: Boolean -> I8 -> Mem -> Mem
callImmFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then callImm mem
    else mem { regs = regs {pc = regs.pc + 3, m = 3} }

-- Call nn
callImm :: Mem -> Mem
callImm mem@{mainMem,regs} =
  mem { mainMem =  mainMem', regs = regs { pc = imm, sp = regs.sp - 2, m = 5 } }
 where
  mainMem' = wr16 (regs.sp-2) (regs.pc + 2) $ mainMem
  imm = rd16 (regs.pc+1) mainMem

-- RET Z
-- RET NZ
-- RET C
-- RET NC
retFlag :: Boolean -> I8 -> Mem -> Regs
retFlag inverse flag mem@{mainMem,regs} =
  if inverse `xor` isSetFlag flag regs.f
    then ret mem
    else regs { m = 1 }

-- RETI
retEnableInterrupt :: Mem -> Regs
retEnableInterrupt { mainMem, regs, svdRegs } =
  restoreRegs svdRegs regs { ime = true, pc = pc', sp = regs.sp + 2, m = 3 }
 where pc' = rd16 regs.sp mainMem

-- RET
ret :: Mem -> Regs
ret { mainMem, regs } =
  regs { pc = pc', sp = regs.sp + 2, m = 3 }
 where pc' = rd16 regs.sp mainMem

--RST XX
callRoutine :: I16 -> Mem -> Mem
callRoutine addr { mainMem, regs, svdRegs } =
  { mainMem : mainMem', regs : regs', svdRegs : saveRegs regs svdRegs }
 where
  regs' = regs { sp = regs.sp - 2, pc = addr, m = 3 }
  mainMem' = wr16 regs.pc (regs.sp - 2) mainMem

-- Misc.
-- =====

--SWAP R
swapReg :: SetReg -> GetReg -> Regs -> Regs
swapReg setReg getReg regs = setReg reg' $ regs { m = 1 }
 where
  reg' = ((0x0F.&.reg) `shl` 4) .|. ((0xF0.&.reg) `zshr` 4)
  reg = getReg regs

--SWAP (HL)
swapMemHL :: Mem -> Mem
swapMemHL mem@{mainMem, regs} =
  mem { mainMem = mainMem', regs = regs { m = 4 } }
 where
  mainMem' = wr8 hlMem' addr mainMem
  hlMem' = ((0x0F.&.hlMem) `shl` 4) .|. ((0xF0.&.hlMem) `zshr` 4)
  hlMem = rd8 addr mainMem
  addr = joinRegs h l regs

--Why a command called 'clear' actually flips the tag
--and not unset it, is a question that shouldn't be directed at me!
--CCF
clearCarryFlag :: Regs -> Regs
clearCarryFlag = changeCarryFlag false

--SCF
setCarryFlag :: Regs -> Regs
setCarryFlag = changeCarryFlag true

changeCarryFlag :: Boolean -> Regs -> Regs
changeCarryFlag isSet regs = regs { f = f', m = 1 }
 where f' =  unsetFlag subtractionFlag
         <<< unsetFlag halfCarryFlag
         <<< (if isSet then setFlag else flipFlag) carryFlag
          $  regs.f

--DAA
adjAForBCDAdd :: Regs -> Regs
adjAForBCDAdd regs = regs { f = adjF regs.f, a = a', m = 1 }
 where
  a' = adjHighNyb <<< adjLowNyb $ regs.a
  adjLowNyb = if (regs.f .&. halfCarryFlag /= 0) || (regs.a .&. 15 > 9)
    then (6 + _) else id
  adjHighNyb = if shouldAdjHighNyb then (0x60 + _) else id
  adjF = if shouldAdjHighNyb then (carryFlag .|. _) else (cmplCarryFlag .&. _)
  shouldAdjHighNyb = (regs.f .&. halfCarryFlag /= 0) || (regs.a > 0x99)

--NOP
nop :: Regs -> Regs
nop = _ { m = 1}

--Invalid opCode
invalidOpCode :: Z80State -> Z80State
--NOTE: LOG THIS! for debugging purposes
invalidOpCode state = state { stop = true }

--STOP
--NOTE: this is just a guess
stop :: Z80State -> Z80State
stop state@{mem=mem@{regs}} =
  state { stop = true
        , mem = mem { regs = regs { m = 1 } }
        }

--HALT
halt :: Z80State -> Z80State
halt state@{mem=mem@{regs}} =
  state { halt = true
        , mem = mem { regs = regs { m = 1 } }
        }

--DI
--EI
setInterrupts :: Boolean -> Regs -> Regs
setInterrupts enable = _ { ime = enable, m = 1 }

--Extended Ops
execExtOps :: Array (Z80State -> Z80State)
           -> Z80State -> Z80State
execExtOps opsMap state@{ mem = mem@{mainMem,regs} } = extOp state'
 where
  --NOTE replace toMaybe, with something to log errors
  extOp = getOpcode opCode opsMap
  state' = state { mem = mem { regs = regs { pc = 65535 .&. (regs.pc + 1) } } }
  opCode = rd8 regs.pc mainMem

-- Helpers
-- =======

saveRegs :: Regs -> SavedRegs -> SavedRegs
saveRegs regs svdRegs =
  svdRegs { a = regs.a, b = regs.b, c = regs.c, d = regs.d
          , e = regs.e, f = regs.f, h = regs.h, l = regs.l
          }

restoreRegs :: SavedRegs -> Regs -> Regs
restoreRegs svdRegs regs =
  regs { a = svdRegs.a, b = svdRegs.b, c = svdRegs.c, d = svdRegs.d
       , e = svdRegs.e, f = svdRegs.f, h = svdRegs.h, l = svdRegs.l
       }

joinRegs :: GetReg -> GetReg -> Regs -> I16
joinRegs msByteReg lsByteReg regs = (msByteReg regs `shl` 8) + lsByteReg regs

absI8 :: I8 -> I8
absI8 i8 = if i8 > 127 then negI8 i8 else i8

negI8 :: I8 -> I8
negI8 i8 = (complement i8 + 1) .&. 255

splitI16 :: I16 -> { ms :: I8, ls :: I8 }
splitI16 x = { ms : (x `zshr` 8) .&. 255
             , ls : x .&. 255
             }

isSetFlag :: I8 -> I8 -> Boolean
isSetFlag f fs = f .&. fs /= 0

setFlag :: I8 -> I8 -> I8
setFlag  f = (f .|. _)

unsetFlag :: I8 -> I8 -> I8
unsetFlag f = ((complement f) .&. _)

flipFlag :: I8 -> I8 -> I8
flipFlag f = (f .^. _)

testZeroFlag :: Int -> Int
testZeroFlag x = if x == 0 then 0x80 else 0

setZeroFlag :: Int -> I8 -> I8
setZeroFlag x = if x == 0
  then (_ .|. zeroFlag)
  else (_ .&. cmplZeroFlag)

testCarryFlag8 :: Int -> I8
testCarryFlag8 x = if x > 255 then carryFlag else 0

setCarryFlag16 :: Int -> I8 -> I8
setCarryFlag16 x = if x > 0xFFFF
  then (_ .|. carryFlag)
  else (_ .&. cmplCarryFlag)

testNegCarryFlag :: Int -> I8
testNegCarryFlag reg = if reg < 0 then carryFlag else 0

testHalfCarryFlag :: I8 -> I8 -> I8 -> I8
testHalfCarryFlag reg1 reg2 sum = if res /= 0 then halfCarryFlag else 0
 where res = 0x10 .&. (reg1 .^. reg2 .^. sum)

setHalfCarryFlag :: I8 -> I8 -> I8 -> I8 -> I8
setHalfCarryFlag reg1 reg2 sum =
  if res /= 0 then (_ .|. halfCarryFlag) else (_ .&. cmplHalfCarryFlag)
 where res = 0x10 .&. (reg1 .^. reg2 .^. sum)

--NOTE: replace fromMaybe with an error report mechanism to trace bad cases
getOpcode :: Int -> Array (Z80State -> Z80State) -> Z80State -> Z80State
getOpcode ix table = fromMaybe id $ table A.!! ix

xor :: Boolean -> Boolean -> Boolean
xor true x = not x
xor false x = x

zeroFlag :: I8
zeroFlag = 0x80
cmplZeroFlag :: I8
cmplZeroFlag = 0x7F
carryFlag :: I8
carryFlag = 0x10
cmplCarryFlag :: I8
cmplCarryFlag = 0xEF
halfCarryFlag :: I8
halfCarryFlag = 0x20
cmplHalfCarryFlag :: I8
cmplHalfCarryFlag = 0xD0
subtractionFlag :: I8
subtractionFlag = 0x40

type GetReg = Regs -> I8
type SetReg = I8 -> Regs -> Regs