-- | Internal implementation of @AtCoder.String@ module.
module AtCoder.Internal.String
  ( -- * Suffix array
    saNaive,
    saDoubling,
    saIsImpl,
    saIs,
    saIsManual,
  )
where

import Control.Monad (unless, when)
import Control.Monad.ST (runST)
import Data.Bit (Bit (..))
import Data.Foldable (for_)
import Data.Vector.Algorithms.Intro qualified as VAI
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)

-- TODO: remove `HasCallStack`?

-- | \(O(n^2)\) Internal implementation of suffix array creation (naive).
--
-- @since 1.0.0.0
{-# INLINE saNaive #-}
saNaive :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
saNaive :: HasCallStack => Vector Int -> Vector Int
saNaive Vector Int
s =
  let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
s
      f :: Int -> Int -> Ordering
f Int
l0 Int
r0
        | Int
l0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r0 = Ordering
GT
        | Bool
otherwise = Int -> Int -> Ordering
inner Int
l0 Int
r0
        where
          inner :: Int -> Int -> Ordering
inner Int
l Int
r
            | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
                let sl :: Int
sl = Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
l
                    sr :: Int
sr = Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
r
                 in if Int
sl Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
sr
                      then Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
sl Int
sr
                      else Int -> Int -> Ordering
inner (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Ordering
LT
            | Bool
otherwise = Ordering
GT
   in -- modify + generate should fuse
      (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify ((Int -> Int -> Ordering)
-> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VAI.sortBy Int -> Int -> Ordering
f) (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n Int -> Int
forall a. a -> a
id

-- | \(O(n \log n)\) Internal implementation of suffix array creation (doubling).
--
-- @since 1.0.0.0
{-# INLINE saDoubling #-}
saDoubling :: (HasCallStack) => VU.Vector Int -> VU.Vector Int
saDoubling :: HasCallStack => Vector Int -> Vector Int
saDoubling Vector Int
s = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
s
  MVector s Int
sa <- Int -> (Int -> Int) -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
VUM.generate Int
n Int -> Int
forall a. a -> a
id
  let loop :: MVector s Int -> MVector s Int -> Int -> ST s ()
loop MVector s Int
rnk MVector s Int
tmp Int
k
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            Vector Int
rnk' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
rnk
            let cmp :: Int -> Int -> Ordering
cmp Int
x Int
y =
                  let rnkX :: Int
rnkX = Vector Int
rnk' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
x
                      rnkY :: Int
rnkY = Vector Int
rnk' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y
                   in if Int
rnkX Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
rnkY
                        then Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
rnkX Int
rnkY
                        else
                          let rx :: Int
rx = if Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n then Vector Int
rnk' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) else (-Int
1)
                              ry :: Int
ry = if Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n then Vector Int
rnk' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) else (-Int
1)
                           in Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
rx Int
ry
            (Int -> Int -> Ordering)
-> MVector (PrimState (ST s)) Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
VAI.sortBy Int -> Int -> Ordering
cmp MVector s Int
MVector (PrimState (ST s)) Int
sa
            Vector Int
sa' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sa
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
tmp (Vector Int
sa' Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0) Int
0
            (Int -> Int -> ST s ()) -> Vector Int -> Vector Int -> ST s ()
forall (m :: * -> *) a b c.
(Monad m, Unbox a, Unbox b) =>
(a -> b -> m c) -> Vector a -> Vector b -> m ()
VU.zipWithM_
              ( \Int
saI_ Int
saI -> do
                  Int
tmpI_ <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
tmp Int
saI_
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
tmp Int
saI (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int
tmpI_ Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if Int -> Int -> Ordering
cmp Int
saI_ Int
saI Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
LT then Int
1 else Int
0
              )
              Vector Int
sa'
              (Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.tail Vector Int
sa')
            MVector s Int -> MVector s Int -> Int -> ST s ()
loop MVector s Int
tmp MVector s Int
rnk (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k)
  MVector s Int
rnkVec <- Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector Int
s
  MVector s Int
tmpVec <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
n
  MVector s Int -> MVector s Int -> Int -> ST s ()
loop MVector s Int
rnkVec MVector s Int
tmpVec Int
1
  MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
sa

-- TODO: use backpermute

-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- @since 1.0.0.0
{-# INLINE saIsImpl #-}
saIsImpl ::
  (HasCallStack) =>
  -- | naive threshould
  Int ->
  -- | doubling threshould
  Int ->
  -- | string
  VU.Vector Int ->
  -- | upper bounds
  Int ->
  -- | suffix array
  VU.Vector Int
saIsImpl :: HasCallStack => Int -> Int -> Vector Int -> Int -> Vector Int
saIsImpl Int
naiveThreshold Int
doublingThreshold Vector Int
s Int
upper = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
s
  let !ls :: Vector Bit
ls = (forall s. ST s (MVector s Bit)) -> Vector Bit
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Bit)) -> Vector Bit)
-> (forall s. ST s (MVector s Bit)) -> Vector Bit
forall a b. (a -> b) -> a -> b
$ do
        MVector s Bit
ls_ <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Bit Int
n
        MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Bit
MVector (PrimState (ST s)) Bit
ls_ (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Bit -> ST s ()) -> Bit -> ST s ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bit
Bit Bool
False
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
          let si :: Int
si = Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
          let si1 :: Int
si1 = Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          Bit
b <-
            if Int
si Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
si1
              then MVector (PrimState (ST s)) Bit -> Int -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Bit
MVector (PrimState (ST s)) Bit
ls_ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              else Bit -> ST s Bit
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bit -> ST s Bit) -> (Bool -> Bit) -> Bool -> ST s Bit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bit
Bit (Bool -> ST s Bit) -> Bool -> ST s Bit
forall a b. (a -> b) -> a -> b
$ Int
si Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
si1
          MVector (PrimState (ST s)) Bit -> Int -> Bit -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Bit
MVector (PrimState (ST s)) Bit
ls_ Int
i Bit
b
        MVector s Bit -> ST s (MVector s Bit)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Bit
ls_

  let (!Vector Int
sumL, !Vector Int
sumS) = (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Vector Int))
 -> (Vector Int, Vector Int))
-> (forall s. ST s (Vector Int, Vector Int))
-> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
sumL_ <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
upper Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
0 :: Int)
        MVector s Int
