Graph stable shuffle

StableShuffle.hs

{-# LANGUAGE PartialTypeSignatures #-}
module StableShuffle where

import           Control.Monad
import           Control.Monad.Random
import           Data.Bifunctor
import           Data.List
import qualified Data.Map.Strict as M
import           Data.Maybe
import           Data.Ord
import           Data.Tuple
import           Test.QuickCheck

{-

> quickCheckWith stdArgs {maxSuccess=10000} prop_dep_order
+++ OK, passed 10000 tests.

> map (second nodeNode) (evalRand (shuffleWithEdges nodes) (mkStdGen 0))
[(2,"Henry"),(4,"Eve"),(4,"Adam"),(4,"Wibble"),(5,"Cain"),(7,"Will"),(7,"Enoch"),(9,"Outlier"),(12,"Abel"),(13,"Seth")]

-}

prop_dep_order :: Int -> Bool
prop_dep_order seed =
  all
    (\node ->
       let nodeId = getNodeIdx (nodeKey node)
        in all
             (\key ->
                let precIdx = getNodeIdx key
                 in precIdx < nodeId)
             (nodePreceding node))
    nodes &&
  sort (nub nodes) == sort (nub (map snd result))
  where
    result = evalRand (shuffleWithEdges nodes) (mkStdGen seed)
    getNodeIdx key =
      case lookup key (map (first nodeKey . swap) result) of
        Nothing -> error "Missing node"
        Just nodeIdx -> nodeIdx

-- INVARIANT: must be in depedency order: parent, then children
--
-- INVARIANT: non-linked nodes should appear after linked nodes (or I
-- should split them up before-hand and change nodePreceding to
-- NonEmpty k)

data Node k a = Node {nodeKey :: k, nodeNode :: a, nodePreceding :: [k]}
  deriving (Eq, Ord, Show)
shuffleWithEdges :: (Ord k, MonadRandom m) => [Node k a] -> m [(Int, Node k a)]
shuffleWithEdges nodes = outer nodes
  where
    outer = fmap (sortBy (comparing fst) . snd) . foldM go (mempty, [])
    go (m, results) node = do
      let precedings = mapMaybe (flip M.lookup m) (nodePreceding node)
          minbound = foldl' max 0 precedings
          maxbound = M.foldl max (minbound + length nodes) m
      idx <- getRandomR (minbound + 1, maxbound + 1)
      pure (M.insert (nodeKey node) idx m, (idx, node) : results)

-- Read: Seth depends on Adam and Eve, etc.
nodes :: [Node String String]
nodes =
  [ n "Adam" []
  , n "Eve" []
  , n "Seth" ["Adam", "Eve"]
  , n "Cain" ["Adam", "Eve"]
  , n "Abel" ["Adam", "Eve"]
  , n "Enoch" ["Cain"]
  , n "Wibble" []
  , n "Henry" []
  , n "Will" ["Henry"]
  , n "Outlier" []
  ]
  where
    n x = Node x x