{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Make.Graph
-- Copyright   :  (c) 2008, Andrea Vezzosi
-- License     :  BSD3
--
-- Maintainer  :  cabal-devel@haskell.org
-- Portability :  portable
--
-- This module defines the pure model behind the dep. analysis framework.


module Make.Graph 
    (
     Graph()
    ,sources
    ,reprs
    ,insRules
    ,insReprs
    ,emptyGraph
    ,toCache
    )
        where

import Make.Rule
import Data.Graph.Inductive hiding (Graph,apply)
import qualified Data.Map as M
import qualified Data.Set as S
import Data.List
import Data.Maybe
import Control.Arrow
import Control.Exception (assert)
import Debug.Trace
import Make.App
import Control.Applicative hiding (empty)
import Control.Monad.State
-- TODO 
-- * record from where dependencies have been generated so that we can rollback them.
-- * handle the case where a new rules is generated for a target already covered, 
--   especially when the sets of targets don't exactly match.

-- | The Dependency Graph, nodes correspond to Rules.  It keeps track
-- of already traversed nodes (see markNode), so we can retrieve which
-- rules are ready to run.  A Graph is always complete, i.e. every
-- mentioned target has a rule to generate it, a target is mentioned
-- if it appears either as a static or dynamic dep, or as a target for
-- a rule.
data Graph m target repr = G {
      targetToRule :: (M.Map target Node), 
      ruleToNode :: (M.Map (Rule m target repr) Node),
      nextNode :: Int,
      graph :: (Gr (Rule m target repr) target), -- ^ nodes are never removed from the graph
      past :: (S.Set Node), -- ^ already processed nodes
      focus :: (S.Set Node), -- ^ nodes that depend only on already processed ones
      announced :: (S.Set Node), -- ^ subset of focus, the nodes already announced as ready by sources
      exprs :: M.Map Node (App (Rule m) target repr (m [(target,repr)])),
    -- ^ map from nodes (i.e. set of targets) to partially evaluated expressions defining them.
      reprs :: M.Map target repr -- ^ stores the applied representations.
      
    } deriving (Show)

invariants :: (Ord target) => Graph m target repr -> Bool
invariants g@G{graph=gr} = and . zipWith number [1..] $
   [-- rules ready to run are disjoint from rules already runned
    past g `S.intersection` focus g == S.empty 
    -- ruleToNode and graph agree on the Node <-> Rule bijection
   ,M.fromList (map swap (labNodes (graph g))) == ruleToNode g
    -- each rule depend on the target from which it can generate dynamic deps
--   ,all (\(t,rs) -> map fst rs `subset` map (fromJust . lab gr) (suc gr (toNode g (getRule g t))) ) (M.toList (dyndep g)) 
    -- nextNode is the next available node
   ,[nextNode g] == newNodes 1 gr 
    -- static dependencies are stored
   ,all (\(deps,static) -> static `subset` deps) . map ((map snd . lpre gr) *** staticdeps) . labNodes $ gr 
    -- a rule is dependent-on only via targets it produces.
   ,all (uncurry subset) . map ((map snd . lsuc gr) *** targets) . labNodes $ gr 
    -- sources depend only on already runned nodes
   ,S.fromList [dep | s <- S.toList (focus g), dep <- pre gr s] `S.isSubsetOf` past g
    -- targetToRule matches with the information stored in the rules.
   ,M.fromList [(t,n)| (n,r) <- labNodes gr, t <- targets r] ==  targetToRule g
    -- past nodes are present in the graph.
   ,past g `S.isSubsetOf` S.fromList (map fst (labNodes gr))
    -- a target has a stored representation iff it's produced by a past node
   ,(S.fromList . map (getRule g) . M.keys . reprs $ g) == past g
    -- announced is a subset of focus
   ,announced g `S.isSubsetOf` focus g
   ]
    where 
          swap (x,y) = (y,x)
          number n True = True
          number n False = trace (show n) False
          subset xs ys = intersect xs ys == xs



emptyGraph :: (Ord target) => Graph m target repr
emptyGraph = s assert G {
                 targetToRule = M.empty, 
                 ruleToNode = M.empty,
                 nextNode = 1,
                 graph = empty,
                 past = S.empty,
                 focus = S.empty,
                 announced = S.empty,
                 exprs = M.empty,
                 reprs = M.empty
               }

-- | Add Rules to the Graph
insRules :: (Functor m1, MonadState (Graph m target repr) m1, Ord target) =>
            [Rule m target repr] -> m1 ()
insRules xs = mapM_ addRule xs >> modify (\g -> s assert g)

-- | Expands the dynamic dependencies using the provided representations.
insReprs xs = mapM_ (uncurry addRepr) xs

-- | Rules ready to run (i.e. their deps have been marked as
-- | processed) and the targets they depend on.
sources :: (MonadState (Graph t t1 t2) m, Ord t1) =>
           m [((Rule t t1 t2, t [(t1, t2)]), [(t1, t2)])]
sources = do g@G{graph=gr} <- get
             sources' g 
    where 
      sources' g@G{graph=gr} = do put g { announced = announced g `S.union` toAnnounce } -- i.e. announced == focus
                                  return .
                                         map (first $ (id &&& getAction g)) . 
                                             map (ruleAndDeps g) . S.toList $ toAnnounce
          where
            toAnnounce = focus g `S.difference` announced g
      getAction g r = let n = toNode g r
                          a = (exprs g) M.! n
                      in runApp a

      

-- | Get back the targets with their computed rapresentation and dependencies.
-- TODO this should be optimized
toCache :: (Ord t) => Graph m t r -> M.Map t (r, [t])
toCache g = M.fromList . concatMap toTargets . 
            map (ruleAndDeps g) . S.toList . past $ g
    where 
      toTargets (r,deps) = [ (t,(getRepr g t,map fst deps)) | t <- targets r ]






addRepr :: (Ord target, Functor m, MonadState (Graph m1 target repr) m) =>
           target -> repr -> m ()
addRepr t r = do
  modify $ \g -> g { reprs = M.insert t r (reprs g) }
  markTarget t
  ns <- dependOn t
  mapM_ (\n -> applyG n t r >>= registerDeps n) ns
  modify $ \g -> s assert g

dependOn :: (Ord target, MonadState (Graph m target r) m1) =>
            target -> m1 [Node]
dependOn t = do 
  g <- get
  let n = getRule g t
  return . map fst . filter ((==t) . snd) $ (lsuc . graph) g n 
  
addRule :: (Functor m1, MonadState (Graph m target repr) m1, Ord target) =>
           Rule m target repr -> m1 ()
addRule r = do 
  g <- get
  let 
      newg = g { nextNode = n+1
                 ,focus = S.insert n (focus g) 
                 ,ruleToNode = M.insert r n (ruleToNode g)
                 ,targetToRule = foldl' (\m t -> M.insert t n m) (targetToRule g) (targets r)
                 ,graph = insNode (n,r) (graph g)
                 ,exprs = M.insert n (expr . action $ r) (exprs g)
                 }
      n = nextNode g
  case M.lookup r (ruleToNode g) of
    Nothing -> do put newg
                  registerDeps n (collect . action $ r)
    Just _  -> return ()

registerDeps :: (Functor m1, MonadState (Graph m target repr) m1, Ord target) =>
                Node -> [(target, Rule m target repr)] -> m1 ()
registerDeps _ [] = return ()
registerDeps r xs = do mapM (addRule . snd) xs
                       modify $ addEdges xs r
                       rs <- catMaybes <$> mapM ((\t -> fmap ((,) t) <$> getReprM t) . fst) xs
                       newdeps <- concat <$> mapM (uncurry $ applyG r) rs
                       registerDeps r newdeps

addEdges :: (Ord t) => [(t, Rule m t r)] -> Node -> Graph m t r -> Graph m t r
addEdges deps to g = g { focus = newf, graph = insEdges [(dep,to,t) | (t,dep) <- from] (graph g) }
    where newf = case filter (`S.notMember` (past g)) (map snd from) of
                   [] -> focus g
                   _  -> S.delete to (focus g) 
          from = map (second $ toNode g) deps

-- | Marks the target as processed, so that it no longer blocks dependent rules.
markTarget :: (MonadState (Graph m1 target repr) m, Ord target) =>
              target -> m ()
markTarget t = modify $ \g@G{targetToRule=m,graph=gr} ->
    case M.lookup t m of
      Nothing -> g -- not there
      Just n | n `S.member` past g -> g -- already marked
             | otherwise -> 
                 g { past = p'
                   , focus = (S.delete n (focus g)) `S.union` 
                     (S.fromList . filter pred . nub . suc gr $ n)
                   , announced = S.delete n (announced g)
                   }
             where 
               pred = all (`S.member` p') . pre gr
               p' = S.insert n (past g)


ruleAndDeps :: (Ord t) => Graph m t r -> Node -> (Rule m t r, [(t, r)])
ruleAndDeps g@G{graph=gr} n = (r, map (id &&& getRepr g) . map fst $ deps)
    where
      (deps,_,r,_) = context gr n 


applyG :: (Eq t, MonadState (Graph m1 t r) m) =>
          Node -> t -> r -> m [(t, Rule m1 t r)]
applyG rule t r = do g@G{exprs=e} <- get 
                     let (ts,expr) = apply t r (e M.! rule)
                     put $ g{exprs=M.insert rule expr e}
                     return ts



-- Utilities

getReprM :: (MonadState (Graph m target repr) m1, Functor m1, Ord target) =>
            target -> m1 (Maybe repr)
getReprM t = do m <- reprs <$> get
                return $ M.lookup t m `asTypeOf` Nothing

getRepr :: (Ord t) => Graph m t r -> t -> r
getRepr g t = M.findWithDefault (error "Graph.getRepr: repr not found.") t (reprs g)

getRule :: (Ord t) => Graph m t r -> t -> Node
getRule g t = maybe (error $ "Graph.getRule: target not found") id (M.lookup t (targetToRule g))

toNode :: (Ord t) => Graph m t r -> Rule m t r -> Node
toNode g r = M.findWithDefault (error "Graph.toNode: Rule not found.") r (ruleToNode g)

fromNode :: Graph m target repr -> Node -> Rule m target repr
fromNode g n = fromMaybe (error $ "Graph.fromNode: Node not found " ++ show n) (lab (graph g) n)

s :: (Ord target) =>  (Bool -> Graph m target repr -> t) -> Graph m target repr -> t
s a g = a (invariants g) g