SplitSpec.hs

module Copar.Algorithm.SplitSpec (spec) where

import           Test.Hspec

import           Control.Monad.ST
import           Data.STRef
import           Data.List                      ( sort )

import           Control.Monad.Reader
import qualified Data.Vector                   as V
import qualified Data.Vector.Unboxed           as VU
import           Lens.Micro.Platform

import           Copar.Algorithm.Types
import           Copar.Algorithm.Split
import           Data.BlockQueue                ( BlockQueue )
import           Data.Partition.Common
import           Data.MorphismEncoding          ( fromEdgeRef
                                                , Encoding
                                                )
import qualified Data.MorphismEncoding         as Encoding
import           Copar.RefinementInterface
import           Copar.Algorithm.Initialize
import qualified Data.BlockQueue               as Queue
import           Copar.Functors.Powerset
import           Copar.Functors.GroupValued
import qualified Data.RefinablePartition       as Partition
import qualified Data.Partition                as Partition
                                                ( toBlocks )

spec :: Spec
spec = do
  collectTouchedBlocksSpec
  updateBlockSpec
  splitBlockSpec
  addBlocksToQueueSpec
  splitSpec


collectTouchedBlocksSpec :: Spec
collectTouchedBlocksSpec = describe "collectTouchedBlocks" $ do
  it "returns an empty list when no predecessors exist" $ do
    withState @Powerset (enc [True, False] [(0, (), 1)])
                        (collectTouchedBlocks (Block 0))
      `shouldBe` []

  it "returns the correct block when predecessors exist" $ do
    withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
                        (map fst <$> collectTouchedBlocks (Block 1))
      `shouldBe` [Block 0]

  it "doesn't return duplicates" $ do
    withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
                        (map fst <$> collectTouchedBlocks (Block 1))
      `shouldBe` [Block 0]

  it "computes the correct v0" $ do
    withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
                        (map snd <$> collectTouchedBlocks (Block 1))
      -- initially, C contains all blocks, so calling update with empty list of
      -- edges to S, the result is: No edges to "outside of C", all edges to C
      -- and no edges to S.
      `shouldBe` [mkPowerF3 False True False]

  it "marks the correct states" $ do
    let res =
          withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
            $ do
                [(b, _)] <- collectTouchedBlocks (Block 1)
                p <- view (_1 . partitionL)
                VU.toList <$> lift (Partition.markedStates p b)
    res `shouldMatchList` [0, 1]

  it "adds the correct edges to toSub" $ do
    let res =
          withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
            $ do
                _  <- collectTouchedBlocks (Block 1)
                ts <- view (_1 . toSubL) >>= V.freeze
                return (ts & each . each %~ fromEdgeRef & V.toList)
    res `shouldBe` [[0], [1], []]


updateBlockSpec :: Spec
updateBlockSpec = describe "updateBlock" $ do
  it "resets toSub to all empty lists"
    $ let
        res =
          withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 2)])
            $ do
                [(b, v0)] <- collectTouchedBlocks (Block 1)
                updateBlock b v0
                view (_1 . toSubL)
                  >>= V.freeze
                  <&> (each . each %~ fromEdgeRef)
                  <&> V.toList
      in  res `shouldBe` [[], [], []]

  -- Visualization of the graph. Can be rendered with ditaa.
  -- @
  --    e1
  --   /-\
  --   | |
  --   v |
  --  +---+  e0
  --  | 0 +-------\
  --  +-+-+       |
  --    |         v
  --    | e2    +---+
  --    |       | 2 |
  --    |       +---+
  --    v         ^
  --  +---+       |
  --  | 1 |-------/
  --  +---+  e3
  -- @
  it "updates lastW correctly"
    $ let res =
            withState @Powerset
                (enc [True, True, False]
                     [(0, (), 2), (0, (), 0), (0, (), 1), (1, (), 2)]
                )
              $ do
                  [(b, v0)] <- collectTouchedBlocks (Block 1)
                  updateBlock b v0
                  lw <- view (_1 . lastWL) >>= lift . V.freeze
                  lift (lw & V.toList & mapM readSTRef)
      in  res
            `shouldBe` [ packWeight True 1
                       , packWeight True 2
                       , packWeight True 2
                       , packWeight False 1
                       ]

  -- The idea here is that the edges from state 0 to block 1 cancel each other
  -- out and thus the state has a total weight of 0 and must be unmarked.
  it "unmarks states where F3 is v0"
    $ let res =
            withState @(GroupValued Int)
                (enc [1, 1, 0, 0]
                     [(0, 1, 2), (0, (-1), 3), (1, 1, 3), (0, 1, 1)]
                )
              $ do
                  [(b, v0)] <- collectTouchedBlocks (Block 1)
                  updateBlock b v0
                  p <- view (_1 . partitionL)
                  lift (Partition.markedStates p b) <&> VU.toList
      in  res `shouldBe` [1]

  it "caches F3 values for all non-v0 states"
    $ let res =
            withState @(GroupValued Int)
                (enc [1, 1, 0, 0]
                     [(0, 1, 2), (0, (-1), 3), (1, 1, 3), (0, 1, 1)]
                )
              $ do
                  [(b, v0)] <- collectTouchedBlocks (Block 1)
                  updateBlock b v0
                  f3 <- view (_1 . f3CacheL) >>= lift . V.freeze
                  return (f3 V.! 1)
      in  res `shouldBe` (mkGroupF3 0 1)


