Shuffle a graph

ShuffleGraph.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Data.Graph.Shuffle where

-- | Shuffle a graph into a randomly sorted list, preserving
-- topological order.

{-

> quickCheckWith stdArgs {maxSuccess=10000} prop_identity
+++ OK, passed 10000 tests.
> quickCheckWith stdArgs {maxSuccess=10000} prop_unique_assignments
+++ OK, passed 10000 tests.
> quickCheckWith stdArgs {maxSuccess=10000} prop_toporder_preserved
+++ OK, passed 10000 tests.

-}

import           Control.Arrow ((&&&))
import           Control.Monad
import           Control.Monad.Random
import           Control.Monad.State.Strict
import           Data.Either
import qualified Data.Graph as G
import           Data.List
import qualified Data.Map.Strict as M
import           Data.Maybe
import           Data.Ord
import           Data.Tuple
import           Test.QuickCheck

--------------------------------------------------------------------------------
-- API

-- | Shuffle the graph, preserving topological order.
shuffleGraph ::
     forall k a. (Ord k) => [(a, k, [k])] -> StdGen -> ([(a, k, [k])], StdGen)
shuffleGraph xs = runRand (shuffleGraphM xs)

-- | Shuffle the graph, preserving topological order.
shuffleGraphM ::
     forall k a m. (Ord k, MonadRandom m)
  => [(a, k, [k])]
  -> m [(a, k, [k])]
shuffleGraphM xs =
  fmap (map (snd . snd) . sortBy (comparing fst)) (randomAssignGraph xs)

-- | Randomly assign positions for nodes in the graph, preserving the
-- topological order for reachable nodes, but otherwise interspersing
-- non-reachable ones.
randomAssignGraph ::
     forall k a m. (Ord k, MonadRandom m)
  => [(a, k, [k])]
  -> m [(Int, (Either Int Int, (a, k, [k])))]
randomAssignGraph nodes =
  fmap
    snd
    (mapAccumM
       assign
       mempty
       (map Left (lefts topSorted) <> map Right (rights topSorted)))
  where
    (topSorted, vertexToNode) = topSortTransposedGraph nodes
    assign hash eitherVertex = do
      idx <- getRandomR (lower + 1, upper + 1)
      let hash' = M.insert key idx hash
      pure (hash', (idx, (eitherVertex, node)))
      where
        lower = foldl' max 0 (mapMaybe (flip M.lookup hash) deps)
        upper = M.foldl max (lower + length nodes) hash
        node@(_, key, deps) = vertexToNode vertex
        vertex =
          case eitherVertex of
            Left v -> v
            Right v -> v

-- | Make a graph, transpose it, toplogically sort it, then split into
-- Left reachable, Right unreachable.
topSortTransposedGraph ::
     Ord key
  => [(node, key, [key])]
  -> ([Either G.Vertex G.Vertex], G.Vertex -> (node, key, [key]))
topSortTransposedGraph nodes = (topSorted, vertexToNode)
  where
    (graph, vertexToNode, _) = G.graphFromEdges nodes
    topSorted = topSort graph
    topSort :: G.Graph -> [Either G.Vertex G.Vertex]
    topSort =
      map
        (\key ->
           if unreachable graph key
             then Right key
             else Left key) .
      G.topSort . G.transposeG
    unreachable g = null . G.reachable g

--------------------------------------------------------------------------------
-- Tests

-- | Identity modulo order (no duplicates or removals).
prop_identity :: ArbitraryGraph -> Int -> Bool
prop_identity (ArbitraryGraph samp) seed =
  sort (map (snd . snd) (fst (runRand (randomAssignGraph samp) (mkStdGen seed)))) ==
  sort samp

-- | Reachable nodes are uniquely assigned, and unreachable nodes are
-- uniquely assigned.
prop_unique_assignments :: ArbitraryGraph -> Int -> Bool
prop_unique_assignments (ArbitraryGraph samp) seed =
  unique
    (map
       (fst &&& (fmap (const ()) . fst . snd))
       (fst (runRand (randomAssignGraph samp) (mkStdGen seed))))
  where
    unique xs = nub xs == xs

-- | Topological order is preserved.
prop_toporder_preserved :: ArbitraryGraph -> Int -> Bool
prop_toporder_preserved (ArbitraryGraph samp) seed =
  all
    (\(i, (_, _, keys)) ->
       let dependencies =
             filter (\(_, (_, k, _)) -> elem k keys) reachableSorted
           beforeMe (j, _) = j < i
        in all beforeMe dependencies)
    reachableSorted
  where
    (assignments, _gen) = runRand (randomAssignGraph samp) (mkStdGen seed)
    sorted = sortBy (comparing fst) assignments
    reachableSorted =
      mapMaybe
        (\(i, (eitherVertex, node)) ->
           case eitherVertex of
             Left {} -> Just (i, node)
             Right {} -> Nothing)
        sorted

-- | For QuickCheck tests.
newtype ArbitraryGraph = ArbitraryGraph [(Int, Int, [Int])]
 deriving (Show, Eq)

instance Arbitrary ArbitraryGraph where
  arbitrary = do
    size <- getSize
    upper <- choose (0, size)
    population <- shuffle [0 .. upper]
    fmap
      ArbitraryGraph
      (foldM
         (\edges key -> do
            candidates <- sublistOf (filter (/= key) population)
            let edges' = (key, key, candidates) : edges
            -- If the edges were cyclic, just give up on this member
            -- of the population.
            if isCyclic edges'
              then pure edges
              else pure edges')
         []
         population)
    where
      isCyclic =
        any
          (\scc ->
             case scc of
               G.CyclicSCC {} -> True
               _ -> False) .
        G.stronglyConnComp

--------------------------------------------------------------------------------
-- Helpers

-- | Handy monadic mapAccumL.
mapAccumM ::
     (Traversable t, Monad m) => (a -> b -> m (a, c)) -> a -> t b -> m (a, t c)
mapAccumM f acc xs =
  fmap
    swap
    (runStateT (traverse (\x -> StateT (\s -> fmap swap (f s x)))
                         xs)
               acc)