-- |
-- Module      :  Data.Algorithm.Assignment
-- Copyright   :  © 2024–present Mark Karpov
-- License     :  BSD 3 clause
--
-- Maintainer  :  Mark Karpov <markkarpov92@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- A solution to the assignment problem.
module Data.Algorithm.Assignment
  ( assign,
  )
where

import Control.Monad (forM_, void, when)
import Control.Monad.Fix (fix)
import Control.Monad.ST (ST, runST)
import Data.Array (Array)
import Data.Array.Base qualified as A
import Data.Array.ST (STUArray)
import Data.Array.ST qualified as ST
import Data.STRef (modifySTRef', newSTRef, readSTRef, writeSTRef)

type CostMatrix s = STUArray s (Int, Int) Int

type MarkMatrix s = STUArray s (Int, Int) Char

type CoverageVector s = STUArray s Int Bool

-- | \(\mathcal{O}(n^4)\). Assign elements from two collections to each
-- other so that the total cost is minimal. The cost of each combination is
-- given the by the first argument and it can be negative. If any of the
-- collections is empty the result is the empty list. The sizes of the
-- collections need not to match. Finally, there is no guarantees on the
-- order of elements in the returned list of pairs.
--
-- See: <https://en.wikipedia.org/wiki/Hungarian_algorithm#Matrix_interpretation>
assign ::
  -- | How to calculate the cost
  (a -> b -> Int) ->
  -- | The first collection
  [a] ->
  -- | The second collection
  [b] ->
  -- | The resulting optimal assignment (no guarantees about order)
  [(a, b)]
assign :: forall a b. (a -> b -> Int) -> [a] -> [b] -> [(a, b)]
assign a -> b -> Int
_ [] [b]
_ = []
assign a -> b -> Int
_ [a]
_ [] = []
assign a -> b -> Int
cost [a]
as [b]
bs = (forall s. ST s [(a, b)]) -> [(a, b)]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [(a, b)]) -> [(a, b)])
-> (forall s. ST s [(a, b)]) -> [(a, b)]
forall a b. (a -> b) -> a -> b
$ do
  let length_a :: Int
length_a = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as
      length_b :: Int
length_b = [b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
bs
      aMinBound :: Int
aMinBound = Int
0
      aMaxBound :: Int
aMaxBound = Int
length_a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      bMinBound :: Int
bMinBound = Int
0
      bMaxBound :: Int
bMaxBound = Int
length_b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      abMaxBound :: Int
abMaxBound = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
aMaxBound Int
bMaxBound
      asArray :: Array Int a
asArray = (Int, Int) -> [a] -> Array Int a
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Int
aMinBound, Int
aMaxBound) [a]
as
      bsArray :: Array Int b
bsArray = (Int, Int) -> [b] -> Array Int b
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Int
bMinBound, Int
bMaxBound) [b]
bs
      matrixBounds :: ((Int, Int), (Int, Int))
matrixBounds = ((Int
aMinBound, Int
bMinBound), (Int
abMaxBound, Int
abMaxBound))
  STUArray s (Int, Int) Int
