ac-library-hs-1.2.2.0: src/AtCoder/Extra/Seq/Raw.hs
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_HADDOCK hide #-}
-- | Base module for implementing dynamic sequences. It internaly uses a splay tree and user has to
-- track the root node change.
--
-- @since 1.2.0.0
module AtCoder.Extra.Seq.Raw
( -- * Seq
Seq (..),
-- * Constructors
newST,
resetST,
newNodeST,
newSeqST,
freeNodeST,
freeSubtreeST,
-- * Metadata
capacity,
lengthST,
-- * Merge/split
mergeST,
merge3ST,
merge4ST,
splitST,
split3ST,
split4ST,
splitLrST,
sliceST,
-- * Read/write
readST,
readMaybeST,
writeST,
modifyST,
exchangeST,
-- * Products
prodST,
prodMaybeST,
prodAllST,
-- * Applications
applyInST,
applyToRootST,
reverseST,
-- * Insert/delete
insertST,
deleteST,
deleteST_,
detachST,
-- * Balancing
rotateST,
splayST,
splayKthST,
-- * Bisection methods
-- ** C++-like
ilowerBoundST,
ilowerBoundM,
ilowerBoundProdST,
ilowerBoundProdM,
-- ** Splits
isplitMaxRightST,
isplitMaxRightM,
isplitMaxRightProdST,
isplitMaxRightProdM,
-- ** Max right
imaxRightST,
imaxRightM,
imaxRightProdST,
imaxRightProdM,
-- * Conversions
freezeST,
-- * Internals
-- | These functions are exported primarily for @Map@ implementations.
splitMaxRightWithST,
maxRightWithST,
updateNodeST,
writeNodeST,
modifyNodeST,
exchangeNodeST,
propNodeST,
applyNodeST,
)
where
import AtCoder.Extra.Pool qualified as P
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.LazySegTree (SegAct (..))
import Control.Monad (unless, when)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bit
import Data.Bits hiding (rotate)
import Data.Coerce (coerce)
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
import Prelude hiding (seq)
-- | Storages of dynamic sequences of monoid values with monoid actions on them through the `SegAct`
-- instance.
--
-- @since 1.2.0.0
data Seq s f a = Seq
{ -- | The maximum number of elements.
--
-- @since 1.2.0.0
nSeq :: {-# UNPACK #-} !Int,
-- | `Pool` for free slot management.
--
-- @since 1.2.0.0
poolSeq :: !(P.Pool s ()),
-- | Decomposed node data storage: left children.
--
-- @since 1.2.0.0
lSeq :: !(VUM.MVector s P.Index),
-- | Decomposed node data storage: right children.
--
-- @since 1.2.0.0
rSeq :: !(VUM.MVector s P.Index),
-- | Decomposed node data storage: parents.
--
-- @since 1.2.0.0
pSeq :: !(VUM.MVector s P.Index),
-- | Decomposed node data storage: subtree sizes.
--
-- @since 1.2.0.0
sSeq :: !(VUM.MVector s Int),
-- | Decomposed node data storage: monoid values.
--
-- @since 1.2.0.0
vSeq :: !(VUM.MVector s a),
-- | Decomposed node data storage: monoid products.
--
-- @since 1.2.0.0
prodSeq :: !(VUM.MVector s a),
-- | Decomposed node data storage: reversed flag of children.
--
-- @since 1.2.0.0
revSeq :: !(VUM.MVector s Bit),
-- | Decomposed node data storage: lazily propagated monoid action. Use @()@ if you don't need
-- monoid actions.
--
-- @since 1.2.0.0
lazySeq :: !(VUM.MVector s f)
}
-- | \(O(n)\) Creates a new `Seq` of length \(n\).
--
-- @since 1.2.0.0
{-# INLINEABLE newST #-}
newST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Int -> ST s (Seq s f a)
newST nSeq = do
poolSeq <- P.new nSeq
lSeq <- VUM.unsafeNew nSeq
rSeq <- VUM.unsafeNew nSeq
pSeq <- VUM.unsafeNew nSeq
sSeq <- VUM.unsafeNew nSeq
vSeq <- VUM.unsafeNew nSeq
prodSeq <- VUM.unsafeNew nSeq
revSeq <- VUM.unsafeNew nSeq
lazySeq <- VUM.unsafeNew nSeq
pure Seq {..}
-- | \(O(1)\) Clears the sequence storage.
--
-- @since 1.2.0.0
{-# INLINE resetST #-}
resetST :: Seq s f a -> ST s ()
resetST Seq {poolSeq} = stToPrim $ P.clear poolSeq
-- | \(O(1)\) Allocates a new sequence of length \(1\).
--
-- @since 1.2.0.0
{-# INLINEABLE newNodeST #-}
newNodeST :: (HasCallStack, Monoid f, VU.Unbox f, VU.Unbox a) => Seq s f a -> a -> ST s P.Index
newNodeST Seq {..} x = do
i <- P.alloc poolSeq ()
VGM.write lSeq (coerce i) P.undefIndex
VGM.write rSeq (coerce i) P.undefIndex
VGM.write pSeq (coerce i) P.undefIndex
VGM.write sSeq (coerce i) 1
VGM.write vSeq (coerce i) x
VGM.write prodSeq (coerce i) x
VGM.write revSeq (coerce i) $ Bit False
VGM.write lazySeq (coerce i) mempty
pure i
-- | \(O(n)\) Allocates a new sequence.
--
-- @since 1.2.0.0
{-# INLINEABLE newSeqST #-}
newSeqST :: (HasCallStack, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> VU.Vector a -> ST s P.Index
newSeqST seq@Seq {..} !xs = do
-- [l, r)
let inner l r
| l >= r = pure P.undefIndex
| l + 1 == r = newNodeST seq $ xs VG.! l
| otherwise = do
let !m = (l + r) `div` 2
rootL <- inner l m
root <- newNodeST seq (xs VG.! m)
rootR <- inner (m + 1) r
unless (P.nullIndex rootL) $ do
VGM.write lSeq (coerce root) rootL
VGM.write pSeq (coerce rootL) root
unless (P.nullIndex rootR) $ do
VGM.write rSeq (coerce root) rootR
VGM.write pSeq (coerce rootR) root
updateNodeST seq root
pure root
inner 0 (VU.length xs)
-- | \(O(1)\) Frees a node.
--
-- @since 1.2.0.0
{-# INLINE freeNodeST #-}
freeNodeST :: Seq s v a -> P.Index -> ST s ()
freeNodeST Seq {poolSeq} = P.free poolSeq
-- | \(O(n)\) Frees a subtree.
--
-- @since 1.2.0.0
{-# INLINEABLE freeSubtreeST #-}
freeSubtreeST :: (VU.Unbox a) => Seq s f a -> P.Index -> ST s ()
freeSubtreeST Seq {lSeq, rSeq, poolSeq} c0
| P.nullIndex c0 = pure ()
| otherwise = do
let inner c = do
cl <- VGM.read lSeq (coerce c)
unless (P.nullIndex cl) (inner cl)
cr <- VGM.read rSeq (coerce c)
unless (P.nullIndex cr) (inner cr)
inner c0
P.free poolSeq c0
-- -------------------------------------------------------------------------------------------------
-- Metadata
-- -------------------------------------------------------------------------------------------------
-- | \(O(1)\) Returns the capacity of the sequence storage.
--
-- @since 1.2.1.0
{-# INLINE capacity #-}
capacity :: Seq s f a -> Int
capacity = nSeq
-- | \(O(1)\) Returns the length of a sequence or a subtree.
--
-- @since 1.2.1.0
{-# INLINE lengthST #-}
lengthST :: Seq s f a -> P.Index -> ST s Int
lengthST Seq {..} i
| P.nullIndex i = pure 0
| otherwise = VGM.read sSeq (coerce i)
-- -------------------------------------------------------------------------------------------------
-- Merge/split
-- -------------------------------------------------------------------------------------------------
{-# INLINE assertRootST #-}
assertRootST :: (HasCallStack) => Seq s f a -> P.Index -> ST s ()
assertRootST Seq {pSeq} i = do
p <- VGM.read pSeq (coerce i)
let !_ = ACIA.runtimeAssert (P.nullIndex p) $ "AtCoder.Extra.Seq.Raw.assertRootST: not a root (node `" ++ show i ++ "`, parent `" ++ show p ++ "`)"
pure ()
{-# INLINE assertRootOrNullST #-}
assertRootOrNullST :: (HasCallStack) => Seq s f a -> P.Index -> ST s ()
assertRootOrNullST Seq {pSeq} i
| P.nullIndex i = pure ()
| otherwise = do
p <- VGM.read pSeq (coerce i)
let !_ = ACIA.runtimeAssert (P.nullIndex p) $ "AtCoder.Extra.Seq.Raw.assertRootOrNullST: not a root (node `" ++ show i ++ "`, parent `" ++ show p ++ "`)"
pure ()
-- | Amortized \(O(\log n)\). Merges two sequences \(l, r\) into one in the given order, ignoring
-- empty sequences.
--
-- ==== Constraints
-- - The vertices must be either null or a root.
--
-- @since 1.2.0.0
{-# INLINEABLE mergeST #-}
mergeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> ST s P.Index
mergeST seq@Seq {pSeq, lSeq} lRoot rRoot
| P.nullIndex lRoot = pure rRoot
| P.nullIndex rRoot = pure lRoot
| otherwise = do
assertRootST seq lRoot
assertRootST seq rRoot
rRoot' <- splayKthST seq rRoot 0
VGM.write lSeq (coerce rRoot') lRoot
VGM.write pSeq (coerce lRoot) rRoot'
updateNodeST seq rRoot'
pure rRoot'
-- | Amortized \(O(\log n)\). Merges three sequences \(l, m, r\) into one in the given order,
-- ignoring empty sequences.
--
-- ==== Constraints
-- - The vertices must be either null or a root.
--
-- @since 1.2.0.0
{-# INLINE merge3ST #-}
merge3ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> P.Index -> ST s P.Index
merge3ST seq a b c = do
r' <- mergeST seq a b
mergeST seq r' c
-- | Amortized \(O(\log n)\). Merges four sequences \(l, b, c, d, m, r\) into one in the given
-- order, ignoring empty sequences.
--
-- ==== Constraints
-- - The vertices must be either null or a root.
--
-- @since 1.2.0.0
{-# INLINE merge4ST #-}
merge4ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> P.Index -> P.Index -> P.Index -> ST s P.Index
merge4ST seq a b c d = do
r' <- mergeST seq a b
r'' <- mergeST seq r' c
mergeST seq r'' d
-- | Amortized \(O(\log n)\). Splits a sequences into two: \([0, k), [k, n)\).
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le k \le n\).
--
-- @since 1.2.0.0
{-# INLINEABLE splitST #-}
splitST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (P.Index, P.Index)
splitST seq@Seq {..} root k = do
assertRootOrNullST seq root
if k == 0
then pure (P.undefIndex, root)
else do
size <- VGM.read sSeq $ coerce root
if k == size
then pure (root, P.undefIndex)
else do
root' <- splayKthST seq root (k - 1)
r <- VGM.exchange rSeq (coerce root') P.undefIndex
VGM.write pSeq (coerce r) P.undefIndex
updateNodeST seq root'
pure (root', r)
-- | Amortized \(O(\log n)\). Splits a sequences into three: \([0, l), [l, r), [r, n)\).
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le l \le r \le n\).
--
-- @since 1.2.0.0
{-# INLINE split3ST #-}
split3ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (P.Index, P.Index, P.Index)
split3ST seq root l r = do
(!root', !c) <- splitST seq root r
(!a, !b) <- splitST seq root' l
pure (a, b, c)
-- | Amortized \(O(\log n)\). Splits a sequences into four: \([0, i), [i, j), [j, k), [k, n)\).
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le i \le j \le k \le n\).
--
-- @since 1.2.0.0
{-# INLINE split4ST #-}
split4ST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> Int -> ST s (P.Index, P.Index, P.Index, P.Index)
split4ST seq root i j k = do
(!root', !d) <- splitST seq root k
(!root'', !c) <- splitST seq root' j
(!a, !b) <- splitST seq root'' i
pure (a, b, c, d)
-- | Amortized \(O(\log n)\). Splits a sequence into three: \([0, \mathrm{root}), \mathrm{root}, [\mathrm{root} + 1, n)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINEABLE splitLrST #-}
splitLrST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s (P.Index, P.Index, P.Index)
splitLrST seq@Seq {..} root = do
assertRootST seq root
s <- do
rootL <- VGM.read lSeq (coerce root)
if P.nullIndex rootL
then pure 0
else VGM.read sSeq (coerce rootL)
split3ST seq root s (s + 1)
-- | Amortized \(O(\log n)\). Captures the root of a subtree of \([l, r)\). Splay the new root after
-- call.
--
-- ==== Constraints
-- - \(0 \le \lt r \le n\). Note that the interval must have positive length.
--
-- @since 1.2.0.0
{-# INLINEABLE sliceST #-}
sliceST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s P.Index
sliceST seq@Seq {..} root l r
| l == 0 = do
size <- VGM.read sSeq (coerce root)
if r == size
then pure root
else do
root' <- splayKthST seq root r
VGM.read lSeq $ coerce root'
| otherwise = do
size <- VGM.read sSeq $ coerce root
if r == size
then do
root' <- splayKthST seq root (l - 1)
VGM.read rSeq $ coerce root'
else do
-- o--l--o--o--r--o
-- [ )
-- * root' (splayed)
-- * rootL (detached from the root)
-- \* rootL' (splayed)
-- * right(rootL'): node that corresponds to [l, r)
root' <- splayKthST seq root r
rootL <- VGM.read lSeq $ coerce root'
-- detach `rootL` from `root'`
VGM.write pSeq (coerce rootL) P.undefIndex
rootL' <- splayKthST seq rootL (l - 1)
-- re-attach `rootL'` to `root'`
VGM.write pSeq (coerce rootL') root'
VGM.write lSeq (coerce root') rootL'
updateNodeST seq root'
VGM.read rSeq $ coerce rootL'
-- -------------------------------------------------------------------------------------------------
-- Modifications
-- -------------------------------------------------------------------------------------------------
-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value.
--
-- ==== Constraints
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE readST #-}
readST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (a, P.Index)
readST seq@Seq {..} root k = do
assertRootST seq root
root' <- splayKthST seq root k
(,root') <$> VGM.read vSeq (coerce root')
-- | Amortized \(O(\log n)\). Reads the \(k\)-th node's monoid value.
--
-- ==== Constraints
-- - The root must be empty or a root.
--
-- @since 1.2.0.0
{-# INLINEABLE readMaybeST #-}
readMaybeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (Maybe (a, P.Index))
readMaybeST seq@Seq {..} root k
| P.nullIndex root = pure Nothing
| otherwise = do
assertRootST seq root
s <- VGM.read sSeq (coerce root)
if 0 <= k && k < s
then do
root' <- splayKthST seq root k
Just . (,root') <$> VGM.read vSeq (coerce root')
else pure Nothing
-- | Amortized \(O(\log n)\). Writes to the \(k\)-th node's monoid value.
--
-- ==== Constraints
-- - The node must be a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE writeST #-}
writeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s P.Index
writeST seq root k v = do
assertRootST seq root
root' <- splayKthST seq root k
writeNodeST seq root' v
pure root'
-- | Amortized \(O(\log n)\). Modifies the \(k\)-th node's monoid value.
--
-- ==== Constraints
-- - The node must be a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE modifyST #-}
modifyST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> (a -> a) -> Int -> ST s P.Index
modifyST seq root f k = do
assertRootST seq root
root' <- splayKthST seq root k
modifyNodeST seq f root'
pure root'
-- | Amortized \(O(\log n)\). Exchanges the \(k\)-th node's monoid value.
--
-- ==== Constraints
-- - The node must be a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE exchangeST #-}
exchangeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s (a, P.Index)
exchangeST seq root k v = do
assertRootST seq root
root' <- splayKthST seq root k
res <- exchangeNodeST seq root' v
pure (res, root')
-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\).
--
-- ==== Constraints
-- - The node must be a root
-- - \(0 \le l \le r \le n\)
--
-- @since 1.2.0.0
{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (a, P.Index)
prodST seq@Seq {sSeq} root l r = do
s <- if P.nullIndex root then pure 0 else VGM.read sSeq (coerce root)
let !_ = ACIA.checkInterval "AtCoder.Extra.Seq.Raw.prodST" l r s
if l == r
then pure (mempty, root)
else unsafeProdST seq root l r
-- | Amortized \(O(\log n)\). Returns the monoid product in an interval \([l, r)\). Returns
-- `Nothing` if an invalid interval is given or for an empty sequence.
--
-- @since 1.2.0.0
{-# INLINEABLE prodMaybeST #-}
prodMaybeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (Maybe (a, P.Index))
prodMaybeST seq@Seq {sSeq} root l r
| P.nullIndex root = pure Nothing
| otherwise = do
s <- VGM.read sSeq (coerce root)
if not (ACIA.testInterval l r s)
then pure Nothing
else
if l == r
then pure $ Just (mempty, root)
else Just <$> unsafeProdST seq root l r
-- | Amortized \(O(\log n)\).
--
-- ==== Constraint
-- - \(0 \le \lt r \le n\). Note that the interval must have positive length.
{-# INLINEABLE unsafeProdST #-}
unsafeProdST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s (a, P.Index)
unsafeProdST seq@Seq {..} root l r = do
assertRootST seq root
target <- sliceST seq root l r
res <- VGM.read prodSeq $ coerce target
splayST seq target True
pure (res, target)
-- | Amortized \(O(\log n)\). Returns the monoid product of the whole sequence. Returns `mempty`
-- for an empty sequence.
--
-- ==== Constraint
-- - The node must be null or a root.
--
-- @since 1.2.0.0
{-# INLINEABLE prodAllST #-}
prodAllST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s a
prodAllST seq@Seq {..} root = do
if P.nullIndex root
then pure mempty
else do
assertRootST seq root
VGM.read prodSeq $ coerce root
-- | Amortized \(O(\log n)\). Given an interval \([l, r)\), applies a monoid action \(f\).
--
-- ==== Constraints
-- - \(0 \le l \le r \le n\)
--
-- @since 1.2.0.0
{-# INLINEABLE applyInST #-}
applyInST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> f -> ST s P.Index
applyInST seq@Seq {..} root l r act = do
-- assertRootOrNullST seq root
s <- if P.nullIndex root then pure 0 else VGM.read sSeq (coerce root)
let !_ = ACIA.checkInterval "AtCoder.Extra.Seq.applyInST" l r s
if l == r
then pure root
else do
root' <- sliceST seq root l r
applyNodeST seq root' act
splayST seq root' True
pure root'
-- | \(O(1)\) Applies a monoid action \(f\) to the root of a sequence.
--
-- @since 1.2.0.0
{-# INLINEABLE applyToRootST #-}
applyToRootST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> f -> ST s ()
applyToRootST seq@Seq {..} root act
| P.nullIndex root = pure ()
| otherwise = do
rootP <- VGM.read pSeq (coerce root)
when (P.nullIndex rootP) $ do
applyNodeST seq root act
-- | Amortized \(O(\log n)\). Reverses the sequence in \([l, r)\).
--
-- ==== Constraints
-- - The monoid action \(f\) must be commutative.
-- - The monoid value \(v\) must be commutative.
--
-- @since 1.2.0.0
{-# INLINEABLE reverseST #-}
reverseST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> Int -> ST s P.Index
reverseST seq@Seq {sSeq} root0 l r
| P.nullIndex root0 = pure P.undefIndex
| otherwise = do
s <- VGM.read sSeq (coerce root0)
if not (ACIA.testInterval l r s)
then pure root0
else
if l == r
then pure root0
else do
root' <- sliceST seq root0 l r
reverseNodeST seq root'
splayST seq root' True
pure root'
-- | Amortized \(O(\log n)\). Inserts a new node at \(k\) with initial monoid value \(v\). This
-- functions for an empty index.
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le k \le n\)
--
-- @since 1.2.0.0
{-# INLINEABLE insertST #-}
insertST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> a -> ST s P.Index
insertST seq root k v = do
if P.nullIndex root
then do
-- `insertST` is actually `insertOrNewNodeST`: it's specifically designed to work for an empty
-- sequence.
newNodeST seq v
else do
(!l, !r) <- splitST seq root k
node <- newNodeST seq v
merge3ST seq l node r
-- | Amortized \(O(\log n)\). Frees the \(k\)-th node and returns the monoid value of it.
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE deleteST #-}
deleteST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s (a, P.Index)
deleteST seq@Seq {..} root i = do
(!l, !m, !r) <- split3ST seq root i (i + 1)
x <- VGM.read vSeq (coerce m)
freeNodeST seq m
root' <- mergeST seq l r
pure (x, root')
-- | Amortized \(O(\log n)\). Frees the \(k\)-th node.
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE deleteST_ #-}
deleteST_ :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index
deleteST_ seq root i = do
(!l, !m, !r) <- split3ST seq root i (i + 1)
freeNodeST seq m
root' <- mergeST seq l r
pure root'
-- | Amortized \(O(\log n)\). Detaches the \(k\)-th node and returns the new root of the original
-- sequence.
--
-- ==== Constraints
-- - The node must be null or a root.
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE detachST #-}
detachST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index
detachST seq root i = do
(!l, !m, !r) <- split3ST seq root i (i + 1)
freeNodeST seq m
root' <- mergeST seq l r
pure root'
-- -------------------------------------------------------------------------------------------------
-- Balancing
-- -------------------------------------------------------------------------------------------------
-- | Amortized \(O(\log n)\). Rotates a child node.
--
-- ==== Constraints
-- - \(0 \le i \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE rotateST #-}
rotateST :: (HasCallStack) => Seq s v a -> P.Index -> ST s ()
rotateST Seq {..} !i = do
p <- VGM.read pSeq $ coerce i
pl <- VGM.read lSeq $ coerce p
c <-
if pl == i
then do
-- p i
-- / \
-- i -> p
-- \ /
-- r r
r <- VGM.exchange rSeq (coerce i) p
VGM.write lSeq (coerce p) r
pure r
else do
-- p i
-- \ /
-- i -> p
-- / \
-- l l
l <- VGM.exchange lSeq (coerce i) p
VGM.write rSeq (coerce p) l
pure l
pp <- VGM.read pSeq $ coerce p
unless (P.nullIndex pp) $ do
-- pp pp
-- / -> /
-- p i
VGM.modify lSeq (\ppl -> if ppl == p then i else ppl) $ coerce pp
-- pp pp
-- \ -> \
-- p i
VGM.modify rSeq (\ppr -> if ppr == p then i else ppr) $ coerce pp
-- set parents
VGM.write pSeq (coerce i) pp
VGM.write pSeq (coerce p) i
unless (P.nullIndex c) $ do
VGM.write pSeq (coerce c) p
-- | Amortized \(O(\log n)\). Moves up a node to be a root.
--
-- @since 1.2.0.0
{-# INLINEABLE splayST #-}
splayST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Bool -> ST s ()
splayST seq@Seq {..} i doneParentProp = do
-- we have no chance to call `splayST` with `doneParentProp == False` though
if doneParentProp
then propNodeST seq i
else propNodeFromRootST seq i
let inner = do
p <- VGM.read pSeq $ coerce i
unless (P.nullIndex p) $ do
pp <- VGM.read pSeq $ coerce p
if P.nullIndex pp
then do
rotateST seq i
updateNodeST seq p
pure ()
else do
pl <- VGM.read lSeq $ coerce p
pr <- VGM.read rSeq $ coerce p
ppl <- VGM.read lSeq $ coerce pp
ppr <- VGM.read rSeq $ coerce pp
if pl == i && ppl == p || pr == i && ppr == p
then do
-- same direction twice
rotateST seq p
rotateST seq i
else do
rotateST seq i
rotateST seq i
updateNodeST seq pp
updateNodeST seq p
inner
inner
updateNodeST seq i
-- | Amortized \(O(\log n)\). Finds \(k\)-th node and splays it. Returns the new root.
--
-- ==== Constraints
-- - \(0 \le k \lt n\)
--
-- @since 1.2.0.0
{-# INLINEABLE splayKthST #-}
splayKthST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> Int -> ST s P.Index
splayKthST seq@Seq {..} root0 k0 = do
size <- VGM.read sSeq $ coerce root0
let !_ = ACIA.checkIndex "AtCoder.Extra.Seq.Raw.splayKthST" k0 size
let inner root k = do
propNodeST seq root
l <- VGM.read lSeq $ coerce root
-- The number of left children = the node's index counting from the leftmost.
sizeL <- if P.nullIndex l then pure 0 else VGM.read sSeq $ coerce l
case compare k sizeL of
EQ -> pure root
LT -> inner l k
GT -> do
r <- VGM.read rSeq $ coerce root
inner r (k - (sizeL + 1))
target <- inner root0 k0
splayST seq target True
pure target
-- -------------------------------------------------------------------------------------------------
-- Bisection methods
-- -------------------------------------------------------------------------------------------------
-- | Amortized \(O(\log n)\).
--
-- @since 1.2.0.0
{-# INLINE ilowerBoundST #-}
ilowerBoundST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> Bool) ->
-- | (r, root)
ST s (Int, P.Index)
ilowerBoundST seq root f = stToPrim $ do
if P.nullIndex root
then pure (0, P.undefIndex)
else do
(!r, !_, !root') <- imaxRightST seq root f
pure (r, root')
-- | Amortized \(O(\log n)\).
--
-- @since 1.2.0.0
{-# INLINE ilowerBoundM #-}
ilowerBoundM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> m Bool) ->
-- | (r, root)
m (Int, P.Index)
ilowerBoundM seq root f = do
if P.nullIndex root
then pure (0, P.undefIndex)
else do
(!r, !_, !root') <- imaxRightM seq root f
pure (r, root')
-- | Amortized \(O(\log n)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINE ilowerBoundProdST #-}
ilowerBoundProdST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product
(Int -> a -> Bool) ->
-- | (r, root)
ST s (Int, P.Index)
ilowerBoundProdST seq root f = do
if P.nullIndex root
then pure (0, P.undefIndex)
else do
(!r, !_, !root') <- imaxRightProdST seq root f
pure (r, root')
-- | Amortized \(O(\log n)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINE ilowerBoundProdM #-}
ilowerBoundProdM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid product
(Int -> a -> m Bool) ->
-- | (r, root)
m (Int, P.Index)
ilowerBoundProdM seq root f = do
if P.nullIndex root
then pure (0, P.undefIndex)
else do
(!r, !_, !root') <- imaxRightProdM seq root f
pure (r, root')
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- @since 1.2.0.0
{-# INLINE isplitMaxRightST #-}
isplitMaxRightST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> Bool) ->
-- | (left, right) sequences where \(f\) holds for the left
ST s (P.Index, P.Index)
isplitMaxRightST seq root f = stToPrim $ isplitMaxRightM seq root (\i x -> pure (f i x))
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- @since 1.2.0.0
{-# INLINEABLE isplitMaxRightM #-}
isplitMaxRightM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> m Bool) ->
-- | (left, right) sequences where \(f\) holds for the left
m (P.Index, P.Index)
isplitMaxRightM seq@Seq {..} root f
| P.nullIndex root = pure (P.undefIndex, P.undefIndex)
| otherwise = do
stToPrim $ assertRootST seq root
(!_, !c, !_) <- imaxRightM seq root f
if P.nullIndex c
then stToPrim $ do
-- `f` does hot hold
splayST seq root True
pure (P.undefIndex, root)
else stToPrim $ do
splayST seq c True
right <- VGM.read rSeq (coerce c)
if P.nullIndex right
then do
-- `f` holds for the whole sequence
pure (c, P.undefIndex)
else do
-- `f` holds for part of the sequence. detach the right child
VGM.write pSeq (coerce right) P.undefIndex
VGM.write rSeq (coerce c) P.undefIndex
updateNodeST seq c
pure (c, right)
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- @since 1.2.0.0
{-# INLINE isplitMaxRightProdST #-}
isplitMaxRightProdST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value
(Int -> a -> Bool) ->
-- | (left, right) sequences where \(f\) holds for the left
ST s (P.Index, P.Index)
isplitMaxRightProdST seq root f = stToPrim $ isplitMaxRightProdM seq root (\i x -> pure (f i x))
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- @since 1.2.0.0
{-# INLINEABLE isplitMaxRightProdM #-}
isplitMaxRightProdM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
-- | \(r\)
(Int -> a -> m Bool) ->
-- | (left, right) sequences where \(f\) holds for the left
m (P.Index, P.Index)
isplitMaxRightProdM seq@Seq {..} root f
| P.nullIndex root = pure (P.undefIndex, P.undefIndex)
| otherwise = do
stToPrim $ assertRootST seq root
(!_, !c, !_) <- imaxRightProdM seq root f
if P.nullIndex c
then stToPrim $ do
-- `f` does hot hold
splayST seq root True
pure (P.undefIndex, root)
else stToPrim $ do
splayST seq c True
right <- VGM.read rSeq (coerce c)
if P.nullIndex right
then do
-- `f` holds for the whole sequence
pure (c, P.undefIndex)
else do
-- `f` holds for part of the sequence. detach the right child
VGM.write pSeq (coerce right) P.undefIndex
VGM.write rSeq (coerce c) P.undefIndex
updateNodeST seq c
pure (c, right)
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v\)
-- where \(f(v)\) holds for every \(v_i (0 \le i \lt k)\). Note that \(f\) works for a single
-- node, not a monoid product.
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINE imaxRightST #-}
imaxRightST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> Bool) ->
-- | (r, left, right)
ST s (Int, P.Index, P.Index)
imaxRightST seq root0 f = stToPrim $ imaxRightM seq root0 (\i x -> pure (f i x))
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \(v_i (0 \le i \le k)\). Note that \(f\) works for a single
-- node, not a monoid product.
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINEABLE imaxRightM #-}
imaxRightM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_i)\) that takes the index and the monoid value
(Int -> a -> m Bool) ->
-- | (r, left, right)
m (Int, P.Index, P.Index)
imaxRightM seq@Seq {..} root0 f = do
let inner offset parent root lastYes
| P.nullIndex root = pure (offset, lastYes, parent)
| otherwise = do
stToPrim $ propNodeST seq root
l <- stToPrim $ VGM.read lSeq (coerce root)
v <- stToPrim $ VGM.read vSeq (coerce root)
pos <- stToPrim $ do
if P.nullIndex l
then pure offset
else (offset +) <$> VGM.read sSeq (coerce l)
!b <- f pos v
if b
then do
r <- stToPrim $ VGM.read rSeq $ coerce root
inner (pos + 1) root r root
else do
inner offset root l lastYes
(!r, !yes, !root') <- inner 0 P.undefIndex root0 P.undefIndex
stToPrim $ splayST seq root' True
pure (r, yes, root')
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINE imaxRightProdST #-}
imaxRightProdST ::
(SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value
(Int -> a -> Bool) ->
-- | (ilowerBound, rightmost node, new root)
ST s (Int, P.Index, P.Index)
imaxRightProdST seq root0 f = imaxRightProdM seq root0 (\i x -> pure (f i x))
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.0.0
{-# INLINEABLE imaxRightProdM #-}
imaxRightProdM ::
(PrimMonad m, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq (PrimState m) f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i, v_0 \dots v_i)\) that takes the index and the monoid value
(Int -> a -> m Bool) ->
-- | (ilowerBound, rightmost node, new root)
m (Int, P.Index, P.Index)
imaxRightProdM seq@Seq {..} root0 f = do
let inner !acc offset parent root lastYes
| P.nullIndex root = pure (offset, lastYes, parent)
| otherwise = do
stToPrim $ propNodeST seq root
l <- stToPrim $ VGM.read lSeq $ coerce root
pos <- stToPrim $ do
if P.nullIndex l
then pure offset
else (offset +) <$> VGM.read sSeq (coerce l)
-- [0, pos]
prodM <- stToPrim $ do
-- detach right child (temporarily) and read the product
rootR <- VGM.exchange rSeq (coerce root) P.undefIndex
updateNodeST seq root
prodRoot <- VGM.read prodSeq (coerce root)
-- attach the right child again
VGM.write rSeq (coerce root) rootR
updateNodeST seq root
pure $! acc <> prodRoot
!b <- f pos prodM
if b
then do
r <- stToPrim $ VGM.read rSeq $ coerce root
inner prodM (pos + 1) root r root
else do
inner acc offset root l lastYes
(!r, !yes, !root') <- inner mempty 0 P.undefIndex root0 P.undefIndex
stToPrim $ splayST seq root' True
pure (r, yes, root')
-- -------------------------------------------------------------------------------------------------
-- Conversions
-- -------------------------------------------------------------------------------------------------
-- | Amortized \(O(n)\). Returns the sequence of monoid values.
--
-- @since 1.2.0.0
{-# INLINEABLE freezeST #-}
freezeST :: (HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s (VU.Vector a)
freezeST seq@Seq {sSeq, lSeq, rSeq, vSeq} root0
| P.nullIndex root0 = pure VU.empty
| otherwise = do
size <- VGM.read sSeq (coerce root0)
res <- VUM.unsafeNew size
let inner i root
| P.nullIndex root = pure i
| otherwise = do
-- visit from left to right
propNodeST seq root
i' <- inner i =<< VGM.read lSeq (coerce root)
vx <- VGM.read vSeq (coerce root)
VGM.write res i' vx
inner (i' + 1) =<< VGM.read rSeq (coerce root)
_ <- inner 0 root0
VU.unsafeFreeze res
-- -------------------------------------------------------------------------------------------------
-- Internals and node methods
-- -------------------------------------------------------------------------------------------------
-- | Amortized \(O(\log n)\).
--
-- @since 1.2.1.0
{-# INLINEABLE splitMaxRightWithST #-}
splitMaxRightWithST ::
(HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate \(f(i)\)
(P.Index -> ST s Bool) ->
-- | (left, right) sequences where \(f\) holds for the left
ST s (P.Index, P.Index)
splitMaxRightWithST seq@Seq {..} root f
| P.nullIndex root = pure (P.undefIndex, P.undefIndex)
| otherwise = do
assertRootST seq root
(!c, !_) <- maxRightWithST seq root f
if P.nullIndex c
then do
-- `f` does hot hold
splayST seq root True
pure (P.undefIndex, root)
else do
splayST seq c True
right <- VGM.read rSeq (coerce c)
if P.nullIndex right
then do
-- `f` holds for the whole sequence
pure (c, P.undefIndex)
else do
-- `f` holds for part of the sequence. detach the right child
VGM.write pSeq (coerce right) P.undefIndex
VGM.write rSeq (coerce c) P.undefIndex
updateNodeST seq c
pure (c, right)
-- | Amortized \(O(\log n)\). Given a monotonious sequence, returns the rightmost node \(v_k\)
-- where \(f(v)\) holds for every \([0, i) (0 \le i \lt k)\).
--
-- ==== Constraints
-- - The node must be a root.
--
-- @since 1.2.1.0
{-# INLINEABLE maxRightWithST #-}
maxRightWithST ::
(HasCallStack, SegAct f a, Eq f, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) =>
-- | Sequence storage
Seq s f a ->
-- | Root node
P.Index ->
-- | User predicate
(P.Index -> ST s Bool) ->
-- | (rightmost node, new root)
ST s (P.Index, P.Index)
maxRightWithST seq@Seq {..} root0 f = do
let inner parent lastYes root
| P.nullIndex root = pure (lastYes, parent)
| otherwise = do
stToPrim $ propNodeST seq root
!b <- f root
if b
then inner root root =<< VGM.read rSeq (coerce root)
else inner root lastYes =<< VGM.read lSeq (coerce root)
(!yes, !root') <- inner P.undefIndex P.undefIndex root0
stToPrim $ splayST seq root' True
pure (yes, root')
-- NOTE(pref): inlining node functions are important for the speed
-- | \(O(1)\) Recomputes the node size and the monoid product.
--
-- @since 1.2.1.0
{-# INLINEABLE updateNodeST #-}
updateNodeST :: (Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s ()
updateNodeST Seq {..} i = do
l <- VGM.read lSeq (coerce i)
r <- VGM.read rSeq (coerce i)
prodM <- VGM.read vSeq (coerce i)
(!size', !prod') <-
if P.nullIndex l
then pure (1, prodM)
else do
sizeL <- VGM.read sSeq (coerce l)
prodL <- VGM.read prodSeq (coerce l)
pure (sizeL + 1, prodL <> prodM)
(!size'', !prod'') <-
if P.nullIndex r
then pure (size', prod')
else do
sizeR <- VGM.read sSeq (coerce r)
prodR <- VGM.read prodSeq (coerce r)
pure (size' + sizeR, prod' <> prodR)
VGM.write sSeq (coerce i) size''
VGM.write prodSeq (coerce i) prod''
-- | \(O(1)\) Writes to the monoid value of a node.
--
-- @since 1.2.1.0
{-# INLINE writeNodeST #-}
writeNodeST :: (Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> a -> ST s ()
writeNodeST seq@Seq {..} root v = do
assertRootST seq root
VGM.write vSeq (coerce root) v
updateNodeST seq root
-- | \(O(1)\) Modifies the monoid value of a node.
--
-- @since 1.2.1.0
{-# INLINE modifyNodeST #-}
modifyNodeST :: (HasCallStack, Monoid a, VU.Unbox a) => Seq s f a -> (a -> a) -> P.Index -> ST s ()
modifyNodeST seq@Seq {..} f root = do
assertRootST seq root
VGM.modify vSeq f $ coerce root
updateNodeST seq root
-- | \(O(1)\) Exchanges the monoid value of a node.
--
-- @since 1.2.1.0
{-# INLINE exchangeNodeST #-}
exchangeNodeST :: (HasCallStack, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> a -> ST s a
exchangeNodeST seq@Seq {..} root v = do
assertRootST seq root
res <- VGM.exchange vSeq (coerce root) v
updateNodeST seq root
pure res
-- | \(O(1)\) Swaps the left and the right children.
{-# INLINE swapLrNodeST #-}
swapLrNodeST :: Seq s f a -> P.Index -> ST s ()
swapLrNodeST Seq {..} i = do
VGM.modifyM lSeq (VGM.exchange rSeq (coerce i)) (coerce i)
-- | \(O(1)\) Reverses the left and the right children, lazily and recursively.
{-# INLINE reverseNodeST #-}
reverseNodeST :: Seq s f a -> P.Index -> ST s ()
reverseNodeST seq@Seq {..} i = do
swapLrNodeST seq i
-- lazily propagate new reverse or cancel:
VGM.modify revSeq (xor (Bit True)) $ coerce i
-- | Amortized \(O(\log n)\). Propgates the lazily propagated values on a node.
--
-- @since 1.2.1.0
{-# INLINE propNodeST #-}
-- NOTE(pref): Although this function is large, inlining it needs for the speed.
propNodeST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> ST s ()
propNodeST seq@Seq {..} i = do
-- action
act <- VGM.exchange lazySeq (coerce i) mempty
-- this is where `Eq f` is required:
when (act /= mempty) $ do
l <- VGM.read lSeq $ coerce i
unless (P.nullIndex l) $ do
applyNodeST seq l act
r <- VGM.read rSeq $ coerce i
unless (P.nullIndex r) $ do
applyNodeST seq r act
-- reverse
Bit b <- VGM.exchange revSeq (coerce i) (Bit False)
when b $ do
l <- VGM.read lSeq $ coerce i
unless (P.nullIndex l) $ do
-- propagate new reverse or cancel:
reverseNodeST seq l
r <- VGM.read rSeq $ coerce i
unless (P.nullIndex r) $ do
-- propagate new reverse or cancel:
reverseNodeST seq r
-- | Amortized \(O(\log n)\). Propagetes from the root to the given node.
{-# INLINE propNodeFromRootST #-}
propNodeFromRootST :: (HasCallStack, SegAct f a, VU.Unbox f, VU.Unbox a, Monoid a) => Seq s f a -> P.Index -> ST s ()
propNodeFromRootST Seq {..} = inner
where
inner i = do
p <- VGM.read pSeq $ coerce i
unless (P.nullIndex p) $ do
inner p
inner i
-- | Amortized \(O(\log n)\). Propgates at a node.
--
-- @since 1.2.1.0
{-# INLINE applyNodeST #-}
applyNodeST :: (HasCallStack, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => Seq s f a -> P.Index -> f -> ST s ()
applyNodeST Seq {..} i act = do
len <- VGM.read sSeq $ coerce i
VGM.modify vSeq (segAct act) $ coerce i
VGM.modify prodSeq (segActWithLength len act) $ coerce i
VGM.modify lazySeq (act <>) $ coerce i