content: {-# LANGUAGE RecordWildCards #-}
  {-|
  Module      : Parsley.Internal.Frontend.Analysis.Dependencies
  Description : Calculate dependencies of a collection of bindings.
  License     : BSD-3-Clause
  Maintainer  : Jamie Willis
  Stability   : experimental
  
  Exposes `dependencyAnalysis`, which is used to calculate information
  regarding the dependencies of each let-bound parser, as well as their
  free-registers.
  
  @since 1.5.0.0
  -}
  module Parsley.Internal.Frontend.Analysis.Dependencies (dependencyAnalysis) where
  
  import Control.Monad                        (unless, forM_)
  import Control.Monad.ST                     (ST)
  import Data.Array                           (Array)
  import Data.Array.MArray                    (readArray, writeArray, newArray, newArray_)
  import Data.Array.ST                        (STArray, runSTUArray, runSTArray)
  import Data.Array.Unboxed                   (UArray)
  import Data.Bifunctor                       (first, second)
  import Data.Dependent.Map                   (DMap)
  import Data.Foldable                        (foldl')
  import Data.Map.Strict                      (Map)
  import Data.Set                             (Set)
  import Data.STRef                           (newSTRef, readSTRef, writeSTRef)
  import Parsley.Internal.Common.Indexed      (Fix, cata, Const1(..), (:*:)(..), zipper)
  import Parsley.Internal.Common.State        (State, MonadState, execState, modify')
  import Parsley.Internal.Core.CombinatorAST  (Combinator(..), traverseCombinator)
  import Parsley.Internal.Core.Identifiers    (IMVar, MVar(..), ΣVar, SomeΣVar(..))
  
  import qualified Data.Dependent.Map as DMap   (foldrWithKey, filterWithKey)
  import qualified Data.Map.Strict    as Map    ((!), empty, insert, findMax, elems, maxView, fromList, fromDistinctAscList)
  import qualified Data.Set           as Set    (toList, insert, union, unions, member, notMember, empty, (\\), fromDistinctAscList, size)
  import qualified Data.Array         as Array  ((!), listArray, bounds, indices)
  import qualified Data.Array.Unboxed as UArray ((!), assocs)
  import qualified Data.List          as List   (partition)
  
  type Graph = Array IMVar [IMVar]
  type PQueue k a = Map k a
  
  {-|
  Given a top-level parser and a collection of its let-bound subjects performs the following tasks:
  
  * Determines which parser depend on which others.
  * Use the previous information to remove any dead bindings.
  * Calculate the direct free registers for each binding.
  * Propogate the free registers according to transitive need via the dependency graph.
  
  Returns the non-dead bindings, the information about each bindings free registers, and the next
  free index for any registers created in code generation.
  
  @since 1.5.0.0
  -}
  -- TODO This actually should be in the backend... dead bindings and the topological ordering can be computed here
  --      but the register stuff should come after register optimisation and instruction peephole
  dependencyAnalysis :: Fix Combinator a -> DMap MVar (Fix Combinator) -> (DMap MVar (Fix Combinator), Map IMVar (Set SomeΣVar))
  dependencyAnalysis toplevel μs =
    let -- Step 1: find roots of the toplevel
        roots = directDependencies toplevel
        -- Step 2: build immediate dependencies
        DependencyMaps{..} = buildDependencyMaps μs
        -- Step 3: Build a call graph
        callees = buildGraph immediateDependencies
        -- Step 4: traverse the call graph, finding unreachable nodes and establishing a topological dfnum for each node
        (dfnums, lives, dead) = topoOrdering roots callees
        -- Step 5: reverse the call graph to make a callers graph
        callers = invertGraph callers dead
        -- Step 6: iterate over the live registers, and propagate the free registers until fix-point
        regs = propagateRegs lives dfnums usedRegisters definedRegisters callees callers
    in (DMap.filterWithKey (\(MVar v) _ -> Set.notMember v dead) μs, regs)
  
  buildGraph :: Map IMVar (Set IMVar) -> Graph
  buildGraph deps = Array.listArray (0, fst (Map.findMax deps)) (map Set.toList (Map.elems deps))
  
  invertGraph :: Graph -> Set IMVar -> Graph
  invertGraph g unreachable = runSTArray $ do
    g' <- newArray (Array.bounds g) []
    forM_ (Array.indices g) $ \n ->
      unless (Set.member n unreachable) $
        forM_ (g Array.! n) $ \s -> do
          preds <- readArray g' s
          writeArray g' s (n : preds)
    return g'
  
  type DFNum = Int
  topoOrdering :: Set IMVar -> Graph -> (UArray IMVar DFNum, [IMVar], Set IMVar)
  topoOrdering roots graph =
    let dfnums :: UArray IMVar DFNum
        dfnums = runSTUArray $ do
          dfnums <- newArray (Array.bounds graph) 0
          nextDfnum <- newSTRef 1
          let dfs v = do seen <- (/= 0) <$> readArray dfnums v
                         unless seen $
                           do dfnum <- readSTRef nextDfnum
                              writeArray dfnums v dfnum
                              writeSTRef nextDfnum (dfnum + 1)
                              forM_ (graph Array.! v) dfs
          -- Assign a DFNum to each IMVar
          forM_ roots dfs
          return dfnums
        -- if something still has dfnum 0, it was not visited, and is dead
        lives, deads :: [(IMVar, DFNum)]
        (lives, deads) = List.partition ((/= 0) . snd) (UArray.assocs dfnums)
        -- The DFNums are unique
    in (dfnums, map fst lives, Set.fromDistinctAscList (map fst deads))
  
  propagateRegs :: [IMVar] -> UArray IMVar DFNum -> Map IMVar (Set SomeΣVar) -> Map IMVar (Set SomeΣVar) -> Graph -> Graph -> Map IMVar (Set SomeΣVar)
  propagateRegs reachables dfnums uses defs callees callers = toMap $ runSTArray $
    do freeRegs <- newArray_ (Array.bounds callees)
       forM_ reachables $ \v -> writeArray freeRegs v ((uses Map.! v) Set.\\ (defs Map.! v))
       let worklist = Map.fromList (map (\v -> (dfnums UArray.! v, v)) reachables)
       maybe (return ()) (unfoldM_ (uncurry (propagate freeRegs))) (Map.maxView worklist)
       return freeRegs
    where
      propagate :: STArray s IMVar (Set SomeΣVar) -> IMVar -> PQueue DFNum IMVar -> ST s (Maybe (IMVar, PQueue DFNum IMVar))
      propagate freeRegs v work = do
        !frees        <- readArray freeRegs v
        !freesCallees <- Set.unions <$> traverse (readArray freeRegs) (callees Array.! v)
        let !frees'   =  frees `Set.union` (freesCallees Set.\\ (defs Map.! v))
        if Set.size frees /= Set.size frees' then do
          writeArray freeRegs v frees'
          return (Map.maxView (addWork (callers Array.! v) work))
        else return (Map.maxView work)
  
      addWork :: [IMVar] -> PQueue DFNum IMVar -> PQueue DFNum IMVar
      addWork vs work = foldl' (flip (\v -> Map.insert (dfnums UArray.! v) v)) work vs
  
      unfoldM_ :: Monad m => (s -> m (Maybe s)) -> s -> m ()
      unfoldM_ f s = f s >>= mapM_ (unfoldM_ f)
  
      toMap arr = Map.fromDistinctAscList (map (\v -> (v, arr Array.! v)) reachables)
  
  -- IMMEDIATE DEPENDENCY MAPS
  data DependencyMaps = DependencyMaps {
    usedRegisters         :: !(Map IMVar (Set SomeΣVar)), -- Leave Lazy
    immediateDependencies :: !(Map IMVar (Set IMVar)), -- Could be Strict
    definedRegisters      :: !(Map IMVar (Set SomeΣVar))
  }
  
  buildDependencyMaps :: DMap MVar (Fix Combinator) -> DependencyMaps
  buildDependencyMaps = DMap.foldrWithKey (\(MVar v) p deps@DependencyMaps{..} ->
    let (uses, defs, ds) = freeRegistersAndDependencies v p
    in deps { usedRegisters = Map.insert v uses usedRegisters
            , immediateDependencies = Map.insert v ds immediateDependencies
            , definedRegisters = Map.insert v defs definedRegisters}) (DependencyMaps Map.empty Map.empty Map.empty)
  
  freeRegistersAndDependencies :: IMVar -> Fix Combinator a -> (Set SomeΣVar,  Set SomeΣVar, Set IMVar)
  freeRegistersAndDependencies v p =
    let frsm :*: depsm = zipper freeRegistersAlg (dependenciesAlg (Just v)) p
        (uses, defs) = runFreeRegisters frsm
        ds = runDependencies depsm
    in (uses, defs, ds)
  
  -- DEPENDENCY ANALYSIS
  newtype Dependencies a = Dependencies { doDependencies :: State (Set IMVar) () }
  runDependencies :: Dependencies a -> Set IMVar
  runDependencies = flip execState Set.empty . doDependencies
  
  directDependencies :: Fix Combinator a -> Set IMVar
  directDependencies = runDependencies . cata (dependenciesAlg Nothing)
  
  {-# INLINE dependenciesAlg #-}
  dependenciesAlg :: Maybe IMVar -> Combinator Dependencies a -> Dependencies a
  dependenciesAlg (Just v) (Let μ@(MVar u)) = Dependencies $ do unless (u == v) (dependsOn μ)
  dependenciesAlg Nothing  (Let μ)          = Dependencies $ do dependsOn μ
  dependenciesAlg _ p                       = Dependencies $ do traverseCombinator (fmap Const1 . doDependencies) p; return ()
  
  dependsOn :: MonadState (Set IMVar) m => MVar a -> m ()
  dependsOn (MVar v) = modify' (Set.insert v)
  
  -- FREE REGISTER ANALYSIS
  newtype FreeRegisters a = FreeRegisters { doFreeRegisters :: State (Set SomeΣVar, Set SomeΣVar) () }
  runFreeRegisters :: FreeRegisters a -> (Set SomeΣVar, Set SomeΣVar)
  runFreeRegisters = flip execState (Set.empty, Set.empty) . doFreeRegisters
  
  {-# INLINE freeRegistersAlg #-}
  freeRegistersAlg :: Combinator FreeRegisters a -> FreeRegisters a
  freeRegistersAlg (GetRegister σ)      = FreeRegisters $ do uses σ
  freeRegistersAlg (PutRegister σ p)    = FreeRegisters $ do uses σ; doFreeRegisters p
  freeRegistersAlg (MakeRegister σ p q) = FreeRegisters $ do defs σ; doFreeRegisters p; doFreeRegisters q
  freeRegistersAlg p                    = FreeRegisters $ do traverseCombinator (fmap Const1 . doFreeRegisters) p; return ()
  
  uses :: MonadState (Set SomeΣVar, vs) m => ΣVar a -> m ()
  uses σ = modify' (first (Set.insert (SomeΣVar σ)))
  
  defs :: MonadState (vs, Set SomeΣVar) m => ΣVar a -> m ()
  defs σ = modify' (second (Set.insert (SomeΣVar σ)))
  