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

import Data.List
import Data.Maybe


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])

lifetimeAnalysis :: Eq a => a -> [BB a] -> [[Bool]]
lifetimeAnalysis target bbs =
    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] -> [[Bool]]
emptyMark bbs = flip map bbs $ \(inss, _) -> map (const False) inss

-- Assumes `target` is known to be alive at `pos`.
markAliveFrom :: Eq a => [BB a] -> a -> (Int, Int) -> [[Bool]] -> [[Bool]]
markAliveFrom bbs target pos topmark = setAt2 (maybe topmark id $ goNoCheck pos topmark) pos True
  where
    goNoCheck :: (Int, Int) -> [[Bool]] -> Maybe [[Bool]]
    goNoCheck (i, j) mark =
        let suc = flip filter (successors bbs (i, j)) $ \(i', j') ->
                      not $ mark !! i' !! j'
            markset = setAt2 mark (i, j) True
        in case catMaybes [go s markset | s <- suc] of
               [] -> Nothing
               l -> Just $ foldOr2 l

    go :: (Int, Int) -> [[Bool]] -> Maybe [[Bool]]
    go (i, j) mark
        | fst (bbs !! i) !! j `contains` Write target = 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]

contains :: Eq a => [a] -> a -> Bool
contains l v = isJust $ find (== v) l

modifyAt2 :: [[a]] -> (Int, Int) -> (a -> a) -> [[a]]
modifyAt2 l (i, j) f = modifyAt l i $ \li -> modifyAt li j f

modifyAt :: [a] -> Int -> (a -> a) -> [a]
modifyAt l i f = let (pre, v : post) = splitAt i l in pre ++ f v : post

setAt2 :: [[a]] -> (Int, Int) -> a -> [[a]]
setAt2 l p v = modifyAt2 l p (const v)

foldOr2 :: [[[Bool]]] -> [[Bool]]
foldOr2 = foldl1 (\m1 m2 -> map (map (uncurry (||)) . uncurry zip) (zip m1 m2))