c <- ((Int, Int), (Int, Int)) -> Int -> ST s (STUArray s (Int, Int) Int)
forall i. Ix i => (i, i) -> Int -> ST s (STUArray s i Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
ST.newArray ((Int, Int), (Int, Int))
matrixBounds Int
0
  STUArray s (Int, Int) Char
m <- ((Int, Int), (Int, Int))
-> Char -> ST s (STUArray s (Int, Int) Char)
forall i. Ix i => (i, i) -> Char -> ST s (STUArray s i Char)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
ST.newArray ((Int, Int), (Int, Int))
matrixBounds Char
noMark
  STUArray s Int Bool
aCoverage <- (Int, Int) -> Bool -> ST s (STUArray s Int Bool)
forall i. Ix i => (i, i) -> Bool -> ST s (STUArray s i Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
ST.newArray (Int
aMinBound, Int
abMaxBound) Bool
False
  STUArray s Int Bool
bCoverage <- (Int, Int) -> Bool -> ST s (STUArray s Int Bool)
forall i. Ix i => (i, i) -> Bool -> ST s (STUArray s i Bool)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
ST.newArray (Int
bMinBound, Int
abMaxBound) Bool
False
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j ->
      STUArray s (Int, Int) Int -> (Int, Int) -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray STUArray s (Int, Int) Int
c (Int
i, Int
j) (a -> b -> Int
cost (Array Int a
asArray Array Int a -> Int -> a
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Int
i) (Array Int b
bsArray Array Int b -> Int -> b
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Int
j))
  if Int
aMaxBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
aMinBound Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bMaxBound Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bMinBound
    then STUArray s (Int, Int) Int -> ST s ()
forall s. CostMatrix s -> ST s ()
normalizePerB STUArray s (Int, Int) Int
c
    else STUArray s (Int, Int) Int -> ST s ()
forall s. CostMatrix s -> ST s ()
normalizePerA STUArray s (Int, Int) Int
c
  STUArray s (Int, Int) Int
-> STUArray s (Int, Int) Char
-> STUArray s Int Bool
-> STUArray s Int Bool
-> ST s ()
forall s.
CostMatrix s
-> MarkMatrix s -> CoverageVector s -> CoverageVector s -> ST s ()
starZeros STUArray s (Int, Int) Int
c STUArray s (Int, Int) Char
m STUArray s Int Bool
aCoverage STUArray s Int Bool
bCoverage
  (ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)]
forall a. (a -> a) -> a
fix ((ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)])
-> (ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)]
forall a b. (a -> b) -> a -> b
$ \ST s [(a, b)]
recurse0 -> do
    Bool
done <- STUArray s (Int, Int) Char
-> STUArray s Int Bool -> STUArray s Int Bool -> ST s Bool
forall s.
MarkMatrix s -> CoverageVector s -> CoverageVector s -> ST s Bool
coverZeros STUArray s (Int, Int) Char
m STUArray s Int Bool
aCoverage STUArray s Int Bool
bCoverage
    if Bool
done
      then STUArray s (Int, Int) Char
-> Array Int a -> Array Int b -> ST s [(a, b)]
forall s a b.
MarkMatrix s -> Array Int a -> Array Int b -> ST s [(a, b)]
recoverResults STUArray s (Int, Int) Char
m Array Int a
asArray Array Int b
bsArray
      else (ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)]
forall a. (a -> a) -> a
fix ((ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)])
-> (ST s [(a, b)] -> ST s [(a, b)]) -> ST s [(a, b)]
forall a b. (a -> b) -> a -> b
$ \ST s [(a, b)]
recurse1 -> do
        Maybe (Int, Int)
r <- STUArray s (Int, Int) Int
-> STUArray s (Int, Int) Char
-> STUArray s Int Bool
-> STUArray s Int Bool
-> ST s (Maybe (Int, Int))
forall s.
CostMatrix s
-> MarkMatrix s
-> CoverageVector s
-> CoverageVector s
-> ST s (Maybe (Int, Int))
primeUncoveredZero STUArray s (Int, Int) Int
c STUArray s (Int, Int) Char
m STUArray s Int Bool
aCoverage STUArray s Int Bool
bCoverage
        case Maybe (Int, Int)
r of
          Maybe (Int, Int)
Nothing -> do
            STUArray s (Int, Int) Int
-> STUArray s Int Bool -> STUArray s Int Bool -> ST s ()
forall s.
CostMatrix s -> CoverageVector s -> CoverageVector s -> ST s ()
adjustCosts STUArray s (Int, Int) Int
c STUArray s Int Bool
aCoverage STUArray s Int Bool
bCoverage
            ST s [(a, b)]
recurse1
          Just (Int, Int)
z0 -> do
            STUArray s (Int, Int) Char -> (Int, Int) -> ST s ()
forall s. MarkMatrix s -> (Int, Int) -> ST s ()
adjustMarks STUArray s (Int, Int) Char
m (Int, Int)
z0
            STUArray s Int Bool -> STUArray s Int Bool -> ST s ()
forall s. CoverageVector s -> CoverageVector s -> ST s ()
clearCoverage STUArray s Int Bool
aCoverage STUArray s Int Bool
bCoverage
            ST s [(a, b)]
