module Stackify (
    stackify
) where

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

import AST (Name)
import Intermediate
import Liveness


-- Note about stackification. Temporaries need to be pushed before a call
-- if they're live after it and they could, conceivably, be wrongly mutated
-- by the called function otherwise. This has a couple of interesting
-- consequences:
-- 1. No temporaries ever need to be pushed from the global context, or
--    "main function", since they can never be wrongly mutated: that
--    requires re-entering the function in which they were defined, but one
--    can never re-enter the "main function".
-- 2. No temporaries ever need to be pushed before a tail call; since no
--    local variables are live after it (obviously). (Global variables are
--    covered by point (1.).)
stackify :: IRProgram -> IRProgram
stackify (IRProgram origbbs gfds datas) =
    let (mainCode, funcCodes) = partitionFunctions gfds origbbs
        resbbs = concatMap stackifyF (mainCode : map snd funcCodes)
    in optimise (IRProgram resbbs gfds datas)
  where
    stackifyF :: [BB] -> [BB]
    stackifyF bbs =
        let livemap = liveness bbs
            bbs' = [stackifyBB (livemap Map.! bid) bb | bb@(BB bid _ _) <- bbs]
        in bbs'

    stackifyBB :: [(Set.Set Int, Set.Set Int)] -> BB -> BB
    stackifyBB live (BB bid inss term) =
        -- Note that no temporaries need to be pushed for a tail call, so
        -- never at all for a terminator.
        BB bid (concatMap (uncurry stackifyIns) (zip live inss)) term

    stackifyIns :: (Set.Set Int, Set.Set Int) -> Instruction -> [Instruction]
    stackifyIns (liveBefore, liveAfter) orig@(d, ins) = case ins of
        ICallC r rs   -> let refs = map RTemp (Set.toList (liveBefore `Set.intersection` liveAfter))
                         in if null refs
                                then [(d, ICallC r rs)]
                                else [(RNone, IPush refs)
                                     ,(d,     ICallC r rs)
                                     ,(RNone, IPop refs)]
        IAssign _     -> [orig]
        IParam _      -> [orig]
        IClosure _    -> [orig]
        IData _       -> [orig]
        IAllocClo _ _ -> [orig]
        IDiscard _    -> [orig]
        IPush _       -> error "Unexpected stack operation before Stackify"
        IPop _        -> error "Unexpected stack operation before Stackify"

optimise :: IRProgram -> IRProgram
optimise (IRProgram bbs gfds datas) =
    IRProgram [BB bid (go inss) term | BB bid inss term <- bbs] gfds datas
  where
    go [] = []
    go ((_, IPop l1) : (_, IPush l2) : inss) | l1 == l2 = go inss
    go (ins : inss) = ins : go inss

-- Returns the global main code, and for each function its local code.
partitionFunctions :: Map.Map Name GlobFuncDef -> [BB] -> ([BB], [(GlobFuncDef, [BB])])
partitionFunctions gfds bbs =
    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
        mainReach = floodFill bbmap 0
        pairs = [(gfd, floodFill bbmap bid) | gfd@(GlobFuncDef bid _ _) <- Map.elems gfds]
        accums = scanl (<>) mainReach (map snd pairs)
        triples = zip3 accums pairs (tail accums)
    in if all (\(from, (_, bbset), to) -> Set.size from + Set.size bbset == Set.size to) triples
            && Set.size (last accums) == length bbs
           then (map (bbmap Map.!) (Set.toList mainReach)
                ,[(gfd, map (bbmap Map.!) (Set.toList bbset)) | (gfd, bbset) <- pairs])
           else error "Non-partitionable BBs in partitionFunctions"
  where
    floodFill bbmap origin = go (Set.singleton origin) origin
      where
        go seen at = let newseen = seen <> Set.fromList (outEdges (bbmap Map.! at))
                     in foldl go newseen (Set.toList (newseen Set.\\ seen))

-- Returns, for each BB, for each instruction the set of temporaries live
-- before and after that instruction.
liveness :: [BB] -> Map.Map Int [(Set.Set Int, Set.Set Int)]
liveness bbs =
    let sets = livenessAnalysis bbs bidOf itemsOf outEdges fread fwrite
    in Map.fromList (zip (map bidOf bbs) sets)
  where
    itemsOf (BB _ inss term) = map Right inss ++ [Left term]
    fread (Right (_, ins)) = icReadTemps ins
    fread (Left term) = termReadTemps term
    fwrite (Right (r, _)) = onlyTemporaries [r]
    fwrite (Left _) = []