{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}

-- |
--
-- == Implementation
--
-- Suffix array construction uses the SAIS algorithm described by
--
-- * Ge Nong, Sen Zhang, and Wai Hong Chan,
-- /\"Linear Suffix Array Construction by Almost Pure Induced-Sorting\"/,
-- 2009 Data Compression Conference,
-- https://doi.org/10.1109/DCC.2009.42
--
-- LCP array construction uses the \(\varPhi\)-algorithm described by
--
-- * Juha Kärkkäinen, Giovanni Manzini, and Simon J. Puglisi
-- /\"Permuted Longest-Common-Prefix Array\"/,
-- Annual Symposium on Combinatorial Pattern Matching, 2009,
-- https://doi.org/10.1007/978-3-642-02441-2_17
--
-- The search algorithms used are described by
--
-- * Udi Manber and Gene Myers,
-- /\"Suffix arrays: a new method for on-line string searches\"/,
-- First annual ACM-SIAM symposium on Discrete algorithms, 1990, pp. 319-327,
-- https://dl.acm.org/doi/10.5555/320176.320218

module Data.Suffix
  (
    -- * Suffix array
    buildSuffixArray
  , SuffixArray(..)
  , search

    -- * Longest common prefix array
  , buildLCPArray
  , LCPArray(..)

    -- * LLCP and RLCP arrays
  , buildLRLCPArray
  , LRLCPArrays(..)
  , searchLRLCP

   -- * Suffix tree
  , foldSuffixTree

    -- * Pull
  , Pull(..)
  , pullFromByteString
  , pullFromPrimArray
  , pullFromArray
  , pullFromArrayLike

    -- * Intn
  , Intn
  ) where

import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.ST (ST, runST)
import Data.Bits ((.&.), (.|.), unsafeShiftL, unsafeShiftR)
import Data.Foldable (for_)
import Data.Int (Int32)
import qualified Data.Primitive.Array as A
import qualified Data.Primitive.PrimArray as PA
import Data.Primitive.Types (Prim)
import Data.Word (Word8)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS.Unsafe

#include "MachDeps.h"

----------
-- Links
----------

-- Nong's site with papers and code:
-- https://code.google.com/archive/p/ge-nong/downloads
-- The implementation here largely mirrors Nong's C++ implementation.

-- Useful explanation of SAIS, see "Constructing Suffix Arrays":
-- https://web.stanford.edu/class/archive/cs/cs166/cs166.1196/

-- Highly optimized C implementations to perhaps adopt optimizations from:
-- https://sites.google.com/site/yuta256/sais
-- https://github.com/IlyaGrebnov/libsais

----------
-- Types
----------

