{-# LANGUAGE BangPatterns, FlexibleContexts #-}

module Cmm.Register.Allocation(
     allocateRegisters
   , generateAllocatedx86
    ) where 

import           Data.Set                   (Set)
import qualified Data.Set        as Set
import           Data.Map                   (Map)
import qualified Data.Map.Strict as Map

import           Control.Monad.IO.Class
import           Control.Monad.Loops        (maximumOnM)
import           Control.Lens               hiding ((#), none)
import           Control.Monad.Trans.State  hiding (State)
import           Control.Monad.Trans        (lift)

import           Control.Concurrent.ParallelIO.Global  (parallel)
import           Data.IORef

import           Text.Printf                (printf)
import           Data.List                  (foldl', find, maximumBy)
import           Data.Maybe                 (fromJust)
import           Data.Ord                   (comparing)
import           Data.Char                  (isAlphaNum)

import           AST                        (MiniJava())
import           Cmm.X86.Backend            (generatex86Gen)
import           Cmm.X86.InstrCore
import           Cmm.DirectedGraph
import           Cmm.Backend                (MachineInstr(..)
                                           , MachineFunction(..)
                                           , MachinePrg(..)
                                           , CodeGen(..))
import           Cmm.LabelGenerator     

import           Cmm.ControlFlowGraph       (createControlFlowGraph)
import           Cmm.ActivityAnalysis       (activityAnalysis, ActivityStorage(..))
import           Cmm.InterferenceGraph      (createInterferenceGraph)
import           Cmm.Register.Core      

-- | parallel register allocation
--
generateAllocatedx86 :: MiniJava -> Bool -> IO X86Prog
generateAllocatedx86 ast inParallel = do
    (x86prog, iorefstate) <- evalNameGenT (generatex86Gen ast)

--  yes, it's that easy
    let runComputation
            | inParallel = parallel
            | otherwise  = sequence

    functions <- runComputation $ map (go c iorefstate) (machinePrgFunctions x86prog)
    return $ replaceFunctions x86prog functions

  where c = X86CodeGen

        go :: X86CodeGen -> IORef ([Temp], [Label]) -> X86Func -> IO X86Func
        go c state f = runWithNameStateT state (allocateRegisters c f)

-- | approximates the graph coloring problem, spills the temps and returns a colored function function
--
allocateRegisters ::
  (CodeGen c p f i, Ord i, Show i, MonadIO m, Show f)
  => c
  -> f
  -> NameGenT m f
allocateRegisters c function = evalStateT (modifyFunction function) regState
    where
        ig :: DirectedGraph Temp
        ig = createInterferenceGraph function

        regState = RegisterState
            { _interferenceGraph = ig
            , _tempStates        = createDefaultTempStates ig
            , _colors            = generalPurposeRegisters c
            , _tempStack         = []
            }

modifyFunction :: (MonadNameGen m, MonadIO m, MachineFunction f i, Ord i, Show i, Show f) => f -> Reg m f
modifyFunction f = do
    usefulFunction <- deleteUnusedTemps f
    coloringPass
    spilled <- getAllSpilled
    case (not . none $ spilled) of
        True -> do
            newFunction <- lift $ machineFunctionSpill usefulFunction spilled  -- modifying the function with asm
            updateInterferenceGraph newFunction            -- create new interferencegraph
            deleteSpilledTemps spilled                     -- delete spilled from our local (temp -> state mapping)
            resetAllSpilledTemps                           -- if you didn't spill every possible temp, this would reset them accordignly
            addNewTempStates                               -- new temps are generated by spilling, we need to add them
            modifyFunction newFunction
        False -> insertRegisterColors f

-- | last step to inserting the new registers, we rename every instruction,
--   filter unused instructions, and allocate enough space for the frame
insertRegisterColors :: (MonadNameGen m, MonadIO m, MachineFunction f i, Ord i, Show i) => f -> Reg m f
insertRegisterColors f = do
    (machineFunctionStackAlloc .
         machineFunctionFilterInstructions .
         machineFunctionRenameByMap f) <$> createTempMapping
    
deleteUnusedTemps :: (MonadNameGen m, MonadIO m, MachineFunction f i, Ord i, Show i) => f -> Reg m f
deleteUnusedTemps f = do
    usedNodes <- nodes <$> (view interferenceGraph <$> get)
    return $ machineFunctionFilterUnusedMoves f usedNodes
    

coloringPass :: (MonadNameGen m, MonadIO m) => Reg m ()
coloringPass = do
    simplify
    mt <- findCurrentMaxChildrenNode
    case mt of
        Nothing  -> reverseStack >> select
        (Just t) -> moveToStack t >> coloringPass

simplify :: (MonadNameGen m, MonadIO m) => Reg m ()
simplify = do
    ns <- allNodes
    k  <- Set.size <$> (view colors <$> get)
    lessThanKNodes <- filterSM (((< k) <$>) . getOutDegree) ns
    mapSM_ moveToStack lessThanKNodes

-- | pick temps from the stack and give them a possible color,
--   if none are possible, mark as spilled
--   terminate when the stack is empty
select :: (MonadNameGen m, MonadIO m) => Reg m ()
select  = do
    mtemp <- getFromStack
    case mtemp of
        Nothing  -> do
            return ()
        (Just t@(NamedTemp color)) -> do
            tempStates %= Map.insert t (Colored t) -- named temps are already colored (this happens because 'ret' uses '%eax' and not a temp)
            select
        (Just temp) -> do
            children <- getColoredChildren temp
            mColor   <- getFreeColor children
            case mColor of
                Nothing      -> setSpilled temp >> select
                (Just color) -> do
                    tempStates %= Map.insert temp color
                    select

-- | get node which is not on the stack, pick the one with the maximum children, otherwise there are not 'Clear' nodes
findCurrentMaxChildrenNode ::  (MonadNameGen m, MonadIO m) => Reg m (Maybe Temp)
findCurrentMaxChildrenNode = do
    allNodes           <- nodes <$> (view interferenceGraph <$> get)
    allAccessableNodes <- filterSM (\n -> not <$> (onStack n)) allNodes
    case (not . none $ allAccessableNodes) of
        True ->  maximumOnM getOutDegree (Set.toList allAccessableNodes)
        False -> return Nothing

-- | check which colors are used, pick one if possible
getFreeColor :: (MonadNameGen m, MonadIO m) => Set Temp -> Reg m (Maybe State)
getFreeColor children = do
    validChildren <- filterSM isColored children
    usedColors    <- mapSM getColor validChildren
    allColors     <- view colors <$> get         
    let possibleColors = Set.difference allColors usedColors
    case none possibleColors of
        True ->  return Nothing
        False -> return . Just . Colored $ 0 `Set.elemAt` possibleColors