splitBlockSpec :: Spec
splitBlockSpec = describe "splitBlock" $ do
  it "splits blocks into marked and unmaked"
    $ let
        res =
          withState @Powerset (enc [True, True, False] [(0, (), 2), (1, (), 0)])
            $ do
                [(b, v0)] <- collectTouchedBlocks (Block 1)
                updateBlock b v0
                splitBlock b
      in  res `shouldMatchList` [Block 0, Block 2]

  it "splits different F3s into different blocks"
    $ let res =
            withState @(GroupValued Int)
                (enc [3, 3, 3, 0]
                     [(0, 1, 3), (1, 2, 3), (2, 3, 3), (0, 2, 0), (1, 1, 1)]
                )
              $ do
                  [(b, v0)] <- collectTouchedBlocks (Block 1)
                  updateBlock b v0
                  splitBlock b
      in  res `shouldMatchList` [Block 0, Block 2, Block 3]

  it "combines equal F3s into the same block"
    $ let res =
            withState @(GroupValued Int)
                (enc [3, 3, 3, 0]
                     [(0, 1, 3), (1, 1, 3), (2, 3, 3), (0, 2, 0), (1, 2, 1)]
                )
              $ do
                  [(b, v0)] <- collectTouchedBlocks (Block 1)
                  updateBlock b v0
                  splitBlock b
      in  res `shouldMatchList` [Block 0, Block 2]


addBlocksToQueueSpec :: Spec
addBlocksToQueueSpec = describe "addBlocksToQueue" $ do
  it "doesn't add the largest block to the queue"
    $ let res =
            withState @(GroupValued Int)
                (enc [1, 1, 2, 3] [(0, 1, 0), (1, 1, 1), (2, 2, 2), (3, 3, 3)])
              $ do
                  lift . Queue.clear =<< view _2
                  addBlocksToQueue (Block 0) (map Block [0, 1, 2])
                  q <- lift . Queue.toList =<< view _2
                  p <- view (_1 . partitionL)
                  lift $ forM q $ \b -> Partition.blockSize p b
      in  res `shouldMatchList` ([1, 1])

  it "does add all new blocks, if the original was already queued"
    $ let res =
            withState @(GroupValued Int)
                (enc [1, 1, 2, 3] [(0, 1, 0), (1, 1, 1), (2, 2, 2), (3, 3, 3)])
              $ do
                  queue <- view _2
                  lift (Queue.clear queue)
                  lift (Queue.enqueue queue (Block 0))
                  addBlocksToQueue (Block 0) (map Block [0, 1, 2])
                  lift (Queue.toList queue)
      in  res `shouldMatchList` (map Block [0, 1, 2])


splitSpec :: Spec
splitSpec = describe "split" $ do
  let encoding =
        (enc
          [10, 10, 10, 1, 1, 20, 20, 0]
          [ (0, 1 , 7)
          , (0, 9 , 0)
          , (1, 2 , 7)
          , (1, 8 , 1)
          , (2, 3 , 7)
          , (2, 7 , 2)
          , (3, 1 , 7)
          , (4, 1 , 4)
          , (5, 20, 6)
          , (6, 20, 5)
          ]
        )

  it "splits all touched blocks correctly"
    $ let res = withState @(GroupValued Int) encoding $ do
            p <- view (_1 . partitionL)
            b <- lift (Partition.blockOfState p 7)
            split b
            Partition.toBlocks <$> lift (Partition.freeze p)
      in  (map sort res)
            `shouldMatchList` [[0], [1], [2], [3], [4], [5, 6], [7]]

  it "adds the correct blocks to the queue"
    $ let
        res = withState @(GroupValued Int) encoding $ do
          p  <- view (_1 . partitionL)
          b  <- lift (Partition.blockOfState p 7)
          b3 <- lift (Partition.blockOfState p 3)
          (lift . Queue.delete b) =<< view _2
          (lift . Queue.delete b3) =<< view _2
          split b
          queuedBlocks <- lift . Queue.toList =<< view _2
          lift (mapM (fmap VU.toList . Partition.statesOfBlock p) queuedBlocks)
      in
        res
          `shouldSatisfy` (\l ->
                            ([0] `elem` l)
                              && ([1] `elem` l)
                              && ([2] `elem` l)
                              && ([5, 6] `elem` (map sort l))
                              && (  ([3] `elem` l && not ([4] `elem` l))
                                 || ([4] `elem` l && not ([3] `elem` l))
                                 )
                          )

withState
  :: RefinementInterface f
  => Encoding (Label f) (F1 f)
  -> (forall s . SplitM s f a)
  -> a
withState e action = runST $ do
  (q, as) <- initialize e True
  queue   <- Queue.empty 20
  mapM_ (Queue.enqueue queue) q
  runSplit as queue action

runSplit :: AlgoState s f -> BlockQueue s -> SplitM s f a -> ST s a
runSplit as queue action = runReaderT action (as, queue)

enc :: [f1] -> [(State, label, State)] -> Encoding label f1
enc f1s transitions = Encoding.new (V.fromList f1s) edges
  where
    edges = V.fromList (map mkEdge transitions)
    mkEdge (from, lab, to) = Encoding.Edge from lab to