Convert pred array to new Stack data structure

src/MA/Algorithm/Initialize.hs

@@ -18,6 +18,7 @@ import qualified MA.RefinementInterface as RI
 import           MA.Algorithm.Types
 import           MA.Coalgebra.RefinementTypes
 import qualified Data.RefinablePartition as Partition
+import qualified Data.Stack as Stack
 
 -- returns (initial queue content, algo state)
 initialize :: forall h s. RefinementInterface h
@@ -26,12 +27,13 @@ initialize :: forall h s. RefinementInterface h
 initialize encoding = do
   toSub <- VM.replicate (size encoding) []
   lastW <- VM.new (length (edges encoding))
-  predMutable <- VM.replicate (size encoding) []
+  predMutable <- V.generateM (size encoding) $ \i ->
+    Stack.make (numPreds encoding i)
   h3Cache <- VM.new (size encoding)
 
   iforM_ (edges encoding) $ \i (Edge x _ y) -> do
     VM.modify toSub (EdgeRef i:) x
-    VM.modify predMutable (EdgeRef i:) y
+    Stack.unsafePush (predMutable V.! y) (EdgeRef i)
 
   forM_ (states encoding) $ \x -> do
     outgoingLabels <- map (label . graph encoding) <$> VM.read toSub x
@@ -42,7 +44,7 @@ initialize encoding = do
 
 
   -- We can use unsafeFreeze here, since we don't need predMutable anymore
-  pred <- V.unsafeFreeze predMutable
+  pred <- V.mapM (Stack.unsafeFreeze) predMutable
 
   -- Initialize partition with one block and assigning each state to that block
   partition <- Partition.make1 (size encoding)

src/MA/Algorithm/Split.hs

@@ -157,7 +157,7 @@ collectTouchedBlocks blockS = do
 
   markedBlocks <- lift $ newSTRef []
 
-  lift $ VU.forM_ statesOfS $ \y -> forM_ (pred as V.! y) $ \e -> do
+  lift $ VU.forM_ statesOfS $ \y -> VU.forM_ (pred as V.! y) $ \e -> do
     let Edge x _ _ = graph (encoding as) e
     b <- Partition.blockOfState (partition as) x
 

src/MA/Algorithm/Types.hs

@@ -17,6 +17,7 @@ import           Data.STRef
 
 import           Data.Vector                    ( Vector )
 import           Data.Vector.Mutable            ( MVector )
+import qualified Data.Vector.Unboxed           as VU
 import           Lens.Micro.TH
 
 import           Data.MorphismEncoding
@@ -27,7 +28,7 @@ data AlgoState s h = AlgoState
   { toSub :: MVector s [EdgeRef]
   , lastW :: MVector s (STRef s (Weight h))
   , encoding :: Encoding (Label h) (H1 h)
-  , pred :: Vector [EdgeRef]
+  , pred :: Vector (VU.Vector EdgeRef)
   , partition :: RefinablePartition s
   , h3Cache :: MVector s (H3 h)
   }