module Intermediate where

import Data.Bits
import Data.List
import Data.Word

import AST
import Defs
import Pretty


data BB = BB Id [IRIns] IRTerm
  deriving (Show, Eq)

data Ref = Temp Size Int | StructTemp Size Int | Argument Size Name | Global Size Name | Constant Size Value
  deriving (Show, Eq, Ord)

data IRProgram = IRProgram [DVar] [IRFunc]
  deriving (Show, Eq)

data IRFunc = IRFunc (Maybe Type) Name [(Type, Name)] [BB] Id
  deriving (Show, Eq)

data IRIns
    = IMov Ref Ref
    | ILea Ref Name
    | IStore Ref Ref
    | ILoad Ref Ref
    | ISet Ref Offset Ref
    | IGet Ref Ref Offset
    | IAri ArithType Ref Ref Ref  -- destination, operand 1, operand 2
    | ICall Name [Ref]
    | ICallr Ref Name [Ref]
    | IResize Ref Ref
    | IDebugger
    | INop
  deriving (Show, Eq)

data IRTerm
    = IJcc CmpType Ref Ref Id Id  -- Id Id == if-yes if-no
    | IJmp Id
    | IRet
    | IRetr Ref
    | IUnreachable
    | ITermNone
  deriving (Show, Eq)

data ArithType
    = AAdd | ASub | AMul | ADiv | AMod
    | AAnd | AOr | AXor
    | AEq | ANeq | AGt | ALt | AGeq | ALeq
  deriving (Show, Eq)

data CmpType
    = CEq | CNeq | CGt | CLt | CGeq | CLeq | CUGt | CULt | CUGeq | CULeq
  deriving (Show, Eq)


instance Pretty BB where
    prettyI i (BB bid inss term) =
        "{{{(" ++ show bid ++ ")\n" ++ indent (i+1) ++
            intercalate ("\n" ++ indent (i+1)) (map pretty inss) ++
            (if null inss then "" else "\n" ++ indent (i+1)) ++
                pretty term ++
            "\n" ++ indent i ++ "}}}"
      where
        indent n = replicate (2*n) ' '

instance Pretty Ref where
    prettyI _ (Temp sz k) = "t" ++ show k ++ pretty_sizeSuffix sz
    prettyI _ (StructTemp sz k) = "s" ++ show k ++ "{" ++ show sz ++ "}"
    prettyI _ (Argument sz n) = "a" ++ n ++ pretty_sizeSuffix sz
    prettyI _ (Global sz n) = "g" ++ n ++ pretty_sizeSuffix sz
    prettyI _ (Constant sz n) = show n ++ pretty_sizeSuffix sz

pretty_sizeSuffix :: Size -> String
pretty_sizeSuffix 1 = "B"
pretty_sizeSuffix 2 = "W"
pretty_sizeSuffix 4 = "D"
pretty_sizeSuffix 8 = "Q"
pretty_sizeSuffix sz = "<" ++ show sz ++ ">"

instance Pretty IRProgram where
    prettyI i (IRProgram vars funcs) =
        intercalate ("\n" ++ indent i) (map (prettyI i) vars) ++
            "\n" ++ indent i ++
            intercalate ("\n" ++ indent i) (map (prettyI (i+1)) funcs) ++
            "\n"
      where
        indent n = replicate (2*n) ' '

instance Pretty IRFunc where
    prettyI i (IRFunc mt n al bbs sid) =
        "irfunc" ++ maybe "" ((' ' :) . prettyI i) mt ++ " " ++ n ++ "(" ++
            intercalate ", "
                (map (\(at,an) -> prettyI i at ++ " " ++ an) al) ++
            ")\n" ++ indent i ++
            intercalate ("\n" ++ indent i) (map (prettyI i) sorted)
      where
        indent n' = replicate (2*n') ' '

        sorted = uncurry (++) $ partition (\(BB bid _ _) -> bid == sid) bbs

instance Pretty IRIns where
    prettyI _ (IMov d s) = "mov " ++ pretty d ++ " <- " ++ pretty s
    prettyI _ (ILea d s) = "lea " ++ pretty d ++ " <- &[" ++ s ++ "]"
    prettyI _ (IStore d s) = "store *" ++ pretty d ++ " <- " ++ pretty s
    prettyI _ (ILoad d s) = "load " ++ pretty d ++ " <- *" ++ pretty s
    prettyI _ (ISet d off s) = "set " ++ pretty d ++ ".[" ++ show off ++ "] <- " ++ pretty s
    prettyI _ (IGet d s off) = "get " ++ pretty d ++ " <- " ++ pretty s ++ ".[" ++ show off ++ "]"
    prettyI _ (IAri at d s1 s2) =
        pretty at ++ " " ++ pretty d ++ " <- " ++ pretty s1 ++ ", " ++ pretty s2
    prettyI _ (ICall n al) =
        "call " ++ n ++ " (" ++ intercalate ", " (map pretty al) ++ ")"
    prettyI _ (ICallr d n al) =
        "call " ++ pretty d ++ " <- " ++ n ++ " (" ++ intercalate ", " (map pretty al) ++ ")"
    prettyI _ (IResize d s) = "resize " ++ pretty d ++ " <- " ++ pretty s
    prettyI _ IDebugger = "debugger"
    prettyI _ INop = "nop"