-- | Suffix array.
newtype SuffixArray i = SuffixArray (PA.PrimArray i)
  deriving (SuffixArray i -> SuffixArray i -> Bool
(SuffixArray i -> SuffixArray i -> Bool)
-> (SuffixArray i -> SuffixArray i -> Bool) -> Eq (SuffixArray i)
forall i. (Eq i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall i. (Eq i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
== :: SuffixArray i -> SuffixArray i -> Bool
$c/= :: forall i. (Eq i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
/= :: SuffixArray i -> SuffixArray i -> Bool
Eq, Eq (SuffixArray i)
Eq (SuffixArray i) =>
(SuffixArray i -> SuffixArray i -> Ordering)
-> (SuffixArray i -> SuffixArray i -> Bool)
-> (SuffixArray i -> SuffixArray i -> Bool)
-> (SuffixArray i -> SuffixArray i -> Bool)
-> (SuffixArray i -> SuffixArray i -> Bool)
-> (SuffixArray i -> SuffixArray i -> SuffixArray i)
-> (SuffixArray i -> SuffixArray i -> SuffixArray i)
-> Ord (SuffixArray i)
SuffixArray i -> SuffixArray i -> Bool
SuffixArray i -> SuffixArray i -> Ordering
SuffixArray i -> SuffixArray i -> SuffixArray i
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall i. (Ord i, Prim i) => Eq (SuffixArray i)
forall i. (Ord i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
forall i.
(Ord i, Prim i) =>
SuffixArray i -> SuffixArray i -> Ordering
forall i.
(Ord i, Prim i) =>
SuffixArray i -> SuffixArray i -> SuffixArray i
$ccompare :: forall i.
(Ord i, Prim i) =>
SuffixArray i -> SuffixArray i -> Ordering
compare :: SuffixArray i -> SuffixArray i -> Ordering
$c< :: forall i. (Ord i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
< :: SuffixArray i -> SuffixArray i -> Bool
$c<= :: forall i. (Ord i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
<= :: SuffixArray i -> SuffixArray i -> Bool
$c> :: forall i. (Ord i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
> :: SuffixArray i -> SuffixArray i -> Bool
$c>= :: forall i. (Ord i, Prim i) => SuffixArray i -> SuffixArray i -> Bool
>= :: SuffixArray i -> SuffixArray i -> Bool
$cmax :: forall i.
(Ord i, Prim i) =>
SuffixArray i -> SuffixArray i -> SuffixArray i
max :: SuffixArray i -> SuffixArray i -> SuffixArray i
$cmin :: forall i.
(Ord i, Prim i) =>
SuffixArray i -> SuffixArray i -> SuffixArray i
min :: SuffixArray i -> SuffixArray i -> SuffixArray i
Ord, Int -> SuffixArray i -> ShowS
[SuffixArray i] -> ShowS
SuffixArray i -> String
(Int -> SuffixArray i -> ShowS)
-> (SuffixArray i -> String)
-> ([SuffixArray i] -> ShowS)
-> Show (SuffixArray i)
forall i. (Show i, Prim i) => Int -> SuffixArray i -> ShowS
forall i. (Show i, Prim i) => [SuffixArray i] -> ShowS
forall i. (Show i, Prim i) => SuffixArray i -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall i. (Show i, Prim i) => Int -> SuffixArray i -> ShowS
showsPrec :: Int -> SuffixArray i -> ShowS
$cshow :: forall i. (Show i, Prim i) => SuffixArray i -> String
show :: SuffixArray i -> String
$cshowList :: forall i. (Show i, Prim i) => [SuffixArray i] -> ShowS
showList :: [SuffixArray i] -> ShowS
Show)

instance NFData (SuffixArray i) where
  rnf :: SuffixArray i -> ()
rnf !SuffixArray i
_ = ()

-- | Longest common prefix array.
newtype LCPArray i = LCPArray (PA.PrimArray i)
  deriving (LCPArray i -> LCPArray i -> Bool
(LCPArray i -> LCPArray i -> Bool)
-> (LCPArray i -> LCPArray i -> Bool) -> Eq (LCPArray i)
forall i. (Eq i, Prim i) => LCPArray i -> LCPArray i -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall i. (Eq i, Prim i) => LCPArray i -> LCPArray i -> Bool
== :: LCPArray i -> LCPArray i -> Bool
$c/= :: forall i. (Eq i, Prim i) => LCPArray i -> LCPArray i -> Bool
/= :: LCPArray i -> LCPArray i -> Bool
Eq, Eq (LCPArray i)
Eq (LCPArray i) =>
(LCPArray i -> LCPArray i -> Ordering)
-> (LCPArray i -> LCPArray i -> Bool)
-> (LCPArray i -> LCPArray i -> Bool)
-> (LCPArray i -> LCPArray i -> Bool)
-> (LCPArray i -> LCPArray i -> Bool)
-> (LCPArray i -> LCPArray i -> LCPArray i)
-> (LCPArray i -> LCPArray i -> LCPArray i)
-> Ord (LCPArray i)
LCPArray i -> LCPArray i -> Bool
LCPArray i -> LCPArray i -> Ordering
LCPArray i -> LCPArray i -> LCPArray i
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall i. (Ord i, Prim i) => Eq (LCPArray i)
forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Bool
forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Ordering
forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> LCPArray i
$ccompare :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Ordering
compare :: LCPArray i -> LCPArray i -> Ordering
$c< :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Bool
< :: LCPArray i -> LCPArray i -> Bool
$c<= :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Bool
<= :: LCPArray i -> LCPArray i -> Bool
$c> :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Bool
> :: LCPArray i -> LCPArray i -> Bool
$c>= :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> Bool
>= :: LCPArray i -> LCPArray i -> Bool
$cmax :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> LCPArray i
max :: LCPArray i -> LCPArray i -> LCPArray i
$cmin :: forall i. (Ord i, Prim i) => LCPArray i -> LCPArray i -> LCPArray i
min :: LCPArray i -> LCPArray i -> LCPArray i
Ord, Int -> LCPArray i -> ShowS
[LCPArray i] -> ShowS
LCPArray i -> String
(Int -> LCPArray i -> ShowS)
-> (LCPArray i -> String)
-> ([LCPArray i] -> ShowS)
-> Show (LCPArray i)
forall i. (Show i, Prim i) => Int -> LCPArray i -> ShowS
forall i. (Show i, Prim i) => [LCPArray i] -> ShowS
forall i. (Show i, Prim i) => LCPArray i -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall i. (Show i, Prim i) => Int -> LCPArray i -> ShowS
showsPrec :: Int -> LCPArray i -> ShowS
$cshow :: forall i. (Show i, Prim i) => LCPArray i -> String
show :: LCPArray i -> String
$cshowList :: forall i. (Show i, Prim i) => [LCPArray i] -> ShowS
showList :: [LCPArray i] -> ShowS
Show)

instance NFData (LCPArray i) where
  rnf :: LCPArray i -> ()
rnf !LCPArray i
_ = ()

-- | LLCP and RLCP arrays.
data LRLCPArrays i = LRLCPArrays
  {-# UNPACK #-} !(PA.PrimArray i)
  {-# UNPACK #-} !(PA.PrimArray i)
  deriving (LRLCPArrays i -> LRLCPArrays i -> Bool
(LRLCPArrays i -> LRLCPArrays i -> Bool)
-> (LRLCPArrays i -> LRLCPArrays i -> Bool) -> Eq (LRLCPArrays i)
forall i. (Eq i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall i. (Eq i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
== :: LRLCPArrays i -> LRLCPArrays i -> Bool
$c/= :: forall i. (Eq i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
/= :: LRLCPArrays i -> LRLCPArrays i -> Bool
Eq, Eq (LRLCPArrays i)
Eq (LRLCPArrays i) =>
(LRLCPArrays i -> LRLCPArrays i -> Ordering)
-> (LRLCPArrays i -> LRLCPArrays i -> Bool)
-> (LRLCPArrays i -> LRLCPArrays i -> Bool)
-> (LRLCPArrays i -> LRLCPArrays i -> Bool)
-> (LRLCPArrays i -> LRLCPArrays i -> Bool)
-> (LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i)
-> (LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i)
-> Ord (LRLCPArrays i)
LRLCPArrays i -> LRLCPArrays i -> Bool
LRLCPArrays i -> LRLCPArrays i -> Ordering
LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall i. (Ord i, Prim i) => Eq (LRLCPArrays i)
forall i. (Ord i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
forall i.
(Ord i, Prim i) =>
LRLCPArrays i -> LRLCPArrays i -> Ordering
forall i.
(Ord i, Prim i) =>
LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
$ccompare :: forall i.
(Ord i, Prim i) =>
LRLCPArrays i -> LRLCPArrays i -> Ordering
compare :: LRLCPArrays i -> LRLCPArrays i -> Ordering
$c< :: forall i. (Ord i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
< :: LRLCPArrays i -> LRLCPArrays i -> Bool
$c<= :: forall i. (Ord i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
<= :: LRLCPArrays i -> LRLCPArrays i -> Bool
$c> :: forall i. (Ord i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
> :: LRLCPArrays i -> LRLCPArrays i -> Bool
$c>= :: forall i. (Ord i, Prim i) => LRLCPArrays i -> LRLCPArrays i -> Bool
>= :: LRLCPArrays i -> LRLCPArrays i -> Bool
$cmax :: forall i.
(Ord i, Prim i) =>
LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
max :: LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
$cmin :: forall i.
(Ord i, Prim i) =>
LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
min :: LRLCPArrays i -> LRLCPArrays i -> LRLCPArrays i
Ord, Int -> LRLCPArrays i -> ShowS
[LRLCPArrays i] -> ShowS
LRLCPArrays i -> String
(Int -> LRLCPArrays i -> ShowS)
-> (LRLCPArrays i -> String)
-> ([LRLCPArrays i] -> ShowS)
-> Show (LRLCPArrays i)
forall i. (Show i, Prim i) => Int -> LRLCPArrays i -> ShowS
forall i. (Show i, Prim i) => [LRLCPArrays i] -> ShowS
forall i. (Show i, Prim i) => LRLCPArrays i -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall i. (Show i, Prim i) => Int -> LRLCPArrays i -> ShowS
showsPrec :: Int -> LRLCPArrays i -> ShowS
$cshow :: forall i. (Show i, Prim i) => LRLCPArrays i -> String
show :: LRLCPArrays i -> String
$cshowList :: forall i. (Show i, Prim i) => [LRLCPArrays i] -> ShowS
showList :: [LRLCPArrays i] -> ShowS
Show)

instance NFData (LRLCPArrays i) where
  rnf :: LRLCPArrays i -> ()
rnf !LRLCPArrays i
_ = ()

-----------------
-- Suffix array
-----------------

-- | \(O(n + k)\). Build a suffix array from a sequence.
--
-- On 64-bit systems, a @SuffixArray Int32@ requires half the memory as a
-- @SuffixArray Int@.
buildSuffixArray
  :: Intn i
  => Int            -- ^ The alphabet size \(k\).
  -> Pull Int       -- ^ Input sequence of length \(n\). Indexing is assumed to
                    --   be \(O(1)\). Elements must be in @[0..k-1]@.
  -> SuffixArray i  -- ^ Output type @i@ can be 'Int' or 'Int32'. If @i@ is
                    --   @Int32@, \(n\) must be @<= (maxBound :: Int32)@.
buildSuffixArray :: forall i. Intn i => Int -> Pull Int -> SuffixArray i
buildSuffixArray !Int
k p :: Pull Int
p@(Pull Int
n Int -> Int
_)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = SuffixArray i
forall a. a
errBuildSANKNegative
  | Bool
otherwise = PrimArray i -> SuffixArray i
forall i. PrimArray i -> SuffixArray i
SuffixArray (PrimArray i -> SuffixArray i) -> PrimArray i -> SuffixArray i
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s i)) -> PrimArray i
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
PA.runPrimArray ((forall s. ST s (MutablePrimArray s i)) -> PrimArray i)
-> (forall s. ST s (MutablePrimArray s i)) -> PrimArray i
forall a b. (a -> b) -> a -> b
$ do
    MutablePrimArray s i
out <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA 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
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
forall i s.
Intn i =>
Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
sais Int
k Pull Int
p MutablePrimArray s i
out Int
n
    MutablePrimArray s i -> ST s (MutablePrimArray s i)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutablePrimArray s i
out
{-# INLINE buildSuffixArray #-}

errBuildSANKNegative :: a
errBuildSANKNegative :: forall a. a
errBuildSANKNegative =
  String -> a
forall a. HasCallStack => String -> a
error String
"Data.Suffix.buildSuffixArray: n and k must be >= 0"

-- Manually specialize saisPrimArray with rules because SPECIALIZE doesn't work
-- https://gitlab.haskell.org/ghc/ghc/-/issues/25117
saisPrimArray
  :: Intn i
  => Int -> Int -> PA.PrimArray i -> PA.MutablePrimArray s i -> Int -> ST s ()
saisPrimArray :: forall i s.
Intn i =>
Int -> Int -> PrimArray i -> MutablePrimArray s i -> Int -> ST s ()
saisPrimArray Int
k Int
n PrimArray i
a MutablePrimArray s i
out Int
outn = Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
forall i s.
Intn i =>
Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
sais Int
k (Int -> (Int -> Int) -> Pull Int
forall a. Int -> (Int -> a) -> Pull a
Pull Int
n (i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> (Int -> i) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
a)) MutablePrimArray s i
out Int
outn
{-# NOINLINE saisPrimArray #-}

saisPrimArrayInt
  :: Int
  -> Int
  -> PA.PrimArray Int
  -> PA.MutablePrimArray s Int
  -> Int
  -> ST s ()
saisPrimArrayInt :: forall s.
Int
-> Int -> PrimArray Int -> MutablePrimArray s Int -> Int -> ST s ()
saisPrimArrayInt Int
k Int
n PrimArray Int
a MutablePrimArray s Int
out Int
outn = Int -> Pull Int -> MutablePrimArray s Int -> Int -> ST s ()
forall i s.
Intn i =>
Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
sais Int
k (Int -> (Int -> Int) -> Pull Int
forall a. Int -> (Int -> a) -> Pull a
Pull Int
n (PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray Int
a)) MutablePrimArray s Int
out Int
outn

saisPrimArrayInt32
  :: Int
  -> Int
  -> PA.PrimArray Int32
  -> PA.MutablePrimArray s Int32
  -> Int
  -> ST s ()
saisPrimArrayInt32 :: forall s.
Int
-> Int
-> PrimArray Int32
-> MutablePrimArray s Int32
-> Int
-> ST s ()
saisPrimArrayInt32 Int
k Int
n PrimArray Int32
a MutablePrimArray s Int32
out Int
outn = Int -> Pull Int -> MutablePrimArray s Int32 -> Int -> ST s ()
forall i s.
Intn i =>
Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
sais Int
k (Int -> (Int -> Int) -> Pull Int
forall a. Int -> (Int -> a) -> Pull a
Pull Int
n (Int32 -> Int
forall i. Intn i => i -> Int
toInt (Int32 -> Int) -> (Int -> Int32) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimArray Int32 -> Int -> Int32
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray Int32
a)) MutablePrimArray s Int32
out Int
outn

{-# RULES
"saisPrimArrayInt"   saisPrimArray = saisPrimArrayInt
"saisPrimArrayInt32" saisPrimArray = saisPrimArrayInt32
#-}

-- Inline sais into buildSuffixArray, saisPrimArrayInt, saisPrimArrayInt32. In
-- buildSuffixArray it should get optimized for whatever input sequence is used.
-- In saisPrimArray it is optimized for PrimArray Int/Int32.
{-# INLINE sais #-}
-- Precondition: n > 0
sais
  :: Intn i
  => Int                      -- alphabet size
  -> Pull Int                 -- input sequence
  -> PA.MutablePrimArray s i  -- buffer, also the output array
  -> Int                      -- length of `out` available
  -> ST s ()
sais :: forall i s.
Intn i =>
Int -> Pull Int -> MutablePrimArray s i -> Int -> ST s ()
sais !Int
k p :: Pull Int
p@(Pull Int
n Int -> Int
at) !MutablePrimArray s i
out !Int
outn = do

  -- L = False, S = True
  BitA
typ <- do
    BitMA s
typm <- Int -> ST s (BitMA s)
forall s. Int -> ST s (BitMA s)
newClearedBitMA Int
n
    Decr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Decr Int
Decr (Int
nInt -> 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
sufId -> do
      Bool
nxt <- BitMA s -> Int -> ST s Bool
forall s. BitMA s -> Int -> ST s Bool
readBitMA BitMA s
typm (Int
sufId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int -> Int
at Int
sufId Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int
at (Int
sufId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool -> Bool -> Bool
|| (Int -> Int
at Int
sufId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Int
at (Int
sufId Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bool -> Bool -> Bool
&& Bool
nxt)) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
        BitMA s -> Int -> ST s ()
forall s. BitMA s -> Int -> ST s ()
setBitMA BitMA s
typm Int
sufId
    BitMA s -> ST s BitA
forall s. BitMA s -> ST s BitA
unsafeFrzBitMA BitMA s
typm
  let isLMS :: Int -> Bool
isLMS Int
i = BitA -> Int -> Bool
indexBitA BitA
typ Int
i Bool -> Bool -> Bool
&& Bool -> Bool
not (BitA -> Int -> Bool
indexBitA BitA
typ (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))

  T2 MutSlice s i
buckets Int
outn1 <- Int
-> MutablePrimArray s i
-> Int
-> Int
-> ST s (T2 (MutSlice s i) Int)
forall i s.
Intn i =>
Int
-> MutablePrimArray s i
-> Int
-> Int
-> ST s (T2 (MutSlice s i) Int)
sharedOrNewSlice Int
n MutablePrimArray s i
out Int
outn Int
k
  Int -> Pull Int -> MutSlice s i -> ST s ()
forall i s. Intn i => Int -> Pull Int -> MutSlice s i -> ST s ()
fillBuckets Int
k Pull Int
p MutSlice s i
buckets
  T2 MutSlice s i
bucketIdx Int
_ <- Int
-> MutablePrimArray s i
-> Int
-> Int
-> ST s (T2 (MutSlice s i) Int)
forall i s.
Intn i =>
Int
-> MutablePrimArray s i
-> Int
-> Int
-> ST s (T2 (MutSlice s i) Int)
sharedOrNewSlice Int
n MutablePrimArray s i
out Int
outn1 Int
k

  MutablePrimArray s i -> Int -> Int -> i -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s i
out Int
0 Int
n (Int -> i
forall i. Intn i => Int -> i
frInt Int
emptyValue)
  MutSlice s i -> Int -> MutSlice s i -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice MutSlice s i
bucketIdx Int
0 MutSlice s i
buckets Int
0 Int
k
  Decr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Decr Int
Decr (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
1) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
sufId ->
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int -> Bool
isLMS Int
sufId) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      MutSlice s i -> Int -> (i -> ST s i) -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM MutSlice s i
bucketIdx (Int -> Int
at Int
sufId) ((i -> ST s i) -> ST s ()) -> (i -> ST s i) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \i
outIdx -> do
        let outIdx' :: i
outIdx' = i
outIdx i -> i -> i
forall a. Num a => a -> a -> a
- i
1
        i
outIdx' i -> ST s () -> ST s i
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (i -> Int
forall i. Intn i => i -> Int
toInt i
outIdx') (Int -> i
forall i. Intn i => Int -> i
frInt Int
sufId)

  let doFillL :: ST s ()
doFillL = Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
forall i s.
Intn i =>
Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
fillL Int
k Pull Int
p BitA
typ MutSlice s i
buckets MutSlice s i
bucketIdx MutablePrimArray s i
out
      doFillS :: ST s ()
doFillS = Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
forall i s.
Intn i =>
Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
fillS Int
k Pull Int
p BitA
typ MutSlice s i
buckets MutSlice s i
bucketIdx MutablePrimArray s i
out

  ST s ()
doFillL
  ST s ()
doFillS

  Int
numLMS <- Incr Int -> Int -> (Int -> Int -> ST s Int) -> ST s Int
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Incr Int
Incr Int
0 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Int
0 ((Int -> Int -> ST s Int) -> ST s Int)
-> (Int -> Int -> ST s Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \Int
j Int
i -> do
    i
sufId <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
    if i
sufId i -> i -> Bool
forall a. Eq a => a -> a -> Bool
/= i
0 Bool -> Bool -> Bool
&& Int -> Bool
isLMS (i -> Int
forall i. Intn i => i -> Int
toInt i
sufId)
    then (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> ST s () -> ST s Int
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out Int
j i
sufId
    else Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
j

  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
numLMS 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
    let !ndiv2 :: Int
ndiv2 = Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
        mapr :: Int -> Int
mapr Int
i = Int
ndiv2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1)

    MutablePrimArray s i -> Int -> Int -> i -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s i
out Int
ndiv2 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
ndiv2) (Int -> i
forall i. Intn i => Int -> i
frInt Int
emptyValue)

    do
      i
sufId0 <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
0
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (Int -> Int
mapr (i -> Int
forall i. Intn i => i -> Int
toInt i
sufId0)) i
0
    Int
lastName <- Incr Int -> Int -> (Int -> Int -> ST s Int) -> ST s Int
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Incr Int
Incr Int
1 (Int
numLMSInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Int
0 ((Int -> Int -> ST s Int) -> ST s Int)
-> (Int -> Int -> ST s Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \Int
prvName Int
i -> do
      Int
prvSufId <- i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> ST s i -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
      Int
sufId <- i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> ST s i -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
      let eqLoop :: Int -> Int -> Bool
eqLoop !Int
i1 !Int
i2
            | Int
i1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int
i2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Bool
False
            | Int -> Bool
isLMS Int
i1 Bool -> Bool -> Bool
|| Int -> Bool
isLMS Int
i2 = Int -> Int
at Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Int
at Int
i2
            | Bool
otherwise = Int -> Int
at Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Int
at Int
i2 Bool -> Bool -> Bool
&& Int -> Int -> Bool
eqLoop (Int
i1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
i2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
          name :: Int
name = if Int -> Int
at Int
prvSufId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Int
at Int
sufId Bool -> Bool -> Bool
&& Int -> Int -> Bool
eqLoop (Int
prvSufIdInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
sufIdInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                 then Int
prvName
                 else Int
prvName Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (Int -> Int
mapr Int
sufId) (Int -> i
forall i. Intn i => Int -> i
frInt Int
name)
      Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
name
    let numNames :: Int
numNames = Int
lastName Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
numNames Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
numLMS) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      Int
lastj <- Decr Int -> Int -> (Int -> Int -> ST s Int) -> ST s Int
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Decr Int
Decr (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
ndiv2) Int
n ((Int -> Int -> ST s Int) -> ST s Int)
-> (Int -> Int -> ST s Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \Int
j Int
i -> do
        i
name <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
        if i -> Int
forall i. Intn i => i -> Int
toInt i
name Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
emptyValue
        then Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
j
        else do
          let j' :: Int
j' = Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
          Int
j' Int -> ST s () -> ST s Int
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out Int
j' i
name

      PrimArray i
newa <- MutablePrimArray (PrimState (ST s)) i
-> Int -> Int -> ST s (PrimArray i)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
PA.freezePrimArray MutablePrimArray s i
MutablePrimArray (PrimState (ST s)) i
out Int
lastj (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lastj)
      Int -> Int -> PrimArray i -> MutablePrimArray s i -> Int -> ST s ()
forall i s.
Intn i =>
Int -> Int -> PrimArray i -> MutablePrimArray s i -> Int -> ST s ()
saisPrimArray Int
numNames Int
numLMS PrimArray i
newa MutablePrimArray s i
out Int
outn1

      Int
_ <- Decr Int -> Int -> (Int -> Int -> ST s Int) -> ST s Int
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Decr Int
Decr (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
1) (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ((Int -> Int -> ST s Int) -> ST s Int)
-> (Int -> Int -> ST s Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \Int
j Int
sufId ->
        if Int -> Bool
isLMS Int
sufId
        then (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> ST s () -> ST s Int
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out Int
j (Int -> i
forall i. Intn i => Int -> i
frInt Int
sufId)
        else Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
j
      Incr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Incr Int
Incr Int
0 (Int
numLMSInt -> 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
        i
j <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
        i
sufId <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
numLMS Int -> Int -> Int
forall a. Num a => a -> a -> a
+ i -> Int
forall i. Intn i => i -> Int
toInt i
j)
        MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out Int
i i
sufId

  MutSlice s i -> Int -> MutSlice s i -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice MutSlice s i
bucketIdx Int
0 MutSlice s i
buckets Int
0 Int
k
  MutablePrimArray s i -> Int -> Int -> i -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s i
out Int
numLMS (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
numLMS) (Int -> i
forall i. Intn i => Int -> i
frInt Int
emptyValue)
  Decr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Decr Int
Decr (Int
numLMSInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
0) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    i
sufId <- MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
    MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out Int
i (Int -> i
forall i. Intn i => Int -> i
frInt Int
emptyValue)
    MutSlice s i -> Int -> (i -> ST s i) -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM MutSlice s i
bucketIdx (Int -> Int
at (i -> Int
forall i. Intn i => i -> Int
toInt i
sufId)) ((i -> ST s i) -> ST s ()) -> (i -> ST s i) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \i
outIdx -> do
      let outIdx' :: i
outIdx' = i
outIdx i -> i -> i
forall a. Num a => a -> a -> a
- i
1
      i
outIdx' i -> ST s () -> ST s i
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (i -> Int
forall i. Intn i => i -> Int
toInt i
outIdx') i
sufId

  ST s ()
doFillL
  ST s ()
doFillS

emptyValue :: Int
emptyValue :: Int
emptyValue = -Int
1

sharedOrNewSlice
  :: Intn i
  => Int
  -> PA.MutablePrimArray s i
  -> Int
  -> Int
  -> ST s (T2 (MutSlice s i) Int)
sharedOrNewSlice :: forall i s.
Intn i =>
Int
-> MutablePrimArray s i
-> Int
-> Int
-> ST s (T2 (MutSlice s i) Int)
sharedOrNewSlice Int
n MutablePrimArray s i
out Int
outn Int
want
  | Int
outn Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
want =
    let !outn' :: Int
outn' = Int
outn Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
want
    in T2 (MutSlice s i) Int -> ST s (T2 (MutSlice s i) Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutSlice s i -> Int -> T2 (MutSlice s i) Int
forall a b. a -> b -> T2 a b
T2 (Int -> MutablePrimArray s i -> MutSlice s i
forall s a. Int -> MutablePrimArray s a -> MutSlice s a
MutSlice Int
outn' MutablePrimArray s i
out) Int
outn')
  | Bool
otherwise = do
    MutablePrimArray s i
a <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
want
    T2 (MutSlice s i) Int -> ST s (T2 (MutSlice s i) Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutSlice s i -> Int -> T2 (MutSlice s i) Int
forall a b. a -> b -> T2 a b
T2 (Int -> MutablePrimArray s i -> MutSlice s i
forall s a. Int -> MutablePrimArray s a -> MutSlice s a
MutSlice Int
0 MutablePrimArray s i
a) Int
outn)
{-# INLINE sharedOrNewSlice #-}

fillBuckets :: Intn i => Int -> Pull Int -> MutSlice s i -> ST s ()
fillBuckets :: forall i s. Intn i => Int -> Pull Int -> MutSlice s i -> ST s ()
fillBuckets Int
k (Pull Int
n Int -> Int
at) MutSlice s i
buckets = do
  MutSlice s i -> Int -> Int -> i -> ST s ()
forall a s. Prim a => MutSlice s a -> Int -> Int -> a -> ST s ()
setMutSlice MutSlice s i
buckets Int
0 Int
k i
0
  Incr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Incr Int
Incr Int
0 (Int
nInt -> 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 ->
    MutSlice s i -> Int -> (i -> i) -> ST s ()
forall a s. Prim a => MutSlice s a -> Int -> (a -> a) -> ST s ()
modifyMutSlice MutSlice s i
buckets (Int -> Int
at Int
i) (i -> i -> i
forall a. Num a => a -> a -> a
+i
1)
  i
_ <- Incr Int -> i -> (i -> Int -> ST s i) -> ST s i
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Incr Int
Incr Int
0 (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) i
0 ((i -> Int -> ST s i) -> ST s i) -> (i -> Int -> ST s i) -> ST s i
forall a b. (a -> b) -> a -> b
$ \i
acc Int
i -> do
    i
v <- MutSlice s i -> Int -> ST s i
forall a s. Prim a => MutSlice s a -> Int -> ST s a
readMutSlice MutSlice s i
buckets Int
i
    let acc' :: i
acc' = i
acc i -> i -> i
forall a. Num a => a -> a -> a
+ i
v
    MutSlice s i -> Int -> i -> ST s ()
forall a s. Prim a => MutSlice s a -> Int -> a -> ST s ()
writeMutSlice MutSlice s i
buckets Int
i i
acc'
    i -> ST s i
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure i
acc'
  () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
{-# INLINE fillBuckets #-}

fillL
  :: Intn i
  => Int
  -> Pull Int
  -> BitA
  -> MutSlice s i
  -> MutSlice s i
  -> PA.MutablePrimArray s i
  -> ST s ()
fillL :: forall i s.
Intn i =>
Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
fillL Int
k (Pull Int
n Int -> Int
at) BitA
typ MutSlice s i
buckets MutSlice s i
bucketIdx MutablePrimArray s i
out = do
  MutSlice s i -> Int -> i -> ST s ()
forall a s. Prim a => MutSlice s a -> Int -> a -> ST s ()
writeMutSlice MutSlice s i
bucketIdx Int
0 i
0
  MutSlice s i -> Int -> MutSlice s i -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice MutSlice s i
bucketIdx Int
1 MutSlice s i
buckets Int
0 (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  MutSlice s i -> Int -> (i -> ST s i) -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM MutSlice s i
bucketIdx (Int -> Int
at (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) ((i -> ST s i) -> ST s ()) -> (i -> ST s i) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \i
outIdx ->
    (i
outIdxi -> i -> i
forall a. Num a => a -> a -> a
+i
1) i -> ST s () -> ST s i
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (i -> Int
forall i. Intn i => i -> Int
toInt i
outIdx) (Int -> i
forall i. Intn i => Int -> i
frInt Int
n i -> i -> i
forall a. Num a => a -> a -> a
- i
1)
  Incr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Incr Int
Incr Int
0 (Int
nInt -> 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
sufId <- i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> ST s i -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sufId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
emptyValue) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      let sufIdL :: Int
sufIdL = Int
sufId Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sufIdL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Bool -> Bool
not (BitA -> Int -> Bool
indexBitA BitA
typ Int
sufIdL)) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        MutSlice s i -> Int -> (i -> ST s i) -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM MutSlice s i
bucketIdx (Int -> Int
at Int
sufIdL) ((i -> ST s i) -> ST s ()) -> (i -> ST s i) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \i
outIdx -> do
          (i
outIdxi -> i -> i
forall a. Num a => a -> a -> a
+i
1) i -> ST s () -> ST s i
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (i -> Int
forall i. Intn i => i -> Int
toInt i
outIdx) (Int -> i
forall i. Intn i => Int -> i
frInt Int
sufIdL)
{-# INLINE fillL #-}

fillS
  :: Intn i
  => Int
  -> Pull Int
  -> BitA
  -> MutSlice s i
  -> MutSlice s i
  -> PA.MutablePrimArray s i
  -> ST s ()
fillS :: forall i s.
Intn i =>
Int
-> Pull Int
-> BitA
-> MutSlice s i
-> MutSlice s i
-> MutablePrimArray s i
-> ST s ()
fillS Int
k (Pull Int
n Int -> Int
at) BitA
typ MutSlice s i
buckets MutSlice s i
bucketIdx MutablePrimArray s i
out = do
  MutSlice s i -> Int -> MutSlice s i -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice MutSlice s i
bucketIdx Int
0 MutSlice s i
buckets Int
0 Int
k
  Decr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Decr Int
Decr (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
0) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    Int
sufId <- i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> ST s i -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
out Int
i
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sufId Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
emptyValue) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
      let sufIdL :: Int
sufIdL = Int
sufId Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sufIdL Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& BitA -> Int -> Bool
indexBitA BitA
typ Int
sufIdL) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        MutSlice s i -> Int -> (i -> ST s i) -> ST s ()
forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM MutSlice s i
bucketIdx (Int -> Int
at Int
sufIdL) ((i -> ST s i) -> ST s ()) -> (i -> ST s i) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \i
outIdx -> do
          let outIdx' :: i
outIdx' = i
outIdx i -> i -> i
forall a. Num a => a -> a -> a
- i
1
          i
outIdx' i -> ST s () -> ST s i
forall a b. a -> ST s b -> ST s a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
out (i -> Int
forall i. Intn i => i -> Int
toInt i
outIdx') (Int -> i
forall i. Intn i => Int -> i
frInt Int
sufIdL)
{-# INLINE fillS #-}

-----------
-- Search
-----------

-- | \(O(m \log n)\). Search for a pattern in a sequence using its suffix array.
--
-- Note: For typical inputs, the worst case is unlikely and the running time is
-- close to \(O(m + \log n)\). To get guaranteed \(O(m + \log n)\) running time
-- consider using 'searchLRLCP' instead.
search
  :: (Ord a, Intn i)
  => Pull a           -- ^ Sequence of length \(n\). Indexing is assumed to
                      --   be \(O(1)\). @compare@ for @a@ is assumed to be
                      --   \(O(1)\).
  -> SuffixArray i    -- ^ Suffix array for the above sequence
  -> Pull a           -- ^ Pattern of length \(m\)
  -> (Int, Int)       -- ^ An @(offset, length)@ pair, denoting a slice of
                      --   the suffix array. Beginning at @offset@, @length@
                      --   suffixes start with the pattern. @length@ is 0 if the
                      --   pattern does not occur in the sequence.
search :: forall a i.
(Ord a, Intn i) =>
Pull a -> SuffixArray i -> Pull a -> (Int, Int)
search s :: Pull a
s@(Pull Int
n Int -> a
_) (SuffixArray PrimArray i
sa) !Pull a
t
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
sa = (Int, Int)
forall a. a
errSearchSizeMismatch
  | Bool
otherwise = let T3 Int
l Int
_ T2 Int Int
_ = Int
-> T2 Int Int
-> (Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int))
-> T3 Int Int (T2 Int Int)
forall s. Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch Int
n (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
0 Int
0) Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtl
                    T3 Int
_ Int
r T2 Int Int
_ = Int
-> T2 Int Int
-> (Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int))
-> T3 Int Int (T2 Int Int)
forall s. Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch Int
n (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
0 Int
0) Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtr
                    !off :: Int
off = Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
                    !len :: Int
len = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
off
                in (Int
off, Int
len)
  where
    doCmpSuffix :: Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
i = Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
forall a.
Ord a =>
Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
cmpSuffix Pull a
s Int
suf Pull a
t Int
i
    {-# NOINLINE doCmpSuffix #-}
    -- ^ Not inlining this apparently helps greatly, benchmark time reduces
    -- by 45-50%.

    nxtl :: Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtl !Int
m (T2 Int
llcp Int
rlcp) = case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
llcp Int
rlcp) of
      T2 Ordering
LT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
      T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
      where
        suf :: Int
suf = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
m)
    nxtr :: Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtr !Int
m (T2 Int
llcp Int
rlcp) = case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
llcp Int
rlcp) of
      T2 Ordering
GT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
      T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
      where
        suf :: Int
suf = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
m)
{-# INLINE search #-}

errSearchSizeMismatch :: a
errSearchSizeMismatch :: forall a. a
errSearchSizeMismatch = String -> a
forall a. HasCallStack => String -> a
error String
"Data.Suffix.search: size mismatch"

cmpSuffix :: Ord a => Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
cmpSuffix :: forall a.
Ord a =>
Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
cmpSuffix (Pull Int
n Int -> a
at) !Int
suf (Pull Int
n2 Int -> a
at2) = Int -> T2 Ordering Int
loop
  where
    loop :: Int -> T2 Ordering Int
loop Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n2 = Ordering -> Int -> T2 Ordering Int
forall a b. a -> b -> T2 a b
T2 Ordering
EQ Int
i
           | Int
sufInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Ordering -> Int -> T2 Ordering Int
forall a b. a -> b -> T2 a b
T2 Ordering
LT Int
i
           | Bool
otherwise = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> a
at (Int
sufInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)) (Int -> a
at2 Int
i) of
             Ordering
EQ -> Int -> T2 Ordering Int
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
             Ordering
o -> Ordering -> Int -> T2 Ordering Int
forall a b. a -> b -> T2 a b
T2 Ordering
o Int
i
{-# INLINE cmpSuffix #-}

getMid :: Int -> Int -> Int
getMid :: Int -> Int -> Int
getMid Int
l Int
h = Word -> Int
w2i (Int -> Word
i2w (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
h) Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1)
  where
    -- To Word to avoid overflow
    i2w :: Int -> Word
    i2w :: Int -> Word
i2w = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    w2i :: Word -> Int
    w2i :: Word -> Int
w2i = Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral

binarySearch :: Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch :: forall s. Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch Int
n s
s0 Int -> s -> E s s
nxt = Int -> Int -> s -> T3 Int Int s
go (-Int
1) Int
n s
s0
  where
    go :: Int -> Int -> s -> T3 Int Int s
go !Int
l !Int
r !s
s
      | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r = Int -> Int -> s -> T3 Int Int s
forall a b c. a -> b -> c -> T3 a b c
T3 Int
l Int
r s
s
      | Bool
otherwise =
        let !m :: Int
m = Int -> Int -> Int
getMid Int
l Int
r
        in case Int -> s -> E s s
nxt Int
m s
s of
          L s
s' -> Int -> Int -> s -> T3 Int Int s
go Int
l Int
m s
s'
          R s
s' -> Int -> Int -> s -> T3 Int Int s
go Int
m Int
r s
s'
{-# INLINE binarySearch #-}

--------------
-- LCP Array
--------------

-- | \(O(n)\). Build a longest common prefix array from a sequence and its
-- suffix array.
--
-- The LCP array has the same length as the sequence, \(n\). The \(0\)-th
-- element of the LCP array is \(0\). The \(i\)-th element of the LCP array for
-- \(0 < i < n\) is the longest common prefix of the \(i\)-th and \((i-1)\)-th
-- suffix in the suffix array.
buildLCPArray
  :: (Eq a, Intn i)
  => Pull a         -- ^ Sequence of length @n@. Indexing is assumed to be
                    --   \(O(1)\). @compare@ for @a@ is assumed to be \(O(1)\).
  -> SuffixArray i  -- ^ Suffix array for the above sequence
  -> LCPArray i
buildLCPArray :: forall a i. (Eq a, Intn i) => Pull a -> SuffixArray i -> LCPArray i
buildLCPArray (Pull Int
n Int -> a
at) (SuffixArray PrimArray i
sa)
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
sa = LCPArray i
forall a. a
errBuildLCPASizeMismatch
  | Bool
otherwise = PrimArray i -> LCPArray i
forall i. PrimArray i -> LCPArray i
LCPArray (PrimArray i -> LCPArray i) -> PrimArray i -> LCPArray i
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s i)) -> PrimArray i
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
PA.runPrimArray ((forall s. ST s (MutablePrimArray s i)) -> PrimArray i)
-> (forall s. ST s (MutablePrimArray s i)) -> PrimArray i
forall a b. (a -> b) -> a -> b
$ do
    MutablePrimArray s i
phi <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA 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
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
phi (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
0)) (-i
1)
    Incr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Incr Int
Incr Int
1 (Int
nInt -> 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 ->
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
phi (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
i)) (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
    MutablePrimArray s i
mplcpa <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
n
    Int
_ <- Incr Int -> Int -> (Int -> Int -> ST s Int) -> ST s Int
forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM (Int -> Int -> Incr Int
Incr Int
0 (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Int
0 ((Int -> Int -> ST s Int) -> ST s Int)
-> (Int -> Int -> ST s Int) -> ST s Int
forall a b. (a -> b) -> a -> b
$ \Int
l Int
i -> do
      Int
j <- i -> Int
forall i. Intn i => i -> Int
toInt (i -> Int) -> ST s i -> ST s Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray s i -> Int -> ST s i
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s i
phi Int
i
      let diff :: Int -> Bool
diff Int
d = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n Bool -> Bool -> Bool
|| Int -> a
at (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> a
at (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d)
          l' :: Int
l' = if Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1
               then Int
0
               else (Int -> Bool) -> (Int -> Int) -> Int -> Int
forall a. (a -> Bool) -> (a -> a) -> a -> a
until Int -> Bool
diff (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
l
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
mplcpa Int
i (Int -> i
forall i. Intn i => Int -> i
frInt Int
l')
      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 -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    PrimArray i
plcpa <- MutablePrimArray (PrimState (ST s)) i -> ST s (PrimArray i)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
PA.unsafeFreezePrimArray MutablePrimArray s i
MutablePrimArray (PrimState (ST s)) i
mplcpa
    MutablePrimArray s i
lcpa <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
n
    Incr Int -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Int -> Int -> Incr Int
Incr Int
0 (Int
nInt -> 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 ->
      MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
lcpa Int
i (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
plcpa (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
i)))
    MutablePrimArray s i -> ST s (MutablePrimArray s i)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutablePrimArray s i
lcpa
{-# INLINE buildLCPArray #-}

errBuildLCPASizeMismatch :: a
errBuildLCPASizeMismatch :: forall a. a
errBuildLCPASizeMismatch = String -> a
forall a. HasCallStack => String -> a
error String
"Data.Suffix.buildLCPArray: size mismatch"

--------------
-- LLCP RLCP
--------------

-- | \(O(n)\). Build LLCP and RLCP arrays from an LCP array, for use in
-- @searchLRLCP@.
buildLRLCPArray :: Intn i => LCPArray i -> LRLCPArrays i
buildLRLCPArray :: forall i. Intn i => LCPArray i -> LRLCPArrays i
buildLRLCPArray (LCPArray PrimArray i
lcpa) = (forall s. ST s (LRLCPArrays i)) -> LRLCPArrays i
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (LRLCPArrays i)) -> LRLCPArrays i)
-> (forall s. ST s (LRLCPArrays i)) -> LRLCPArrays i
forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s i
llcpa <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
n
  MutablePrimArray s i
rlcpa <- Int -> ST s (MutablePrimArray s i)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
n
  let go :: Int -> Int -> ST s Int
go Int
l Int
r
        | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r = if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
                       then Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
                       else 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
$! i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
lcpa Int
r)
        | Bool
otherwise = do
          let m :: Int
m = Int -> Int -> Int
getMid Int
l Int
r
          Int
llcp <- Int -> Int -> ST s Int
go Int
l Int
m
          MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
llcpa Int
m (Int -> i
forall i. Intn i => Int -> i
frInt Int
llcp)
          Int
rlcp <- Int -> Int -> ST s Int
go Int
m Int
r
          MutablePrimArray s i -> Int -> i -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s i
rlcpa Int
m (Int -> i
forall i. Intn i => Int -> i
frInt Int
rlcp)
          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 -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
llcp Int
rlcp
  Int
_ <- Int -> Int -> ST s Int
go (-Int
1) Int
n
  PrimArray i -> PrimArray i -> LRLCPArrays i
forall i. PrimArray i -> PrimArray i -> LRLCPArrays i
LRLCPArrays
    (PrimArray i -> PrimArray i -> LRLCPArrays i)
-> ST s (PrimArray i) -> ST s (PrimArray i -> LRLCPArrays i)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState (ST s)) i -> ST s (PrimArray i)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
PA.unsafeFreezePrimArray MutablePrimArray s i
MutablePrimArray (PrimState (ST s)) i
llcpa
    ST s (PrimArray i -> LRLCPArrays i)
-> ST s (PrimArray i) -> ST s (LRLCPArrays i)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MutablePrimArray (PrimState (ST s)) i -> ST s (PrimArray i)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
PA.unsafeFreezePrimArray MutablePrimArray s i
MutablePrimArray (PrimState (ST s)) i
rlcpa
  where
    n :: Int
n = PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
lcpa
{-# SPECIALIZE buildLRLCPArray :: LCPArray Int -> LRLCPArrays Int #-}
{-# SPECIALIZE buildLRLCPArray :: LCPArray Int32 -> LRLCPArrays Int32 #-}

-- | \(O(m + \log n)\). Search for a pattern in a sequence using its suffix
-- array and LLCP and RLCP arrays.
searchLRLCP
  :: (Ord a, Intn i)
  => Pull a           -- ^ Sequence of length \(n\). Indexing is assumed to
                      --   be \(O(1)\). @compare@ for @a@ is assumed to be
                      --   \(O(1)\).
  -> SuffixArray i    -- ^ Suffix array for the above sequence
  -> LRLCPArrays i    -- ^ LLCP and RLCP arrays for the above sequence
  -> Pull a           -- ^ Pattern sequence of length \(m\)
  -> (Int, Int)       -- ^ An @(offset, length)@ pair, denoting a slice of
                      --   the suffix array. Beginning at @offset@, @length@
                      --   suffixes start with the pattern. @length@ is 0 if the
                      --   pattern does not occur in the sequence.
searchLRLCP :: forall a i.
(Ord a, Intn i) =>
Pull a -> SuffixArray i -> LRLCPArrays i -> Pull a -> (Int, Int)
searchLRLCP s :: Pull a
s@(Pull Int
n Int -> a
_) (SuffixArray PrimArray i
sa) (LRLCPArrays PrimArray i
llcpa PrimArray i
rlcpa) !Pull a
t
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
sa Bool -> Bool -> Bool
||
    Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
llcpa Bool -> Bool -> Bool
||
    Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
rlcpa
  = (Int, Int)
forall a. a
errSearchLRLCPSizeMismatch
  | Bool
otherwise = let T3 Int
l Int
_ T2 Int Int
_ = Int
-> T2 Int Int
-> (Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int))
-> T3 Int Int (T2 Int Int)
forall s. Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch Int
n (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
0 Int
0) Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtl
                    T3 Int
_ Int
r T2 Int Int
_ = Int
-> T2 Int Int
-> (Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int))
-> T3 Int Int (T2 Int Int)
forall s. Int -> s -> (Int -> s -> E s s) -> T3 Int Int s
binarySearch Int
n (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
0 Int
0) Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtr
                    !off :: Int
off = Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
                    !len :: Int
len = Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
off
                in (Int
off, Int
len)
  where
    doCmpSuffix :: Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
i = Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
forall a.
Ord a =>
Pull a -> Int -> Pull a -> Int -> T2 Ordering Int
cmpSuffix Pull a
s Int
suf Pull a
t Int
i
    {-# NOINLINE doCmpSuffix #-}
    -- ^ Not inlining this apparently helps greatly, benchmark time reduces
    -- by 45-50%.

    nxtl :: Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtl !Int
m (T2 Int
llcp Int
rlcp)
      | Int
llcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rlcp =
        let mllcp :: Int
mllcp = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
llcpa Int
m)
        in if Int
llcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
mllcp
           then T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
mllcp)
           else case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
llcp of
             T2 Ordering
LT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
             T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
      | Bool
otherwise =
        let mrlcp :: Int
mrlcp = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
rlcpa Int
m)
        in if Int
rlcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
mrlcp
           then T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
mrlcp Int
rlcp)
           else case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
rlcp of
             T2 Ordering
LT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
             T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
      where
        suf :: Int
suf = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
m)

    nxtr :: Int -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
nxtr !Int
m (T2 Int
llcp Int
rlcp)
      | Int
llcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rlcp =
        let mllcp :: Int
mllcp = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
llcpa Int
m)
        in if Int
llcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
mllcp
           then T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
mllcp)
           else case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
llcp of
             T2 Ordering
GT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
             T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
      | Bool
otherwise =
        let mrlcp :: Int
mrlcp = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
rlcpa Int
m)
        in if Int
rlcp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
mrlcp
           then T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
mrlcp Int
rlcp)
           else case Int -> Int -> T2 Ordering Int
doCmpSuffix Int
suf Int
rlcp of
             T2 Ordering
GT Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. a -> E a b
L (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
llcp Int
lcp')
             T2 Ordering
_ Int
lcp' -> T2 Int Int -> E (T2 Int Int) (T2 Int Int)
forall a b. b -> E a b
R (Int -> Int -> T2 Int Int
forall a b. a -> b -> T2 a b
T2 Int
lcp' Int
rlcp)
      where
        suf :: Int
suf = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
m)
{-# INLINE searchLRLCP #-}

errSearchLRLCPSizeMismatch :: a
errSearchLRLCPSizeMismatch :: forall a. a
errSearchLRLCPSizeMismatch = String -> a
forall a. HasCallStack => String -> a
error String
"Data.Suffix.searchLRLCP: size mismatch"

----------------
-- Suffix tree
----------------

-- | \(O(n)\). Bottom-up strict monadic fold of the suffix tree formed by the
-- given suffix array and LCP array.
--
-- The tree folded over can be considered to be equivalent to the following
-- definition.
--
-- @
-- data SuffixTree
--   = Leaf
--       Int          -- ^ The suffix index
--   | Branch
--       Int          -- ^ The string depth of this node
--       [SuffixTree] -- ^ The children of this node
-- @
--
-- This tree has exactly \(n+1\) leaves and at most \(n\) internal nodes.
-- The suffix with index \(n\), i.e. the empty suffix, is a leaf in this tree,
-- though it is not present in the suffix array.
--
-- The calls to the leaf function and the internal node finalizer taken together
-- form a post-order traversal of the tree. For an internal node, the values
-- from its children are combined left-to-right.
--
-- The \(O(n)\) bound assumes that monadic bind and all given functions are
-- \(O(1)\).
foldSuffixTree
  :: (Intn i, Monad m)
  => (Int -> m a)           -- ^ Leaf. The @Int@ is the suffix index.
  -> (Int -> m b)           -- ^ Internal node, initialize. The @Int@ is the
                            --   node's string depth.
  -> (Int -> b -> a -> m b) -- ^ Internal node, combine with the value from a
                            --   child. The @Int@ is the node's string depth.
  -> (Int -> b -> m a)      -- ^ Internal node, finalize. The @Int@ is the
                            --   node's string depth.
  -> SuffixArray i          -- ^ A suffix array derived from a sequence
  -> LCPArray i             -- ^ The LCP array for the same sequence and suffix
                            --   array
  -> m a
foldSuffixTree :: forall i (m :: * -> *) a b.
(Intn i, Monad m) =>
(Int -> m a)
-> (Int -> m b)
-> (Int -> b -> a -> m b)
-> (Int -> b -> m a)
-> SuffixArray i
-> LCPArray i
-> m a
foldSuffixTree
  Int -> m a
leaf Int -> m b
branchInit Int -> b -> a -> m b
branchCombine Int -> b -> m a
branchFinish (SuffixArray PrimArray i
sa) (LCPArray PrimArray i
lcpa)
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
lcpa = m a
forall a. a
errFoldSTSizeMismatch
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> m a
leaf Int
n
  | Bool
otherwise = do
    !a
aLeft <- Int -> m a
leaf Int
n
    !b
b <- Int -> m b
branchInit Int
0
    !b
b' <- Int -> b -> a -> m b
branchCombine Int
0 b
b a
aLeft
    Int -> T3 Int b (Stack b) -> m a
down Int
0 (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
0 b
b' Stack b
forall a. Stack a
Nil)
  where
    n :: Int
n = PrimArray i -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray i
sa
    down :: Int -> T3 Int b (Stack b) -> m a
down !Int
i (T3 Int
d b
b Stack b
stk1) = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
lcpCur Int
lcpNxt of
      Ordering
LT | Bool
addLeft -> do
           !a
aLeft <- Int -> m a
leaf (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)))
           !b
b' <- Int -> b -> a -> m b
branchCombine Int
d b
b a
aLeft
           !b
b1 <- Int -> m b
branchInit Int
lcpNxt
           Int -> T3 Int b (Stack b) -> m a
down (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
lcpNxt b
b1 (Int -> b -> Stack b -> Stack b
forall a. Int -> a -> Stack a -> Stack a
Push Int
d b
b' Stack b
stk1))
         | Bool
otherwise -> do
           !b
b1 <- Int -> m b
branchInit Int
lcpNxt
           Int -> T3 Int b (Stack b) -> m a
down (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
lcpNxt b
b1 (Int -> b -> Stack b -> Stack b
forall a. Int -> a -> Stack a -> Stack a
Push Int
d b
b Stack b
stk1))
      Ordering
EQ | Bool
addLeft -> do
           !a
aLeft <- Int -> m a
leaf (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)))
           !b
b' <- Int -> b -> a -> m b
branchCombine Int
d b
b a
aLeft
           Int -> T3 Int b (Stack b) -> m a
down (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
d b
b' Stack b
stk1)
         | Bool
otherwise -> Int -> T3 Int b (Stack b) -> m a
down (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
d b
b Stack b
stk1)
      Ordering
GT -> do
        !b
b' <-
              if Bool
addLeft
              then do
                !a
aLeft <- Int -> m a
leaf (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)))
                Int -> b -> a -> m b
branchCombine Int
d b
b a
aLeft
              else b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b
        !a
aRight <- Int -> m a
leaf (i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
sa Int
i))
        !b
b'' <- Int -> b -> a -> m b
branchCombine Int
d b
b' a
aRight
        !a
a <- Int -> b -> m a
branchFinish Int
d b
b''
        Int -> Int -> a -> Stack b -> m a
up (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
lcpNxt a
a Stack b
stk1
      where
        lcpPrv :: Int
lcpPrv = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
lcpa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) else Int
n
        lcpCur :: Int
lcpCur = i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
lcpa Int
i)
        lcpNxt :: Int
lcpNxt = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 then i -> Int
forall i. Intn i => i -> Int
toInt (PrimArray i -> Int -> i
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray i
lcpa (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)) else (-Int
1)
        addLeft :: Bool
addLeft = Int
lcpPrv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
lcpCur

    up :: Int -> Int -> a -> Stack b -> m a
up !Int
i !Int
dep = a -> Stack b -> m a
go
      where
        go :: a -> Stack b -> m a
go !a
a stk :: Stack b
stk@(Push Int
d b
b Stack b
stk1) = do
          case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
d Int
dep of
            Ordering
LT -> do
              !b
b1 <- Int -> m b
branchInit Int
dep
              !b
b1' <- Int -> b -> a -> m b
branchCombine Int
dep b
b1 a
a
              Int -> T3 Int b (Stack b) -> m a
down Int
i (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
dep b
b1' Stack b
stk)
            Ordering
EQ -> do
              !b
b' <- Int -> b -> a -> m b
branchCombine Int
d b
b a
a
              Int -> T3 Int b (Stack b) -> m a
down Int
i (Int -> b -> Stack b -> T3 Int b (Stack b)
forall a b c. a -> b -> c -> T3 a b c
T3 Int
d b
b' Stack b
stk1)
            Ordering
GT -> do
              !b
b' <- Int -> b -> a -> m b
branchCombine Int
d b
b a
a
              !a
a' <- Int -> b -> m a
branchFinish Int
d b
b'
              a -> Stack b -> m a
go a
a' Stack b
stk1
        go !a
a Stack b
Nil = a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
{-# INLINE foldSuffixTree #-}

errFoldSTSizeMismatch :: a
errFoldSTSizeMismatch :: forall a. a
errFoldSTSizeMismatch = String -> a
forall a. HasCallStack => String -> a
error String
"Data.Suffix.foldSuffixTree: size mismatch"

data Stack a
  = Push {-# UNPACK #-} !Int !a !(Stack a)
  | Nil

---------
-- Pull
---------

-- | A pull array. Serves as a simple interface for array-like structures.
--
-- Note: For good performance, create @Pull@s right before supplying them to
-- functions in this module.
data Pull a = Pull
  !Int       -- ^ Length \(n\). \(n\) must be \(\geq 0\).
  (Int -> a) -- ^ Indexing function. Must be valid for inputs in @[0 .. n-1]@

instance Functor Pull where
  fmap :: forall a b. (a -> b) -> Pull a -> Pull b
fmap a -> b
f (Pull Int
n Int -> a
at) = Int -> (Int -> b) -> Pull b
forall a. Int -> (Int -> a) -> Pull a
Pull Int
n (a -> b
f (a -> b) -> (Int -> a) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
at)
  a
x <$ :: forall a b. a -> Pull b -> Pull a
<$ Pull Int
n Int -> b
_ = Int -> (Int -> a) -> Pull a
forall a. Int -> (Int -> a) -> Pull a
Pull Int
n (a -> Int -> a
forall a b. a -> b -> a
const a
x)

-- | Pull from a @ByteString@ (from the bytestring package).
pullFromByteString :: BS.ByteString -> Pull Word8
pullFromByteString :: ByteString -> Pull Word8
pullFromByteString =
  (ByteString -> Int)
-> (ByteString -> Int -> Word8) -> ByteString -> Pull Word8
forall arr a. (arr -> Int) -> (arr -> Int -> a) -> arr -> Pull a
pullFromArrayLike
    ByteString -> Int
BS.length
#ifdef CHECKS
    BS.index
#else
    ByteString -> Int -> Word8
BS.Unsafe.unsafeIndex
#endif
{-# INLINE pullFromByteString #-}

-- | Pull from a @PrimArray@ (from the primitive package).
pullFromPrimArray :: Prim a => PA.PrimArray a -> Pull a
pullFromPrimArray :: forall a. Prim a => PrimArray a -> Pull a
pullFromPrimArray = (PrimArray a -> Int)
-> (PrimArray a -> Int -> a) -> PrimArray a -> Pull a
forall arr a. (arr -> Int) -> (arr -> Int -> a) -> arr -> Pull a
pullFromArrayLike PrimArray a -> Int
forall a. Prim a => PrimArray a -> Int
PA.sizeofPrimArray PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPA
{-# INLINE pullFromPrimArray #-}

-- | Pull from an @Array@ (from the primitive package).
pullFromArray :: A.Array a -> Pull a
pullFromArray :: forall a. Array a -> Pull a
pullFromArray =
  (Array a -> Int) -> (Array a -> Int -> a) -> Array a -> Pull a
forall arr a. (arr -> Int) -> (arr -> Int -> a) -> arr -> Pull a
pullFromArrayLike
    Array a -> Int
forall a. Array a -> Int
A.sizeofArray
#ifdef CHECKS
    (\a i -> check "pullFromArray"
                   (0 <= i && i < A.sizeofArray a)
                   (A.indexArray a i))
#else
    Array a -> Int -> a
forall a. Array a -> Int -> a
A.indexArray
#endif
{-# INLINE pullFromArray #-}

-- | Pull elements from any array-like structure.
pullFromArrayLike
  :: (arr -> Int)      -- ^ Size function
  -> (arr -> Int -> a) -- ^ Indexing function
  -> arr               -- ^ The structure
  -> Pull a
pullFromArrayLike :: forall arr a. (arr -> Int) -> (arr -> Int -> a) -> arr -> Pull a
pullFromArrayLike arr -> Int
size arr -> Int -> a
index !arr
a = Int -> (Int -> a) -> Pull a
forall a. Int -> (Int -> a) -> Pull a
Pull (arr -> Int
size arr
a) (arr -> Int -> a
index arr
a)
{-# INLINE pullFromArrayLike #-}

----------
-- Intn
----------

-- | Allows for a choice between @Int@ and @Int32@.
class (Prim i, Integral i) => Intn i where
  toInt :: i -> Int
  frInt :: Int -> i

instance Intn Int where
  toInt :: Int -> Int
toInt = Int -> Int
forall a. a -> a
id
  frInt :: Int -> Int
frInt = Int -> Int
forall a. a -> a
id

instance Intn Int32 where
  toInt :: Int32 -> Int
toInt = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
  frInt :: Int -> Int32
frInt = Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

--------------
-- Bit array
--------------

newtype BitMA s = BitMA (PA.MutablePrimArray s Int)
newtype BitA = BitA (PA.PrimArray Int)

wshift, wmask :: Int
#if SIZEOF_HSWORD == 4
wshift = 5
wmask = 31
#elif SIZEOF_HSWORD == 8
wshift :: Int
wshift = Int
6
wmask :: Int
wmask = Int
63
#else
#error "unsupported word size"
#endif

newClearedBitMA :: Int -> ST s (BitMA s)
newClearedBitMA :: forall s. Int -> ST s (BitMA s)
newClearedBitMA Int
n = do
  let wn :: Int
wn = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wmask) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
wshift
  MutablePrimArray s Int
a <- Int -> ST s (MutablePrimArray s Int)
forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
wn
  MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s Int
a Int
0 Int
wn Int
0
  BitMA s -> ST s (BitMA s)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutablePrimArray s Int -> BitMA s
forall s. MutablePrimArray s Int -> BitMA s
BitMA MutablePrimArray s Int
a)
{-# INLINE newClearedBitMA #-}

setBitMA :: BitMA s -> Int -> ST s ()
setBitMA :: forall s. BitMA s -> Int -> ST s ()
setBitMA (BitMA MutablePrimArray s Int
a) Int
i = do
  let j :: Int
j = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
wshift
  Int
x <- MutablePrimArray s Int -> Int -> ST s Int
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s Int
a Int
j
  MutablePrimArray s Int -> Int -> Int -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s Int
a Int
j (Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
wmask)))
{-# INLINE setBitMA #-}

readBitMA :: BitMA s -> Int -> ST s Bool
readBitMA :: forall s. BitMA s -> Int -> ST s Bool
readBitMA (BitMA MutablePrimArray s Int
a) Int
i = do
  let j :: Int
j = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
wshift
  Int
x <- MutablePrimArray s Int -> Int -> ST s Int
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s Int
a Int
j
  Bool -> ST s Bool
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> ST s Bool) -> Bool -> ST s Bool
forall a b. (a -> b) -> a -> b
$ Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
wmask)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
{-# INLINE readBitMA #-}

unsafeFrzBitMA :: BitMA s -> ST s BitA
unsafeFrzBitMA :: forall s. BitMA s -> ST s BitA
unsafeFrzBitMA (BitMA MutablePrimArray s Int
a) = PrimArray Int -> BitA
BitA (PrimArray Int -> BitA) -> ST s (PrimArray Int) -> ST s BitA
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState (ST s)) Int -> ST s (PrimArray Int)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
PA.unsafeFreezePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
a
{-# INLINE unsafeFrzBitMA #-}

indexBitA :: BitA -> Int -> Bool
indexBitA :: BitA -> Int -> Bool
indexBitA (BitA PrimArray Int
a) Int
i =
  let j :: Int
j = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
wshift
  in PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray Int
a Int
j Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
wmask)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
{-# INLINE indexBitA #-}

-------------
-- MutSlice
-------------

data MutSlice s a = MutSlice
  {-# UNPACK #-} !Int
  {-# UNPACK #-} !(PA.MutablePrimArray s a)

readMutSlice :: Prim a => MutSlice s a -> Int -> ST s a
readMutSlice :: forall a s. Prim a => MutSlice s a -> Int -> ST s a
readMutSlice (MutSlice Int
off MutablePrimArray s a
a) Int
i = MutablePrimArray s a -> Int -> ST s a
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s a
a (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)
{-# INLINE readMutSlice #-}

writeMutSlice :: Prim a => MutSlice s a -> Int -> a -> ST s ()
writeMutSlice :: forall a s. Prim a => MutSlice s a -> Int -> a -> ST s ()
writeMutSlice (MutSlice Int
off MutablePrimArray s a
a) Int
i = MutablePrimArray s a -> Int -> a -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s a
a (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)
{-# INLINE writeMutSlice #-}

modifyMutSlice :: Prim a => MutSlice s a -> Int -> (a -> a) -> ST s ()
modifyMutSlice :: forall a s. Prim a => MutSlice s a -> Int -> (a -> a) -> ST s ()
modifyMutSlice (MutSlice Int
off MutablePrimArray s a
a) Int
i = MutablePrimArray s a -> Int -> (a -> a) -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> (a -> a) -> ST s ()
modifyPA MutablePrimArray s a
a (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)
{-# INLINE modifyMutSlice #-}

setMutSlice :: Prim a => MutSlice s a -> Int -> Int -> a -> ST s ()
setMutSlice :: forall a s. Prim a => MutSlice s a -> Int -> Int -> a -> ST s ()
setMutSlice (MutSlice Int
off MutablePrimArray s a
a) Int
i = MutablePrimArray s a -> Int -> Int -> a -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s a
a (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)
{-# INLINE setMutSlice #-}

copyMutSlice
  :: Prim a => MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice :: forall a s.
Prim a =>
MutSlice s a -> Int -> MutSlice s a -> Int -> Int -> ST s ()
copyMutSlice (MutSlice Int
doff MutablePrimArray s a
dst) Int
di (MutSlice Int
soff MutablePrimArray s a
src) Int
si =
  MutablePrimArray s a
-> Int -> MutablePrimArray s a -> Int -> Int -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a
-> Int -> MutablePrimArray s a -> Int -> Int -> ST s ()
copyMutPA MutablePrimArray s a
dst (Int
doffInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
di) MutablePrimArray s a
src (Int
soffInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
si)
{-# INLINE copyMutSlice #-}

modifyMutSliceM :: Prim a => MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM :: forall a s.
Prim a =>
MutSlice s a -> Int -> (a -> ST s a) -> ST s ()
modifyMutSliceM (MutSlice Int
off MutablePrimArray s a
a) Int
i = MutablePrimArray s a -> Int -> (a -> ST s a) -> ST s ()
forall a s.
Prim a =>
MutablePrimArray s a -> Int -> (a -> ST s a) -> ST s ()
modifyPAM MutablePrimArray s a
a (Int
offInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)
{-# INLINE modifyMutSliceM #-}

----------------------
-- Primitive helpers
----------------------

indexPA :: Prim a => PA.PrimArray a -> Int -> a
indexPA :: forall a. Prim a => PrimArray a -> Int -> a
indexPA PrimArray a
a Int
i =
#ifdef CHECKS
  check "indexPA" (0 <= i && i < PA.sizeofPrimArray a) $
#endif
  PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
PA.indexPrimArray PrimArray a
a Int
i
{-# INLINE indexPA #-}

newPA :: Prim a => Int -> ST s (PA.MutablePrimArray s a)
newPA :: forall a s. Prim a => Int -> ST s (MutablePrimArray s a)
newPA Int
n =
#ifdef CHECKS
  check "newPA" (n >= 0) $
#endif
  Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
PA.newPrimArray Int
n
{-# INLINE newPA #-}

readPA :: Prim a => PA.MutablePrimArray s a -> Int -> ST s a
readPA :: forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s a
a Int
i = do
#ifdef CHECKS
  sz <- PA.getSizeofMutablePrimArray a
  check "readPA" (0 <= i && i < sz) $
#endif
    MutablePrimArray (PrimState (ST s)) a -> Int -> ST s a
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
PA.readPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
a Int
i
{-# INLINE readPA #-}

writePA :: Prim a => PA.MutablePrimArray s a -> Int -> a -> ST s ()
writePA :: forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s a
a Int
i a
x = do
#ifdef CHECKS
  sz <- PA.getSizeofMutablePrimArray a
  check "writePA" (0 <= i && i < sz) $
#endif
    MutablePrimArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
PA.writePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
a Int
i a
x
{-# INLINE writePA #-}

setPA :: Prim a => PA.MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA :: forall a s.
Prim a =>
MutablePrimArray s a -> Int -> Int -> a -> ST s ()
setPA MutablePrimArray s a
a Int
i Int
n a
x = do
#ifdef CHECKS
  sz <- PA.getSizeofMutablePrimArray a
  check "setPA" (0 <= i && 0 <= n && i + n <= sz) $
#endif
    MutablePrimArray (PrimState (ST s)) a -> Int -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
PA.setPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
a Int
i Int
n a
x
{-# INLINE setPA #-}

copyMutPA
  :: Prim a
  => PA.MutablePrimArray s a -> Int
  -> PA.MutablePrimArray s a -> Int
  -> Int
  -> ST s ()
copyMutPA :: forall a s.
Prim a =>
MutablePrimArray s a
-> Int -> MutablePrimArray s a -> Int -> Int -> ST s ()
copyMutPA MutablePrimArray s a
dst Int
dstoff MutablePrimArray s a
src Int
srcoff Int
n = do
#ifdef CHECKS
  dstSz <- PA.getSizeofMutablePrimArray dst
  srcSz <- PA.getSizeofMutablePrimArray src
  check "copyMutPA"
        ( 0 <= dstoff &&
          dstoff + n <= dstSz &&
          0 <= srcoff &&
          srcoff + n <= srcSz ) $
#endif
    MutablePrimArray (PrimState (ST s)) a
-> Int
-> MutablePrimArray (PrimState (ST s)) a
-> Int
-> Int
-> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
PA.copyMutablePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
dst Int
dstoff MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
src Int
srcoff Int
n
{-# INLINE copyMutPA #-}

modifyPA
  :: Prim a => PA.MutablePrimArray s a -> Int -> (a -> a) -> ST s ()
modifyPA :: forall a s.
Prim a =>
MutablePrimArray s a -> Int -> (a -> a) -> ST s ()
modifyPA MutablePrimArray s a
a Int
i a -> a
f = MutablePrimArray s a -> Int -> ST s a
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s a
a Int
i ST s a -> (a -> 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
>>= MutablePrimArray s a -> Int -> a -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s a
a Int
i (a -> ST s ()) -> (a -> a) -> a -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
{-# INLINE modifyPA #-}

modifyPAM
  :: Prim a => PA.MutablePrimArray s a -> Int -> (a -> ST s a) -> ST s ()
modifyPAM :: forall a s.
Prim a =>
MutablePrimArray s a -> Int -> (a -> ST s a) -> ST s ()
modifyPAM MutablePrimArray s a
a Int
i a -> ST s a
f = MutablePrimArray s a -> Int -> ST s a
forall a s. Prim a => MutablePrimArray s a -> Int -> ST s a
readPA MutablePrimArray s a
a Int
i ST s a -> (a -> ST s a) -> ST s a
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
>>= a -> ST s a
f ST s a -> (a -> 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
>>= MutablePrimArray s a -> Int -> a -> ST s ()
forall a s. Prim a => MutablePrimArray s a -> Int -> a -> ST s ()
writePA MutablePrimArray s a
a Int
i
{-# INLINE modifyPAM #-}

#ifdef CHECKS
check :: String -> Bool -> a -> a
check msg b x = if not b then error ("Data.Suffix." ++ msg) else x
#endif

----------
-- Utils
----------

data T2 a b = T2 !a !b
data T3 a b c = T3 !a !b !c
data E a b = L !a | R !b

foldlM :: (Foldable f, Monad m) => f a -> b -> (b -> a -> m b) -> m b
foldlM :: forall (f :: * -> *) (m :: * -> *) a b.
(Foldable f, Monad m) =>
f a -> b -> (b -> a -> m b) -> m b
foldlM f a
xs b
z0 b -> a -> m b
f = (a -> (b -> m b) -> b -> m b) -> (b -> m b) -> f a -> b -> m b
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (b -> m b) -> b -> m b
c b -> m b
forall {f :: * -> *} {a}. Applicative f => a -> f a
z f a
xs b
z0
  where
    z :: a -> f a
z !a
y = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
y
    c :: a -> (b -> m b) -> b -> m b
c a
x b -> m b
k !b
y = b -> a -> m b
f b
y a
x m b -> (b -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> m b
k
    {-# INLINE c #-}
{-# INLINE foldlM #-}

data Incr a where
  Incr :: !Int -> !Int -> Incr Int

instance Foldable Incr where
  foldr :: forall a b. (a -> b -> b) -> b -> Incr a -> b
foldr a -> b -> b
f b
z (Incr Int
i0 Int
j) = Int -> b
go Int
i0
    where
      go :: Int -> b
go Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
j = b
z
           | Bool
otherwise = a -> b -> b
f a
Int
i (Int -> b
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
  {-# INLINE foldr #-}

data Decr a where
  Decr :: !Int -> !Int -> Decr Int

instance Foldable Decr where
  foldr :: forall a b. (a -> b -> b) -> b -> Decr a -> b
foldr a -> b -> b
f b
z (Decr Int
i0 Int
j) = Int -> b
go Int
i0
    where
      go :: Int -> b
go Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
j = b
z
           | Bool
otherwise = a -> b -> b
f a
Int
i (Int -> b
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
  {-# INLINE foldr #-}