Split.hs 5.41 KB
Newer Older
Hans-Peter Deifel's avatar
Hans-Peter Deifel committed
1
2
{-# LANGUAGE BangPatterns #-}

3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
module MA.Algorithm.Split
  ( SplitM
  , split
  ) where

import           Prelude hiding (pred)

import           Control.Monad
import           Control.Monad.ST
import           Control.Monad.ST.Unsafe (unsafeSTToIO)
import           Data.Function (on)
import           Data.List (delete,maximumBy)
import           Data.Maybe (maybeToList)
import           Data.STRef
import           System.IO.Unsafe (unsafeDupablePerformIO)

import           Control.Monad.Extra (unlessM, whenM, ifM)
import           Control.Monad.Reader
import           Data.Tuple.Extra (snd3)
import           Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
25
import qualified Data.Vector.Unboxed as VU
26
27
28
29
30
31
32

import           Data.Algorithm.PossibleMajorityCandidate
import           Data.BlockQueue (BlockQueue)
import qualified Data.BlockQueue as Queue
import           Data.MorphismEncoding
import           Data.RefinablePartition (State, Block)
import qualified Data.RefinablePartition as Partition
33
34
import           MA.RefinementInterface (RefinementInterface)
import qualified MA.RefinementInterface as RI
35
36
import           MA.Coalgebra.RefinementTypes
import           MA.Algorithm.Types
37
import           MA.Algorithm.Internal
38
39
40

type SplitM s h = ReaderT (AlgoState s h, BlockQueue s) (ST s)

41
42
split :: RefinementInterface h => Block -> SplitM s h ()
split blockS = do
43
  (as, _) <- ask
44
  touchedBlocks <- collectTouchedBlocks blockS
45
46
47
48
49

  forM_ touchedBlocks $ \(b, v0) -> do
    updateBlock b v0
    whenM (lift $ Partition.hasMarked (partition as) b) $
      splitBlock b
50
{-# SPECIALIZE split :: Block -> SplitM s TheFunctor () #-}
51
52
53
54

updateBlock :: forall s h. RefinementInterface h => Block -> H3 h -> SplitM s h ()
updateBlock b v0 = ask >>= \(as, _) -> lift $ do
  markB <- Partition.markedStates (partition as) b
55
  VU.forM_ markB $ \x -> do
56
57
    -- We can use `head` here, since states are only marked if they have at
    -- least one edge into S => toSub[x] can't be empty.
58
    !pc <- (fromEdgeRef . head <$> VM.read (toSub as) x)
59
            >>= VM.read (lastW as)
60
61
62
63
64
65
66
67
68
69
70
71
72

    !labelsToS <- {-# SCC readLabels #-} VM.read (toSub as) x >>= (mapM $ \e -> do
      let Edge _ !lab _ = graph (encoding as) e
      return $! lab)

    !pc' <- readSTRef pc
    (!wxS, !vx, !wxCwithoutS) <- {-# SCC riupdate #-} return $! (((RI.update @h) $! labelsToS) $! pc')
    writeSTRef pc $! {-# SCC wxCwithoutS #-} wxCwithoutS
    !ps <- newSTRef $! {-# SCC wxS #-} wxS

    VM.read (toSub as) x >>= \(!edges) -> forM_ edges $ \(EdgeRef !e) ->
      {-# SCC writelastw #-} (VM.write (lastW as) $! e) $! ps

73
74
75
76
    VM.write (toSub as) x []

    if vx == v0
      then Partition.unmark (partition as) x
77
78
      else VM.write (h3Cache as) x $! vx
{-# SPECIALIZE updateBlock :: Block -> H3 TheFunctor -> SplitM s TheFunctor () #-}
79
80
81
82
83
84
85

-- b must have at least one marked state
splitBlock :: RefinementInterface h => Block -> SplitM s h ()
splitBlock b = ask >>= \(as, queue) -> lift $ do
  -- b has marked states, so b1 is guaranteed to be non-empty
  (Just b1, bunmarked) <- Partition.splitMarked (partition as) b

Hans-Peter Deifel's avatar
Hans-Peter Deifel committed
86
87
88
89
90
91
  -- NOTE: We need to use unsafePerformIO here, because all vector sortBy
  -- functions expect a pure predicate. Since our predicate is only monadic
  -- because we need to _read_ from a mutable vector and doesn't have side
  -- effects, this should be safe.
  let unsafeH3 = unsafeDupablePerformIO . unsafeSTToIO . VM.read (h3Cache as)

92
93
  !pmc <- (possibleMajorityCandidateBy' unsafeH3) <$>
          Partition.unsafeStatesOfBlock (partition as) b1
94
95

  -- the pmc occurs in b1, so b1' has to be non-empty
96
97
  (Just b1', b2) <- Partition.splitBy (partition as) b1
                    ((==pmc) . unsafeH3)
98
99
100

  blocks <- ((b1':maybeToList bunmarked) ++) <$> case b2 of
    Nothing -> return []
Hans-Peter Deifel's avatar
Hans-Peter Deifel committed
101
    Just b2' -> Partition.groupBy (partition as) b2' unsafeH3
102
103
104
105
106
107

  let enqueue = Queue.enqueue queue

  ifM (b `Queue.elem` queue) (mapM_ enqueue blocks) $
    deleteLargest (Partition.blockSize (partition as)) (maybeAdd b blocks)
      >>= mapM_ enqueue
108
{-# SPECIALIZE splitBlock :: Block -> SplitM s TheFunctor () #-}
109
110
111
112
113
114
115
116

-- | Remove one largest element from the list
--
-- TODO This could probably be more efficient
deleteLargest :: Eq e => (e -> ST s Int) -> [e] -> ST s [e]
deleteLargest sizeFunction lst = do
  zipWithSize <- traverse (\x -> (,x) <$> sizeFunction x) lst
  return (delete (snd (maximumBy (compare `on` fst) zipWithSize)) lst)
117
{-# INLINE deleteLargest #-}
118
119
120
121
122
123

-- | Add element to list if it isn't already there
maybeAdd :: Eq e => e -> [e] -> [e]
maybeAdd e lst
  | e `elem` lst = lst
  | otherwise    = e : lst
124
{-# INLINE maybeAdd #-}
125

126
127
collectTouchedBlocks :: forall s h. RefinementInterface h => Block -> SplitM s h [(Block, H3 h)]
collectTouchedBlocks blockS = do
128
129
  (as, _) <- ask

130
131
  statesOfS <- lift $ Partition.statesOfBlock (partition as) blockS

132
133
  markedBlocks <- lift $ newSTRef []

134
  lift $ VU.forM_ statesOfS $ \y -> forM_ (pred as V.! y) $ \e -> do
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    let Edge x _ _ = graph (encoding as) e
    b <- Partition.blockOfState (partition as) x

    unlessM (Partition.hasMarked (partition as) b) $ do
      wCx <- readSTRef =<< VM.read (lastW as) (fromEdgeRef e)
      let v0 = snd3 $ RI.update @h [] wCx
      modifySTRef markedBlocks ((b, v0):)

    whenM (null <$> VM.read (toSub as) x) $
      Partition.mark (partition as) x

    VM.modify (toSub as) (e:) x

  lift $ readSTRef markedBlocks
149
{-# SPECIALIZE collectTouchedBlocks :: Block -> SplitM s TheFunctor [(Block, H3 TheFunctor)] #-}