{-# LANGUAGE LambdaCase, FlexibleInstances #-}
module Intermediate where

import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe

import AST
import Util


data IRProgram =
    IRProgram [BB]
              (Map.Map Name GlobFuncDef)
              [Value]  -- data table

data GlobFuncDef =
    GlobFuncDef Int  -- BB id of implementation
                Int  -- number of arguments
                [Name]  -- closure slots

data BB = BB Int [Instruction] Terminator

type Instruction = (Ref, InsCode)

data Ref
    = RConst Int
    | RTemp Int
    | RSClo Name  -- static closure object of a function
    | RNone
  deriving (Eq, Ord)

data InsCode
    = IAssign Ref
    | IParam Int  -- first param is self-recurse link
    | IClosure Int
    | IData Int
    | ICallC Ref [Ref]
    | IAllocClo Name [Ref]
    | IDiscard Ref
    | IPush [Ref]  -- pushes references on the stack; should be matched with an IPop with the same number of references
    | IPop [Ref]
  deriving Eq

data Terminator
    = IBr Ref Int Int
    | IJmp Int
    | IRet Ref
    | ITailC Ref [Ref]
    | IExit
    | IUnknown
  deriving Eq


bidOf :: BB -> Int
bidOf (BB i _ _) = i

inssOf :: BB -> [Instruction]
inssOf (BB _ i _) = i

termOf :: BB -> Terminator
termOf (BB _ _ t) = t

class AllRefs a where
    allRefs :: a -> [Ref]

instance AllRefs IRProgram where
    allRefs (IRProgram bbs _ _) = allRefs bbs

instance AllRefs BB where
    allRefs (BB _ inss term) =
        sortUniq $ concatMap (allRefs . snd) inss ++ allRefs term

instance AllRefs [BB] where
    allRefs = sortUniq . concatMap allRefs

instance AllRefs InsCode where
    allRefs (IAssign r) = [r]
    allRefs (IParam _) = []
    allRefs (IClosure _) = []
    allRefs (IData _) = []
    allRefs (ICallC r rs) = sortUniq (r : rs)
    allRefs (IAllocClo _ rs) = sortUniq rs
    allRefs (IDiscard r) = [r]
    allRefs (IPush rs) = sortUniq rs
    allRefs (IPop rs) = sortUniq rs

instance AllRefs Terminator where
    allRefs (IBr r _ _) = [r]
    allRefs (IJmp _) = []
    allRefs (IRet r) = [r]
    allRefs (ITailC r rs) = sortUniq (r : rs)
    allRefs IExit = []
    allRefs IUnknown = []

outEdges :: BB -> [Int]
outEdges (BB _ _ term) = outEdgesT term

outEdgesT :: Terminator -> [Int]
outEdgesT (IBr _ a b) = [a, b]
outEdgesT (IJmp a) = [a]
outEdgesT (IRet _) = []
outEdgesT (ITailC _ _) = []
outEdgesT IExit = []
outEdgesT IUnknown = []

icReadTemps :: InsCode -> [Int]
icReadTemps = \case
    IAssign r -> onlyTemporaries [r]
    IParam _ -> []
    IClosure _ -> []
    IData _ -> []
    ICallC r rs -> onlyTemporaries (r : rs)
    IAllocClo _ rs -> onlyTemporaries rs
    IDiscard r -> onlyTemporaries [r]
    IPush rs -> onlyTemporaries rs
    IPop _ -> []

termReadTemps :: Terminator -> [Int]
termReadTemps = \case
    IBr r _ _ -> onlyTemporaries [r]
    IJmp _ -> []
    IRet r -> onlyTemporaries [r]
    ITailC r rs -> onlyTemporaries (r : rs)
    IExit -> []
    IUnknown -> []

bbReadTemps :: BB -> [Int]
bbReadTemps (BB _ inss term) = onlyTemporaries (concatMap (allRefs . snd) inss ++ allRefs term)

bbWrittenTemps :: BB -> [Int]
bbWrittenTemps (BB _ inss _) = concatMap insWrittenTemps inss

insWrittenTemps :: Instruction -> [Int]
insWrittenTemps (_, IPop rs) = onlyTemporaries rs
insWrittenTemps (d, _) = onlyTemporaries [d]

onlyTemporaries :: [Ref] -> [Int]
onlyTemporaries rs = [i | RTemp i <- rs]


instance Show IRProgram where
    show (IRProgram bbs gfds datas) = intercalate "\n" $
        ["IRPROGRAM", "Data Table:"] ++ map (("- " ++) . show) datas ++
        ["Global functions:"] ++ map (\(n, gfd) -> "- " ++ n ++ ": " ++ show gfd) (Map.assocs gfds) ++
        ["Blocks:"] ++ [intercalate "\n" (map (genericShowBB bbannot icshow termshow) bbs)]
      where
        annotate s "" = s
        annotate s a = s ++ "   ; " ++ a
        refAnnot rs = intercalate ", " . catMaybes $
                        [case Map.lookup name gfds of
                             Nothing -> Nothing
                             Just (GlobFuncDef i _ _) -> Just (name ++ " = BB " ++ show i)
                        | RSClo name <- nub rs]
        safeIndex l i = if 0 <= i && i < length l then Just (l !! i) else Nothing
        icshow ins@(IData n) = annotate (show ins) (maybe "??" show (datas `safeIndex` n))
        icshow ins = annotate (show ins) (refAnnot (allRefs ins))
        termshow term = annotate (show term) (refAnnot (allRefs term))
        bidToName = Map.fromList [(bid, n) | (n, GlobFuncDef bid _ _) <- Map.assocs gfds]
        bbannot bid = maybe "" ("entry point of " ++) (Map.lookup bid bidToName)

instance Show GlobFuncDef where
    show (GlobFuncDef bbid na []) = "GFD at BB " ++ show bbid ++ " (" ++ show na ++ ")"
    show (GlobFuncDef bbid na cs) =
        "GFD at BB " ++ show bbid ++ " (" ++ show na ++ ") (closure slots: " ++ intercalate ", " cs ++ ")"

instance Show BB where
    show = genericShowBB (const "") show show

instance Show Ref where
    show (RConst n) = show n
    show (RTemp n) = "t" ++ show n
    show (RSClo name) = "SC(\"" ++ name ++ "\")"
    show RNone = "<<NONE>>"

instance Show InsCode where
    show (IAssign r) = "assign " ++ show r
    show (IParam n) = "param " ++ show n
    show (IClosure n) = "closure " ++ show n
    show (IData n) = "data " ++ show n
    show (ICallC r as) = "callc " ++ show r ++ " " ++ show as
    show (IAllocClo name vs) = "alloc-closure \"" ++ name ++ "\" " ++ show vs
    show (IDiscard r) = "discard " ++ show r
    show (IPush rs) = "push " ++ show rs
    show (IPop rs) = "pop " ++ show rs

instance Show Terminator where
    show (IBr r b1 b2) = "br " ++ show r ++ " " ++ show b1 ++ " " ++ show b2
    show (IJmp b) = "jmp " ++ show b
    show (IRet r) = "ret " ++ show r
    show (ITailC r as) = "tailc " ++ show r ++ " " ++ show as
    show IExit = "exit"
    show IUnknown = "<<UNKNOWN>>"

genericShowBB :: (Int -> String) -> (InsCode -> String) -> (Terminator -> String) -> BB -> String
genericShowBB bbannot icshow termshow (BB i inss term) =
    "BB " ++ show i ++ (case bbannot i of { "" -> "" ; s -> "   ; " ++ s }) ++
    concatMap (\(r, ic) -> case r of
        RNone -> "\n  " ++ icshow ic
        _ -> "\n  " ++ show r ++ " <- " ++ icshow ic) inss ++
    "\n  " ++ termshow term