{-# LANGUAGE RankNTypes #-}

module LifetimeAnalysis(lifetimeAnalysis, Access(..), unAccess) where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Maybe
import Data.Vector ((!))
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV

import Utils


data Access a = Write a | Read a
  deriving (Show, Eq)

unAccess :: Access a -> a
unAccess (Write x) = x
unAccess (Read x) = x

type BB a = ([[Access a]], [Int])

type Mark = V.Vector (V.Vector Bool)
type MutMark s = MV.MVector s (MV.MVector s Bool)

lifetimeAnalysis :: Eq a => a -> [BB a] -> [[Bool]]
lifetimeAnalysis target bbs = V.toList $ V.map V.toList $
    foldOr2 $ map (\p -> markAliveFrom bbs target p (emptyMark bbs)) occurrences
  where
    occurrences :: [(Int, Int)]
    occurrences = do
        (i, bb) <- zip [0..] bbs
        (j, ins) <- zip [0..] (fst bb)
        if ins `contains` Write target || ins `contains` Read target
            then [(i, j)]
            else []
    
emptyMark :: [BB a] -> Mark
emptyMark bbs = V.fromList $ flip map bbs $ \(inss, _) -> V.replicate (length inss) False

-- Assumes `target` is known to be alive at `pos`.
markAliveFrom :: Eq a => [BB a] -> a -> (Int, Int) -> Mark -> Mark
markAliveFrom bbs target pos topmark = runST $ do
    maybem <- goNoCheck pos topmark
    m <- case maybem of
            Nothing -> thaw2 topmark
            Just m' -> return m'
    setAt2 m pos True
    freeze2 m
  where
    goNoCheck :: (Int, Int) -> MutMark s -> ST s (Maybe (MutMark s))
    goNoCheck (i, j) mark = do
        let suc = flip filter (successors bbs (i, j)) $ \(i', j') ->
                      not $ mark ! i' ! j'
        setAt2 mark (i, j) True
        res <- mapM (\s -> go s mark) suc
        case catMaybes res of
            [] -> return Nothing
            mutmarks -> Just <$> mvFoldOr2 mutmarks

    go :: (Int, Int) -> MutMark s -> ST s (Maybe (MutMark s))
    go (i, j) mark
        | fst (bbs !! i) !! j `contains` Write target = return Nothing
        | fst (bbs !! i) !! j `contains` Read target = Just $ setAt2 mark (i, j) True
        | otherwise = goNoCheck (i, j) mark

successors :: [BB a] -> (Int, Int) -> [(Int, Int)]
successors bbs (i, j) =
    let (inss, nexts) = bbs !! i
    in if j < length inss - 1
        then [(i, j + 1)]
        else [(n, 0) | n <- nexts]

setAt2 :: MutMark s -> (Int, Int) -> Bool -> ST s ()
setAt2 l (i, j) v = do
    row <- MV.read l i
    MV.write row j v

foldOr2 :: [Mark] -> Mark
foldOr2 marks =
    let (dim1, dim2) = (V.length (head marks), V.length (head marks ! 0))
    in V.generate dim1 $ \i -> V.generate dim2 $ \j ->
            foldl1 (||) $ map (\m -> m ! i ! j) marks

mvFoldOr2 :: [MutMark s] -> ST s (MutMark s)
mvFoldOr2 marks = do
    let dim1 = MV.length (head marks)
    dim2 <- liftM MV.length $ MV.read (head marks) 0
    mvGenerate dim1 $ \i -> mvGenerate dim2 $ \j ->
        foldl1 (||) <$> mapM (\m -> MV.read m i >>= \row -> MV.read row j) marks

thaw2 :: Mark -> ST s (MutMark s)
thaw2 = join . fmap V.thaw . V.mapM V.thaw

freeze2 :: MutMark s -> ST s Mark
freeze2 = join . fmap (V.mapM V.freeze) . V.freeze

mvGenerate :: Int -> (Int -> ST s a) -> ST s (MV.MVector s a)
mvGenerate sz f = do
    v <- MV.new sz
    mapM_ (\i -> f i >>= \value -> MV.write v i value) [0..sz-1]
    return v