recurse0
{-# INLINEABLE assign #-}

normalizePerA :: CostMatrix s -> ST s ()
normalizePerA :: forall s. CostMatrix s -> ST s ()
normalizePerA CostMatrix s
c = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- CostMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Int -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds CostMatrix s
c
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    STRef s Int
minValueRef <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
forall a. Bounded a => a
maxBound
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j ->
      CostMatrix s -> (Int, Int) -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CostMatrix s
c (Int
i, Int
j) ST s Int -> (Int -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
minValueRef ((Int -> Int) -> ST s ()) -> (Int -> Int -> Int) -> Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
    Int
minValue <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
minValueRef
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
minValue Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
        CostMatrix s -> (Int, Int) -> (Int -> Int) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray' CostMatrix s
c (Int
i, Int
j) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
minValue)
{-# INLINE normalizePerA #-}

normalizePerB :: CostMatrix s -> ST s ()
normalizePerB :: forall s. CostMatrix s -> ST s ()
normalizePerB CostMatrix s
c = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- CostMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Int -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds CostMatrix s
c
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
    STRef s Int
minValueRef <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
forall a. Bounded a => a
maxBound
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      CostMatrix s -> (Int, Int) -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CostMatrix s
c (Int
i, Int
j) ST s Int -> (Int -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
minValueRef ((Int -> Int) -> ST s ()) -> (Int -> Int -> Int) -> Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
    Int
minValue <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
minValueRef
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
minValue Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
        CostMatrix s -> (Int, Int) -> (Int -> Int) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray' CostMatrix s
c (Int
i, Int
j) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
minValue)
{-# INLINE normalizePerB #-}

starZeros ::
  CostMatrix s ->
  MarkMatrix s ->
  CoverageVector s ->
  CoverageVector s ->
  ST s ()
starZeros :: forall s.
CostMatrix s
-> MarkMatrix s -> CoverageVector s -> CoverageVector s -> ST s ()
starZeros CostMatrix s
c MarkMatrix s
m CoverageVector s
aCoverage CoverageVector s
bCoverage = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- CostMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Int -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds CostMatrix s
c
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Int
x <- CostMatrix s -> (Int, Int) -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CostMatrix s
c (Int
i, Int
j)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        Bool
aCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
aCoverage Int
i
        Bool
bCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
bCoverage Int
j
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
aCovered Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
bCovered) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          MarkMatrix s -> (Int, Int) -> Char -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray MarkMatrix s
m (Int
i, Int
j) Char
starMark
          CoverageVector s -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray CoverageVector s
aCoverage Int
i Bool
True
          CoverageVector s -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray CoverageVector s
bCoverage Int
j Bool
True
  CoverageVector s -> CoverageVector s -> ST s ()
forall s. CoverageVector s -> CoverageVector s -> ST s ()
clearCoverage CoverageVector s
aCoverage CoverageVector s
bCoverage
{-# INLINE starZeros #-}

coverZeros ::
  MarkMatrix s ->
  CoverageVector s ->
  CoverageVector s ->
  ST s Bool
coverZeros :: forall s.
MarkMatrix s -> CoverageVector s -> CoverageVector s -> ST s Bool
coverZeros MarkMatrix s
m CoverageVector s
_aCoverage CoverageVector s
bCoverage = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  STRef s Int
nRef <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
0
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Char
x <- MarkMatrix s -> (Int, Int) -> ST s Char
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray MarkMatrix s
m (Int
i, Int
j)
      Bool
bCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
bCoverage Int
j
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
starMark Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
bCovered) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        CoverageVector s -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray CoverageVector s
bCoverage Int
j Bool
True
        STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
nRef (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int
n <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
nRef
  Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
aMaxBound)
{-# INLINE coverZeros #-}

recoverResults ::
  MarkMatrix s ->
  Array Int a ->
  Array Int b ->
  ST s [(a, b)]
recoverResults :: forall s a b.
MarkMatrix s -> Array Int a -> Array Int b -> ST s [(a, b)]
recoverResults MarkMatrix s
m Array Int a
as Array Int b
bs = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  STRef s [(a, b)]
resultRef <- [(a, b)] -> ST s (STRef s [(a, b)])
forall a s. a -> ST s (STRef s a)
newSTRef []
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Char
x <- MarkMatrix s -> (Int, Int) -> ST s Char
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray MarkMatrix s
m (Int
i, Int
j)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
starMark) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        case (,) (a -> b -> (a, b)) -> Maybe a -> Maybe (b -> (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Array Int a
as Array Int a -> Int -> Maybe a
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> Maybe e
A.!? Int
i) Maybe (b -> (a, b)) -> Maybe b -> Maybe (a, b)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Array Int b
bs Array Int b -> Int -> Maybe b
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> Maybe e
A.!? Int
j) of
          Maybe (a, b)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Just (a
a, b
b) -> STRef s [(a, b)] -> ([(a, b)] -> [(a, b)]) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s [(a, b)]
resultRef ((a
a, b
b) :)
  STRef s [(a, b)] -> ST s [(a, b)]
forall s a. STRef s a -> ST s a
readSTRef STRef s [(a, b)]
resultRef
{-# INLINE recoverResults #-}

primeUncoveredZero ::
  CostMatrix s ->
  MarkMatrix s ->
  CoverageVector s ->
  CoverageVector s ->
  ST s (Maybe (Int, Int))
primeUncoveredZero :: forall s.
CostMatrix s
-> MarkMatrix s
-> CoverageVector s
-> CoverageVector s
-> ST s (Maybe (Int, Int))
primeUncoveredZero CostMatrix s
c MarkMatrix s
m CoverageVector s
aCoverage CoverageVector s
bCoverage = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  STRef s (Maybe (Int, Int))
primedRef <- Maybe (Int, Int) -> ST s (STRef s (Maybe (Int, Int)))
forall a s. a -> ST s (STRef s a)
newSTRef Maybe (Int, Int)
forall a. Maybe a
Nothing
  ST s Bool -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s Bool -> ST s ())
-> ((Int -> ST s Bool) -> ST s Bool)
-> (Int -> ST s Bool)
-> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
forall s. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' Int
aMinBound Int
aMaxBound ((Int -> ST s Bool) -> ST s ()) -> (Int -> ST s Bool) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s Bool) -> ST s Bool
forall s. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' Int
bMinBound Int
bMaxBound ((Int -> ST s Bool) -> ST s Bool)
-> (Int -> ST s Bool) -> ST s Bool
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Int
x <- CostMatrix s -> (Int, Int) -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CostMatrix s
c (Int
i, Int
j)
      if Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then do
          Bool
aCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
aCoverage Int
i
          Bool
bCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
bCoverage Int
j
          if Bool -> Bool
not Bool
aCovered Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
bCovered
            then Bool
False Bool -> ST s () -> ST s Bool
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ STRef s (Maybe (Int, Int)) -> Maybe (Int, Int) -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Maybe (Int, Int))
primedRef ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
i, Int
j))
            else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  Maybe (Int, Int)