instance Pretty IRTerm where
    prettyI _ (IJcc ct s1 s2 did1 did2) =
        pretty ct ++ " " ++ pretty s1 ++ ", " ++ pretty s2 ++ " -> " ++ show did1 ++ " | " ++ show did2
    prettyI _ (IJmp did) = "jmp " ++ show did
    prettyI _ IRet = "ret"
    prettyI _ (IRetr ref) = "retr " ++ pretty ref
    prettyI _ IUnreachable = "unreachable"
    prettyI _ ITermNone = "?NONE?"
 
instance Pretty ArithType where
    prettyI _ AAdd = "add"
    prettyI _ ASub = "sub"
    prettyI _ AMul = "mul"
    prettyI _ ADiv = "div"
    prettyI _ AMod = "mod"
    prettyI _ AAnd = "and"
    prettyI _ AOr = "or"
    prettyI _ AXor = "xor"
    prettyI _ AEq = "eq"
    prettyI _ ANeq = "neq"
    prettyI _ AGt = "gt"
    prettyI _ ALt = "lt"
    prettyI _ AGeq = "geq"
    prettyI _ ALeq = "leq"

instance Pretty CmpType where
    prettyI _ CEq = "jeq"
    prettyI _ CNeq = "jne"
    prettyI _ CGt = "jg"
    prettyI _ CLt = "jl"
    prettyI _ CGeq = "jge"
    prettyI _ CLeq = "jle"
    prettyI _ CUGt = "jug"
    prettyI _ CULt = "jul"
    prettyI _ CUGeq = "juge"
    prettyI _ CULeq = "jule"


blockIdOf :: BB -> Id
blockIdOf (BB bid _ _) = bid

refSize :: Ref -> Size
refSize (Temp sz _) = sz
refSize (StructTemp sz _) = sz
refSize (Argument sz _) = sz
refSize (Global sz _) = sz
refSize (Constant sz _) = sz

isStructTemp :: Ref -> Bool
isStructTemp (StructTemp _ _) = True
isStructTemp _ = False

isConstant :: Ref -> Bool
isConstant (Constant _ _) = True
isConstant _ = False


evaluateArith :: ArithType -> Value -> Value -> Value
evaluateArith at a b = case at of
    AAdd -> a + b
    ASub -> a - b
    AMul -> a * b
    ADiv -> if b == 0 then error "Division by zero detected" else a `div` b
    AMod -> if b == 0 then error "Modulo by zero detected" else a `mod` b
    AAnd -> a .&. b
    AOr -> a .|. b
    AXor -> a `xor` b
    AEq -> if a == b then 1 else 0
    ANeq -> if a /= b then 1 else 0
    AGt -> if a > b then 1 else 0
    ALt -> if a < b then 1 else 0
    AGeq -> if a >= b then 1 else 0
    ALeq -> if a <= b then 1 else 0

evaluateCmp :: CmpType -> Value -> Value -> Bool
evaluateCmp ct a b = case ct of
    CEq -> a == b
    CNeq -> a /= b
    CGt -> a > b
    CLt -> a < b
    CGeq -> a >= b
    CLeq -> a <= b
    CUGt -> (fromIntegral a :: Word64) > (fromIntegral b :: Word64)
    CULt -> (fromIntegral a :: Word64) < (fromIntegral b :: Word64)
    CUGeq -> (fromIntegral a :: Word64) >= (fromIntegral b :: Word64)
    CULeq -> (fromIntegral a :: Word64) <= (fromIntegral b :: Word64)

isCommutative :: ArithType -> Bool
isCommutative AAdd = True
isCommutative AMul = True
isCommutative AAnd = True
isCommutative AOr = True
isCommutative AXor = True
isCommutative AEq = True
isCommutative ANeq = True
isCommutative ASub = False
isCommutative ADiv = False
isCommutative AMod = False
isCommutative AGt = False
isCommutative ALt = False
isCommutative AGeq = False
isCommutative ALeq = False

isIMov :: IRIns -> Bool
isIMov (IMov _ _) = True
isIMov _ = False

isILoad :: IRIns -> Bool
isILoad (ILoad _ _) = True
isILoad _ = False

isIAri :: IRIns -> Bool
isIAri (IAri _ _ _ _) = True
isIAri _ = False

isIResize :: IRIns -> Bool
isIResize (IResize _ _) = True
isIResize _ = False

jumpTargets :: IRTerm -> [Id]
jumpTargets (IJcc _ _ _ i2 i1) = [i1, i2]
jumpTargets (IJmp i) = [i]
jumpTargets _ = []