content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ViewPatterns #-}
  module WTO (
    wtoTests
    ) where
  
  import Control.Applicative
  import Control.Monad ( replicateM, unless )
  import qualified Control.Monad.State.Strict as St
  import qualified Data.Foldable as F
  import qualified Data.Map.Strict as M
  import qualified Data.Set as S
  
  import Prelude
  
  import qualified Test.QuickCheck as QC
  import qualified Test.Tasty as T
  import qualified Test.Tasty.QuickCheck as T
  
  import Lang.Crucible.Analysis.Fixpoint.Components
  
  wtoTests :: T.TestTree
  wtoTests = T.testGroup "WeakTopologicalOrdering" [
    T.testProperty "prop_reachableInWTO" prop_reachableInWTO,
    T.testProperty "prop_validWTO" prop_validWTO
    ]
  
  -- | Test that all reachable nodes in the graph are present in the
  -- weak topological ordering
  prop_reachableInWTO :: RandomGraph -> Bool
  prop_reachableInWTO gp = dfsReachedNodes dfs == S.fromList (concatMap F.toList wto)
    where
      dfs = reachable gp
      (root, sf) = toCFG gp
      wto = weakTopologicalOrdering sf root
  
  -- | Test that the weak topological ordering property holds for the
  -- WTO we compute.
  --
  -- That property is defined in terms of a relation w(c), which
  -- evaluates to the set of heads of the nested components containing
  -- the vertex c.
  --
  -- For every edge u->v:
  --
  --   (u < v AND v not in w(u)) OR (v <= u AND v in w(u))
  --
  -- where v <= u means that u->v is a backedge.
  prop_validWTO :: RandomGraph -> Bool
  prop_validWTO gp@(RG edges) =
    and [ (isBackEdge e && (v `vertexInWC` u)) || not (v `vertexInWC` u)
        | e@(u, v) <- edges
        , isReachableEdge e
        ]
    where
      dfs = reachable gp
      (root, sf) = toCFG gp
      wto = weakTopologicalOrdering sf root
      cchs = indexContainingComponentHeads wto
      isReachableEdge (src, _) = src `S.member` dfsReachedNodes dfs
      isBackEdge e = e `S.member` dfsBackEdges dfs
      vertexInWC v c = maybe False (S.member v) $ M.lookup c cchs
  
  -- | This is the w(c) relation from the WTO criteria
  --
  -- The map keys are the cs
  indexContainingComponentHeads :: (Ord n) => [WTOComponent n] -> M.Map n (S.Set n)
  indexContainingComponentHeads cs = St.execState (mapM_ (go []) cs) M.empty
    where
      go heads c =
        case c of
          Vertex v -> St.modify' $ M.insert v (S.fromList heads)
          SCC (SCCData { wtoHead = h
                       , wtoComps = cs'
                       }) -> do
            let heads' = h : heads
            St.modify' $ M.insert h (S.fromList heads')
            mapM_ (go heads') cs'
  
  newtype NodeId = NID Int
    deriving (Eq, Show)
  
  instance QC.Arbitrary NodeId where
    arbitrary = QC.sized mkNodeId
      where
        mkNodeId n = NID <$> QC.choose (0, n)
  
  newtype RandomGraph = RG [(Int, Int)]
    deriving (Show)
  
  instance QC.Arbitrary RandomGraph where
    arbitrary = QC.sized mkRandomGraph
  
  -- | Make an arbitrary graph by deciding on a number of edges and then
  -- generating random edges.  Note that we always increment the size so
  -- that we don't get empty graphs.
  --
  -- The graphs are not all connected.
  mkRandomGraph :: Int -> QC.Gen RandomGraph
  mkRandomGraph ((+1) -> sz) = do
    nEdges <- QC.choose (2, 2*sz)
    srcs <- replicateM nEdges (QC.choose (0, sz))
    dsts <- replicateM nEdges (QC.choose (0, sz))
    return $! RG (unique (zip srcs dsts))
  
  -- | A DFS result; it additionally contains a list of back edges
  -- discovered during its traversal
  data DFS = DFS { dfsReachedNodes :: S.Set Int
                 , dfsBackEdges :: S.Set (Int, Int)
                 , dfsStart :: S.Set Int
                 , dfsFinish :: S.Set Int
                 }
  
  -- | Compute the DFS of a random graph from its root.
  reachable :: RandomGraph -> DFS
  reachable gp =
    St.execState (go root) s0
    where
      s0 = DFS { dfsReachedNodes = S.empty
               , dfsBackEdges = S.empty
               , dfsStart = S.empty
               , dfsFinish = S.empty
               }
      (root, sf) = toCFG gp
      go n = do
        markDiscovered n
        F.forM_ (sf n) $ \successor -> do
          disc <- isDiscovered successor
          case disc of
            False -> go successor
            True -> do
              fin <- isFinished successor
              unless fin $ addBackedge (n, successor)
        markFinished n
  
  type M a = St.State DFS a
  
  markDiscovered :: Int -> M ()
  markDiscovered v = St.modify' $ \s ->
    s { dfsReachedNodes = S.insert v (dfsReachedNodes s)
      , dfsStart = S.insert v (dfsStart s)
      }
  
  markFinished :: Int -> M ()
  markFinished v = St.modify' $ \s ->
    s { dfsFinish = S.insert v (dfsFinish s) }
  
  addBackedge :: (Int, Int) -> M ()
  addBackedge be = St.modify' $ \s -> s { dfsBackEdges = S.insert be (dfsBackEdges s) }
  
  isDiscovered :: Int -> M Bool
  isDiscovered n = S.member n <$> St.gets dfsReachedNodes
  
  isFinished :: Int -> M Bool
  isFinished n = S.member n <$> St.gets dfsFinish
  
  unique :: (Ord a) => [a] -> [a]
  unique = S.toList . S.fromList
  
  -- | Return the root and the successor function for the graph.
  --
  -- Not defined for empty graphs
  toCFG :: RandomGraph -> (Int, (Int -> [Int]))
  toCFG (RG []) = error "Empty graph"
  toCFG (RG edges@((s0, _) : _)) =
    (s0, \n -> [ dst | (src, dst) <- edges, n == src])
  
  
  