r <- STRef s (Maybe (Int, Int)) -> ST s (Maybe (Int, Int))
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe (Int, Int))
primedRef
  case Maybe (Int, Int)
r of
    Maybe (Int, Int)
Nothing -> Maybe (Int, Int) -> ST s (Maybe (Int, Int))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Int, Int)
forall a. Maybe a
Nothing
    Just (Int
i, Int
j) -> do
      MarkMatrix s -> (Int, Int) -> Char -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray MarkMatrix s
m (Int
i, Int
j) Char
primeMark
      Maybe Int
mj' <- MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
forall s. MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
findInA MarkMatrix s
m Char
starMark Int
i
      case Maybe Int
mj' of
        Maybe Int
Nothing -> Maybe (Int, Int) -> ST s (Maybe (Int, Int))
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
i, Int
j))
        Just Int
j' -> do
          CoverageVector s -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray CoverageVector s
aCoverage Int
i Bool
True
          CoverageVector s -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray CoverageVector s
bCoverage Int
j' Bool
False
          CostMatrix s
-> MarkMatrix s
-> CoverageVector s
-> CoverageVector s
-> ST s (Maybe (Int, Int))
forall s.
CostMatrix s
-> MarkMatrix s
-> CoverageVector s
-> CoverageVector s
-> ST s (Maybe (Int, Int))
primeUncoveredZero CostMatrix s
c MarkMatrix s
m CoverageVector s
aCoverage CoverageVector s
bCoverage
{-# INLINEABLE primeUncoveredZero #-}

adjustMarks :: MarkMatrix s -> (Int, Int) -> ST s ()
adjustMarks :: forall s. MarkMatrix s -> (Int, Int) -> ST s ()
adjustMarks MarkMatrix s
m (Int, Int)
z0 = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  let go :: (Int, Int) -> [(Int, Int)] -> ST s [(Int, Int)]
go (Int
_, Int
j) [(Int, Int)]
acc = do
        Maybe Int
mi' <- MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
forall s. MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
findInB MarkMatrix s
m Char
starMark Int
j
        case Maybe Int
mi' of
          Maybe Int
Nothing -> [(Int, Int)] -> ST s [(Int, Int)]
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return [(Int, Int)]
acc
          Just Int
i' -> do
            Maybe Int
mj' <- MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
forall s. MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
findInA MarkMatrix s
m Char
primeMark Int
i'
            case Maybe Int
mj' of
              Maybe Int
Nothing -> [Char] -> ST s [(Int, Int)]
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Algorithm.Assignment.adjustMarks"
              Just Int
j' -> (Int, Int) -> [(Int, Int)] -> ST s [(Int, Int)]
go (Int
i', Int
j') ((Int
i', Int
j) (Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
: (Int
i', Int
j') (Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
: [(Int, Int)]
acc)
  [(Int, Int)]
path <- (Int, Int) -> [(Int, Int)] -> ST s [(Int, Int)]
go (Int, Int)
z0 [(Int, Int)
z0]
  [(Int, Int)] -> ((Int, Int) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, Int)]
path (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, Int
j) -> do
    let adjust :: Char -> Char
adjust Char
x =
          if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
starMark
            then Char
noMark
            else Char
starMark
    MarkMatrix s -> (Int, Int) -> (Char -> Char) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray' MarkMatrix s
m (Int
i, Int
j) Char -> Char
adjust
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      let resetPrime :: Char -> Char
resetPrime Char
x =
            if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
primeMark
              then Char
noMark
              else Char
x
      MarkMatrix s -> (Int, Int) -> (Char -> Char) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray' MarkMatrix s
m (Int
i, Int
j) Char -> Char
resetPrime
{-# INLINE adjustMarks #-}

adjustCosts ::
  CostMatrix s ->
  CoverageVector s ->
  CoverageVector s ->
  ST s ()
adjustCosts :: forall s.
CostMatrix s -> CoverageVector s -> CoverageVector s -> ST s ()
adjustCosts CostMatrix s
c CoverageVector s
aCoverage CoverageVector s
bCoverage = do
  ((Int
aMinBound, Int
bMinBound), (Int
aMaxBound, Int
bMaxBound)) <- CostMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Int -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds CostMatrix s
c
  STRef s Int
minUncoveredValueRef <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
forall a. Bounded a => a
maxBound
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Bool
aCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
aCoverage Int
i
      Bool
bCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
bCoverage Int
j
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
aCovered Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
bCovered) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        CostMatrix s -> (Int, Int) -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CostMatrix s
c (Int
i, Int
j) ST s Int -> (Int -> ST s ()) -> ST s ()
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
minUncoveredValueRef ((Int -> Int) -> ST s ()) -> (Int -> Int -> Int) -> Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
  Int
minUncoveredValue <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
minUncoveredValueRef
  Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
aMinBound Int
aMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
    Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
bMinBound Int
bMaxBound ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Bool
aCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
aCoverage Int
i
      Bool
bCovered <- CoverageVector s -> Int -> ST s Bool
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray CoverageVector s
bCoverage Int
j
      if Bool -> Bool
not Bool
aCovered Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
bCovered
        then CostMatrix s -> (Int, Int) -> (Int -> Int) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray CostMatrix s
c (Int
i, Int
j) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
minUncoveredValue)
        else
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
aCovered Bool -> Bool -> Bool
&& Bool
bCovered) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
            CostMatrix s -> (Int, Int) -> (Int -> Int) -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