sumS_ <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
upper Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
0 :: Int)
        (Bit -> Int -> ST s ()) -> Vector Bit -> Vector Int -> ST s ()
forall (m :: * -> *) a b c.
(Monad m, Unbox a, Unbox b) =>
(a -> b -> m c) -> Vector a -> Vector b -> m ()
VU.zipWithM_
          ( \(Bit !Bool
b) !Int
si -> do
              if Bool -> Bool
not Bool
b
                then MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
sumS_ (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
si
                else MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
sumL_ (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
si Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          )
          Vector Bit
ls
          Vector Int
s
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
upper] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
          Int
l <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
sumL_ Int
i
          MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
sumS_ (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l) Int
i
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
upper) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
si <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
sumS_ Int
i
            MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
sumL_ (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
si) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        Vector Int
sumL' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sumL_
        Vector Int
sumS' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sumS_
        (Vector Int, Vector Int) -> ST s (Vector Int, Vector Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector Int
sumL', Vector Int
sumS')

  MVector s Int
sa <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
n
  let induce :: Vector Int -> ST s ()
induce Vector Int
lms = do
        MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set MVector s Int
MVector (PrimState (ST s)) Int
sa (-Int
1)
        MVector s Int
buf <- Vector Int -> ST s (MVector (PrimState (ST s)) Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector Int
sumS
        Vector Int -> (Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector Int
lms ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
d -> do
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
i <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
buf (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
d
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf (Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
d) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sa Int
i Int
d
        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
sumL (MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf)
        do
          Int
i <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
buf (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf (Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sa Int
i (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        do
          -- TODO: try VUM.forM_
          [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
            Int
v <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
sa Int
i
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Bool -> Bool
not (Bit -> Bool
unBit (Vector Bit
ls Vector Bit -> Int -> Bit
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
              Int
j <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
buf (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf (Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sa Int
j (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
sumL (MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf)
        -- TODO: try foldr
        [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
          Int
v <- MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
sa Int
i
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 Bool -> Bool -> Bool
&& Bit -> Bool
unBit (Vector Bit
ls Vector Bit -> Int -> Bit
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Int
j <- Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1 (Int -> Int) -> ST s Int -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
buf (Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
buf (Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sa Int
j (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

  let (!Vector Int
lmsMap, !Int
m) = (forall s. ST s (Vector Int, Int)) -> (Vector Int, Int)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int, Int)) -> (Vector Int, Int))
-> (forall s. ST s (Vector Int, Int)) -> (Vector Int, Int)
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
lmsMap_ <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (-Int
1 :: Int)
        Int
len <-
          (Int -> Int -> (Bit, Bit) -> ST s Int)
-> Int -> Vector (Bit, Bit) -> ST s Int
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> Int -> b -> m a) -> a -> Vector b -> m a
VU.ifoldM'
            ( \Int
iAcc Int
i (Bit !Bool
b1, Bit !Bool
b2) -> do
                if Bool -> Bool
not Bool
b1 Bool -> Bool -> Bool
&& Bool
b2
                  then do
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
lmsMap_ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
iAcc
                    Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                  else do
                    Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iAcc
            )
            (Int
0 :: Int)
            (Vector (Bit, Bit) -> ST s Int) -> Vector (Bit, Bit) -> ST s Int
forall a b. (a -> b) -> a -> b
$ Vector Bit -> Vector Bit -> Vector (Bit, Bit)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Bit
ls (Vector Bit -> Vector Bit
forall a. Unbox a => Vector a -> Vector a
VU.tail Vector Bit
ls)
        (,Int
len) (Vector Int -> (Vector Int, Int))
-> ST s (Vector Int) -> ST s (Vector Int, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
lmsMap_

  let lms :: Vector Int
lms = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        MVector s Int
lms_ <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
m
        (Int -> Int -> (Bit, Bit) -> ST s Int)
-> Int -> Vector (Bit, Bit) -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> Int -> b -> m a) -> a -> Vector b -> m ()
VU.ifoldM'_
          ( \Int
iAcc Int
i (Bit !Bool
b1, Bit !Bool
b2) -> do
              if Bool -> Bool
not Bool
b1 Bool -> Bool -> Bool
&& Bool
b2
                then do
                  MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
lms_ Int
iAcc (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                else do
                  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iAcc
          )
          (Int
0 :: Int)
          (Vector (Bit, Bit) -> ST s ()) -> Vector (Bit, Bit) -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vector Bit -> Vector Bit -> Vector (Bit, Bit)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Bit
ls (Vector Bit -> Vector Bit
forall a. Unbox a => Vector a -> Vector a
VU.tail Vector Bit
ls)
        MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
lms_

  Vector Int -> ST s ()
induce Vector Int
lms

  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    Vector Int
sa' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sa
    MVector s Int
sortedLms <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
m
    (Int -> Int -> ST s Int) -> Int -> Vector Int -> ST s ()
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m ()
VU.foldM'_
      ( \Int
iAcc Int
v -> do
          if Vector Int
lmsMap Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1
            then do
              MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sortedLms Int
iAcc Int
v
              Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
iAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            else do
              Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iAcc
      )
      (Int
0 :: Int)
      Vector Int
sa'
    MVector s Int
recS <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew @_ @Int Int
m
    (\Int
i -> MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
recS Int
i Int
0) (Int -> ST s ()) -> (Int -> Int) -> Int -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector Int
lmsMap VG.!) (Int -> ST s ()) -> ST s Int -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Int -> Int -> ST s Int
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Int
MVector (PrimState (ST s)) Int
sortedLms Int
0

    Int
recUpper' <- do
      Vector Int
sortedLms' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sortedLms
      (Int -> (Int, Int) -> ST s Int)
-> Int -> Vector (Int, Int) -> ST s Int
forall (m :: * -> *) b a.
(Monad m, Unbox b) =>
(a -> b -> m a) -> a -> Vector b -> m a
VU.foldM'
        ( \Int
recUpper (!Int
l, !Int
r) -> do
            let lmsMapL :: Int
lmsMapL = Vector Int
lmsMap Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
l
            let lmsMapR :: Int
lmsMapR = Vector Int
lmsMap Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
r
            let endL :: Int
endL = if Int
lmsMapL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m then Vector Int
lms Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
lmsMapL Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) else Int
n
            let endR :: Int
endR = if Int
lmsMapR Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m then Vector Int
lms Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! (Int
lmsMapR Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) else Int
n
            let same :: Bool
same
                  | Int
endL Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
endR Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r = Bool
False
                  | Bool
otherwise = Int -> Int -> Bool
inner Int
l Int
r
                  where
                    inner :: Int -> Int -> Bool
inner Int
x Int
y
                      | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
endL Bool -> Bool -> Bool
&& Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y = Int -> Int -> Bool
inner (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                      | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
|| Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
y = Bool
False
                      | Bool
otherwise = Bool
True
            let recUpper' :: Int
recUpper' = if Bool
same then Int
recUpper else Int
recUpper Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
recS Int
lmsMapR Int
recUpper'
            Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
recUpper'
        )
        (Int
0 :: Int)
        (Vector (Int, Int) -> ST s Int) -> Vector (Int, Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
sortedLms' (Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector a
VU.tail Vector Int
sortedLms')

    Vector Int
recS' <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
recS
    let recSa :: Vector Int
recSa = HasCallStack => Int -> Int -> Vector Int -> Int -> Vector Int
Int -> Int -> Vector Int -> Int -> Vector Int
saIsManual Int
naiveThreshold Int
doublingThreshold Vector Int
recS' Int
recUpper'

    Vector Int -> (Int -> Int -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector Int
recSa ((Int -> Int -> ST s ()) -> ST s ())
-> (Int -> Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Int
x -> do
      MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
sortedLms Int
i (Vector Int
lms Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
x)
    Vector Int -> ST s ()
induce (Vector Int -> ST s ()) -> ST s (Vector Int) -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
sortedLms

  MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
sa

-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- SA-IS, linear-time suffix array construction.
-- Reference:
-- G. Nong, S. Zhang, and W. H. Chan,
-- Two Efficient Algorithms for Linear Time Suffix Array Construction
--
-- @since 1.0.0.0
{-# INLINE saIs #-}
saIs ::
  (HasCallStack) =>
  -- | string
  VU.Vector Int ->
  -- | upper bounds
  Int ->
  -- | suffix array
  VU.Vector Int
saIs :: HasCallStack => Vector Int -> Int -> Vector Int
saIs = HasCallStack => Int -> Int -> Vector Int -> Int -> Vector Int
Int -> Int -> Vector Int -> Int -> Vector Int
saIsManual Int
10 Int
40

-- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting).
--
-- SA-IS, linear-time suffix array construction.
-- Reference:
-- G. Nong, S. Zhang, and W. H. Chan,
-- Two Efficient Algorithms for Linear Time Suffix Array Construction
--
-- @since 1.0.0.0
{-# INLINE saIsManual #-}
saIsManual ::
  (HasCallStack) =>
  -- | naive threshold
  Int ->
  -- | doubling threshold
  Int ->
  -- | string
  VU.Vector Int ->
  -- | upper bounds
  Int ->
  -- | suffix array
  VU.Vector Int
saIsManual :: HasCallStack => Int -> Int -> Vector Int -> Int -> Vector Int
saIsManual Int
naiveThreshold Int
doublingThreshold Vector Int
s Int
upper
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector Int
forall a. Unbox a => Vector a
VU.empty
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Int -> Vector Int
forall a. Unbox a => a -> Vector a
VU.singleton Int
0
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
&& Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector Int
s Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
1 = Int -> [Int] -> Vector Int
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [Int
0, Int
1]
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 = Int -> [Int] -> Vector Int
forall a. Unbox a => Int -> [a] -> Vector a
VU.fromListN Int
2 [Int
1, Int
0]
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
naiveThreshold = HasCallStack => Vector Int -> Vector Int
Vector Int -> Vector Int
saNaive Vector Int
s
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
doublingThreshold = HasCallStack => Vector Int -> Vector Int
Vector Int -> Vector Int
saDoubling Vector Int
s
  | Bool
otherwise = HasCallStack => Int -> Int -> Vector Int -> Int -> Vector Int
Int -> Int -> Vector Int -> Int -> Vector Int
saIsImpl Int
naiveThreshold Int
doublingThreshold Vector Int
s Int
upper
  where
    n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
s