module Stackify (
    stackify
) where

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

import Intermediate


data Funcinfo =
    FuncInfo { fiGFD :: GlobFuncDef
             , fiBBs :: [BB]
             , fiTemps :: [Int] }

stackify :: IRProgram -> IRProgram
stackify (IRProgram origbbs gfds datas) =
    let infos = [FuncInfo gfd bbs [i | RTemp i <- allRefs bbs]
                | (gfd, bbs) <- partitionFunctions gfds origbbs]
    -- LIVENESS ANALYSIS within the function!

partitionFunctions :: Map.Map Name GlobFuncDef -> [BB] -> [(GlobFuncDef, [BB])]
partitionFunctions gfds bbs =
    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
        pairs = [(gfd, floodFill bbmap bid) | gfd@(GlobFuncDef bid _ _) <- Map.elems gfds]
        accums = scanl (<>) Set.empty (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 [(gfd, map (bbmap Map.!) (Set.toList bbset)) | (gfd, bbset) <- pairs]
           else error "Non-partitionable BBs in partitionFunctions"
  where
    floodFill bbmap origin = go Set.empty origin
      where
        go seen at = let newseen = seen <> Set.fromList (outEdges (bbmap Map.! at))
                     in foldl go newseen (Set.toList (newseen Set.\\ seen))