ST.modifyArray CostMatrix s
c (Int
i, Int
j) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
minUncoveredValue)
{-# INLINE adjustCosts #-}

clearCoverage ::
  CoverageVector s ->
  CoverageVector s ->
  ST s ()
clearCoverage :: forall s. CoverageVector s -> CoverageVector s -> ST s ()
clearCoverage CoverageVector s
aCoverage CoverageVector s
bCoverage = do
  let clearOne :: a Int Bool -> ST s ()
clearOne a Int Bool
v = do
        (Int
from, Int
to) <- a Int Bool -> ST s (Int, Int)
forall i. Ix i => a i Bool -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds a Int Bool
v
        Int -> Int -> (Int -> ST s ()) -> ST s ()
forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
from Int
to ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
          a Int Bool -> Int -> Bool -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
ST.writeArray a Int Bool
v Int
i Bool
False
  CoverageVector s -> ST s ()
forall {a :: * -> * -> *} {s}.
MArray a Bool (ST s) =>
a Int Bool -> ST s ()
clearOne CoverageVector s
aCoverage
  CoverageVector s -> ST s ()
forall {a :: * -> * -> *} {s}.
MArray a Bool (ST s) =>
a Int Bool -> ST s ()
clearOne CoverageVector s
bCoverage
{-# INLINE clearCoverage #-}

findInA ::
  MarkMatrix s ->
  Char ->
  Int ->
  ST s (Maybe Int)
findInA :: forall s. MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
findInA MarkMatrix s
m Char
mark Int
i = do
  ((Int
_aMinBound, Int
bMinBound), (Int
_aMaxBound, Int
bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  STRef s (Maybe Int)
starredRef <- Maybe Int -> ST s (STRef s (Maybe Int))
forall a s. a -> ST s (STRef s a)
newSTRef Maybe Int
forall a. Maybe a
Nothing
  ST s Bool -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s Bool -> ST s ())
-> ((Int -> ST s Bool) -> ST s Bool)
-> (Int -> ST s Bool)
-> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
forall s. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' Int
bMinBound Int
bMaxBound ((Int -> ST s Bool) -> ST s ()) -> (Int -> ST s Bool) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
    Char
x <- MarkMatrix s -> (Int, Int) -> ST s Char
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray MarkMatrix s
m (Int
i, Int
j)
    if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
mark
      then Bool
False Bool -> ST s () -> ST s Bool
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ STRef s (Maybe Int) -> Maybe Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Maybe Int)
starredRef (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
j)
      else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  STRef s (Maybe Int) -> ST s (Maybe Int)
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe Int)
starredRef
{-# INLINE findInA #-}

findInB ::
  MarkMatrix s ->
  Char ->
  Int ->
  ST s (Maybe Int)
findInB :: forall s. MarkMatrix s -> Char -> Int -> ST s (Maybe Int)
findInB MarkMatrix s
m Char
mark Int
j = do
  ((Int
aMinBound, Int
_bMinBound), (Int
aMaxBound, Int
_bMaxBound)) <- MarkMatrix s -> ST s ((Int, Int), (Int, Int))
forall i. Ix i => STUArray s i Char -> ST s (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
ST.getBounds MarkMatrix s
m
  STRef s (Maybe Int)
starredRef <- Maybe Int -> ST s (STRef s (Maybe Int))
forall a s. a -> ST s (STRef s a)
newSTRef Maybe Int
forall a. Maybe a
Nothing
  ST s Bool -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s Bool -> ST s ())
-> ((Int -> ST s Bool) -> ST s Bool)
-> (Int -> ST s Bool)
-> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
forall s. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' Int
aMinBound Int
aMaxBound ((Int -> ST s Bool) -> ST s ()) -> (Int -> ST s Bool) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    Char
x <- MarkMatrix s -> (Int, Int) -> ST s Char
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
ST.readArray MarkMatrix s
m (Int
i, Int
j)
    if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
mark
      then Bool
False Bool -> ST s () -> ST s Bool
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ STRef s (Maybe Int) -> Maybe Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Maybe Int)
starredRef (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i)
      else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  STRef s (Maybe Int) -> ST s (Maybe Int)
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe Int)
starredRef
{-# INLINE findInB #-}

countFromTo :: Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo :: forall s. Int -> Int -> (Int -> ST s ()) -> ST s ()
countFromTo Int
start Int
end Int -> ST s ()
action = Int -> ST s ()
go Int
start
  where
    go :: Int -> ST s ()
go !Int
n = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
end) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int -> ST s ()
action Int
n
      Int -> ST s ()
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE countFromTo #-}

countFromTo' :: Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' :: forall s. Int -> Int -> (Int -> ST s Bool) -> ST s Bool
countFromTo' Int
start Int
end Int -> ST s Bool
action = Int -> ST s Bool
go Int
start
  where
    go :: Int -> ST s Bool
go !Int
n =
      if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
end
        then do
          Bool
r <- Int -> ST s Bool
action Int
n
          if Bool
r then Int -> ST s Bool
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        else Bool -> ST s Bool
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
{-# INLINE countFromTo' #-}

noMark, starMark, primeMark :: Char
noMark :: Char
noMark = Char
'n'
starMark :: Char
starMark = Char
's'
primeMark :: Char
primeMark = Char
'p'