{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_HADDOCK hide #-}

-- | Base module of a dynamic segment tree.
--
-- @since 1.2.1.0
module AtCoder.Extra.DynSegTree.Raw
  ( -- * Dynamic segment tree
    DynSegTree (..),

    -- * Re-exports
    P.Index (..),

    -- * Constructors
    newST,
    newRootST,
    newNodeST,
    newSeqST,

    -- * Accessing elements
    modifyMST,

    -- * Products
    prodST,
    -- prodMaybe,

    -- * Tree operations
    resetIntervalST,

    -- * Binary searches
    maxRightM,
    -- -- * Conversions
    -- freezeST,
  )
where

import AtCoder.Extra.Pool qualified as P
import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
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 (read)

-- | A dynamic segment tree that covers a half-open interval \([l_0, r_0)\). Is is dynamic in that
-- the nodes are instantinated as needed.
--
-- @since 1.2.1.0
data DynSegTree s a = DynSegTree
  { -- | The maximum number of nodes allocated
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Int
capacityDst :: {-# UNPACK #-} !Int,
    -- | Whether the data is persistent or not
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Bool
isPersistentDst :: {-# UNPACK #-} !Bool,
    -- | Left index boundary (inclusive)
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Int
l0Dst :: {-# UNPACK #-} !Int,
    -- | Right index boundary (exclusive)
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Int
r0Dst :: {-# UNPACK #-} !Int,
    -- | Initial monoid value assignment \(g: (l, r) \rightarrow a\)
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Int -> Int -> a
initialProdDst :: !(Int -> Int -> a),
    -- | `Pool` for free slot management.
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> Pool s ()
poolDst :: !(P.Pool s ()),
    -- | Decomposed node storage: left children
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> MVector s Index
lDst :: !(VUM.MVector s P.Index),
    -- | Decomposed node storage: right children
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> MVector s Index
rDst :: !(VUM.MVector s P.Index),
    -- | Decomposed node storage: monoid value
    --
    -- @since 1.2.1.0
    forall s a. DynSegTree s a -> MVector s a
xDst :: !(VUM.MVector s a)
  }

-- | \(O(n)\)
--
-- @since 1.2.1.0
{-# INLINEABLE newST #-}
newST :: (HasCallStack, VU.Unbox a) => Bool -> Int -> Int -> Int -> (Int -> Int -> a) -> ST s (DynSegTree s a)
newST :: forall a s.
(HasCallStack, Unbox a) =>
Bool
-> Int -> Int -> Int -> (Int -> Int -> a) -> ST s (DynSegTree s a)
newST Bool
isPersistentDst Int
capacityDst Int
l0Dst Int
r0Dst Int -> Int -> a
initialProdDst = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
l0Dst Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0Dst) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.DynSegTree.Raw.newST: given invalid interval " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
l0Dst, Int
r0Dst)
  Pool s ()
poolDst <- Int -> ST s (Pool (PrimState (ST s)) ())
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int -> m (Pool (PrimState m) a)
P.new Int
capacityDst
  MVector s Index
lDst <- Int -> ST s (MVector (PrimState (ST s)) Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
capacityDst
  MVector s Index
rDst <- Int -> ST s (MVector (PrimState (ST s)) Index)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
capacityDst
  MVector s a
xDst <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
capacityDst
  DynSegTree s a -> ST s (DynSegTree s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
isPersistentDst :: Bool
capacityDst :: Int
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..}

-- | \(O(1)\)
--
-- @since 1.2.1.0
{-# INLINE newRootST #-}
newRootST :: (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> ST s P.Index
newRootST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> ST s Index
newRootST dst :: DynSegTree s a
dst@DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} = do
  DynSegTree s a -> Int -> Int -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Int -> Int -> ST s Index
newNodeInST DynSegTree s a
dst Int
l0Dst Int
r0Dst

-- | \(O(1)\)
--
-- @since 1.2.1.0
{-# INLINE newNodeST #-}
newNodeST :: (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> a -> ST s P.Index
newNodeST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> a -> ST s Index
newNodeST DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} !a
x = do
  Index
i <- Pool (PrimState (ST s)) () -> () -> ST s Index
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> a -> m Index
P.alloc Pool s ()
Pool (PrimState (ST s)) ()
poolDst ()
  MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) Index
P.undefIndex
  MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) Index
P.undefIndex
  MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) a
x
  Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
i

-- | \(O(L)\)
--
-- @since 1.2.1.0
{-# INLINEABLE newSeqST #-}
newSeqST :: (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> VU.Vector a -> ST s P.Index
newSeqST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Vector a -> ST s Index
newSeqST dst :: DynSegTree s a
dst@DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} !Vector a
xs = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
l0Dst Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
r0Dst Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs) String
"AtCoder.Extra.DynSegTree.Raw: mismatch between the bounds and the input vector: the bounds must be [0, n)"
  -- run DFS and allocate nodes from left to right
  let dfs :: Int -> Int -> ST s Index
dfs Int
l Int
r
        | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
P.undefIndex
        | Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = DynSegTree s a -> a -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> a -> ST s Index
newNodeST DynSegTree s a
dst (a -> ST s Index) -> a -> ST s Index
forall a b. (a -> b) -> a -> b
$ Vector a
xs Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
l
        | Bool
otherwise = do
            let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            Index
lRoot <- Int -> Int -> ST s Index
dfs Int
l Int
m
            Index
rRoot <- Int -> Int -> ST s Index
dfs Int
m Int
r
            a
xlRoot <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
lRoot)
            a
xrRoot <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
rRoot)
            let !x :: a
x = a
xlRoot a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xrRoot
            Index
root <- DynSegTree s a -> a -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> a -> ST s Index
newNodeST DynSegTree s a
dst a
x
            MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
root) Index
lRoot
            MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
root) Index
rRoot
            Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
root
  Int -> Int -> ST s Index
dfs Int
0 (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs)

-- | \(O(1)\)
--
-- ==== Constraints
-- - The interval must be non-empty
--
-- @since 1.2.1.0
{-# INLINE newNodeInST #-}
newNodeInST :: (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> Int -> Int -> ST s P.Index
newNodeInST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Int -> Int -> ST s Index
newNodeInST dst :: DynSegTree s a
dst@DynSegTree {Int -> Int -> a
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
initialProdDst :: Int -> Int -> a
initialProdDst} Int
l Int
r = do
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
l) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.DynSegTree.Raw.nodeNodeInST: not empty or negative interval: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
l, Int
r)
  DynSegTree s a -> a -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> a -> ST s Index
newNodeST DynSegTree s a
dst (a -> ST s Index) -> a -> ST s Index
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
l Int
r

-- | \(O(\log L)\)
--
-- @since 1.2.1.0
{-# INLINEABLE modifyMST #-}
modifyMST :: forall m a. (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DynSegTree (PrimState m) a -> P.Index -> (a -> m a) -> Int -> m P.Index
modifyMST :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DynSegTree (PrimState m) a -> Index -> (a -> m a) -> Int -> m Index
modifyMST dst :: DynSegTree (PrimState m) a
dst@DynSegTree {Bool
Int
MVector (PrimState m) a
MVector (PrimState m) Index
Pool (PrimState m) ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool (PrimState m) ()
lDst :: MVector (PrimState m) Index
rDst :: MVector (PrimState m) Index
xDst :: MVector (PrimState m) a
..} Index
root a -> m a
f Int
i = do
  Index
root' <- ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ if Index -> Bool
P.nullIndex Index
root then DynSegTree (PrimState m) a -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> ST s Index
newRootST DynSegTree (PrimState m) a
dst else DynSegTree (PrimState m) a -> Index -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree (PrimState m) a
dst Index
root
  Index -> Int -> Int -> m Index
inner Index
root' Int
l0Dst Int
r0Dst
  where
    !()
_ = HasCallStack => String -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> ()
ACIA.checkIndexBounded String
"AtCoder.Extra.DynSegTree.Raw.modifyMST" Int
i Int
l0Dst Int
r0Dst
    -- `c` is already cloned or newly allocated
    inner :: P.Index -> Int -> Int -> m P.Index
    inner :: Index -> Int -> Int -> m Index
inner Index
c Int
l Int
r
      | Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = do
          -- let !_ = ACIA.runtimeAssert (i == l) ""
          MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
xDst a -> m a
f (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
          Index -> m Index
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c
      | Bool
otherwise = do
          let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2

          -- one of left or right child is updated, not both
          if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m
            then do
              Index
cl <- ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ do
                Index
j <- MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
                if Index -> Bool
P.nullIndex Index
j
                  then DynSegTree (PrimState m) a -> Int -> Int -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Int -> Int -> ST s Index
newNodeInST DynSegTree (PrimState m) a
dst Int
l Int
m
                  else DynSegTree (PrimState m) a -> Index -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree (PrimState m) a
dst Index
j
              ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST (PrimState m))) Index
-> Int -> Index -> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c) Index
cl
              Index
_ <- Index -> Int -> Int -> m Index
inner Index
cl Int
l Int
m
              () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            else do
              Index
cr <- ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ do
                Index
j <- MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
                if Index -> Bool
P.nullIndex Index
j
                  then DynSegTree (PrimState m) a -> Int -> Int -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Int -> Int -> ST s Index
newNodeInST DynSegTree (PrimState m) a
dst Int
m Int
r
                  else DynSegTree (PrimState m) a -> Index -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree (PrimState m) a
dst Index
j
              ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST (PrimState m))) Index
-> Int -> Index -> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c) Index
cr
              Index
_ <- Index -> Int -> Int -> m Index
inner Index
cr Int
m Int
r
              () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

          ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
            Index
cl <- MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
            a
clx <- if Index -> Bool
P.nullIndex Index
cl then a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> a -> ST (PrimState m) a
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
l Int
m else MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
cl)
            Index
cr <- MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
            a
crx <- if Index -> Bool
P.nullIndex Index
cr then a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> a -> ST (PrimState m) a
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
m Int
r else MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
cr)
            MVector (PrimState (ST (PrimState m))) a
-> Int -> a -> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c) (a -> ST (PrimState m) ()) -> a -> ST (PrimState m) ()
forall a b. (a -> b) -> a -> b
$! a
clx a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
crx
          Index -> m Index
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c

-- | \(O(\log L)\)
--
-- @since 1.2.1.0
{-# INLINEABLE prodST #-}
prodST :: forall a s. (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> P.Index -> Int -> Int -> ST s a
prodST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> Int -> Int -> ST s a
prodST DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} Index
root0 Int
ql0 Int
qr0
  | Int
ql0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
qr0 Bool -> Bool -> Bool
|| Index -> Bool
P.nullIndex Index
root0 = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
  | Bool
otherwise = Index -> Int -> Int -> Int -> Int -> a -> ST s a
inner Index
root0 Int
l0Dst Int
r0Dst Int
ql0 Int
qr0 a
forall a. Monoid a => a
mempty
  where
    !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkIntervalBounded String
"AtCoder.Extra.DynSegTree.Raw.prodST" Int
ql0 Int
qr0 Int
l0Dst Int
r0Dst
    -- left to right
    -- - l, r: node interval
    -- - ql, qr: queried interval
    inner :: P.Index -> Int -> Int -> Int -> Int -> a -> ST s a
    inner :: Index -> Int -> Int -> Int -> Int -> a -> ST s a
inner Index
c Int
l Int
r Int
ql_ Int
qr_ !a
x
      | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
      | Index -> Bool
P.nullIndex Index
c = do
          a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> Int -> Int -> a
initialProdDst Int
ql Int
qr
      | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
ql Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
qr = do
          a
cx <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
          a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
cx
      | Bool
otherwise = do
          let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
          Index
cl <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
          a
x' <- Index -> Int -> Int -> Int -> Int -> a -> ST s a
inner Index
cl Int
l Int
m Int
ql Int
qr a
x
          Index
cr <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
          Index -> Int -> Int -> Int -> Int -> a -> ST s a
inner Index
cr Int
m Int
r Int
ql Int
qr a
x'
      where
        -- shrink target interval to node interval
        ql :: Int
ql = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
ql_ Int
l
        qr :: Int
qr = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
qr_ Int
r
        len :: Int
len = Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ql

-- | \(O(\log L)\) Resets an interval \([l, r)\) to initial monoid values.
--
-- @since 1.2.1.0
{-# INLINEABLE resetIntervalST #-}
resetIntervalST ::
  forall a s.
  (HasCallStack, Monoid a, VU.Unbox a) =>
  DynSegTree s a ->
  P.Index ->
  Int ->
  Int ->
  ST s P.Index
resetIntervalST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> Int -> Int -> ST s Index
resetIntervalST dst :: DynSegTree s a
dst@DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} Index
root Int
ql0 Int
qr0
  | Int
ql0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
qr0 = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
root
  | Index -> Bool
P.nullIndex Index
root = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
P.undefIndex
  | Int
ql0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l0Dst Bool -> Bool -> Bool
&& Int
qr0 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r0Dst = do
      -- for the case of non-persistent segment tere, we should update the root in-place:
      Index
root' <- DynSegTree s a -> Index -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree s a
dst Index
root
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
root') (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
l0Dst Int
r0Dst
      MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
root') Index
P.undefIndex
      MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
root') Index
P.undefIndex
      Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
root'
  | Bool
otherwise = Index -> Int -> Int -> Int -> Int -> ST s Index
inner Index
root Int
l0Dst Int
r0Dst Int
ql0 Int
qr0
  where
    !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkIntervalBounded String
"AtCoder.Extra.DynSegTree.Raw.resetIntervalST" Int
ql0 Int
qr0 Int
l0Dst Int
r0Dst

    -- replace interval with null
    inner :: P.Index -> Int -> Int -> Int -> Int -> ST s P.Index
    inner :: Index -> Int -> Int -> Int -> Int -> ST s Index
inner Index
c Int
l Int
r Int
ql_ Int
qr_
      -- TODO: shall we allocate new node?
      | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c
      | Index -> Bool
P.nullIndex Index
c = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
P.undefIndex
      -- NOTE: we're returning `undefIndex`, but we can instead free the subtree if it's not persistent
      | Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Bool -> Bool -> Bool
&& Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
qr = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
P.undefIndex
      | Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c
      | Bool
otherwise = do
          let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
          Index
c' <- DynSegTree s a -> Index -> ST s Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree s a
dst Index
c
          MVector (PrimState (ST s)) Index
-> (Index -> ST s Index) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector s Index
MVector (PrimState (ST s)) Index
lDst (\Index
i -> Index -> Int -> Int -> Int -> Int -> ST s Index
inner Index
i Int
l Int
m Int
ql Int
qr) (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c')
          Index
cl <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c')
          MVector (PrimState (ST s)) Index
-> (Index -> ST s Index) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector s Index
MVector (PrimState (ST s)) Index
rDst (\Index
i -> Index -> Int -> Int -> Int -> Int -> ST s Index
inner Index
i Int
m Int
r Int
ql Int
qr) (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c')
          Index
cr <- MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c')
          a
clx <- if Index -> Bool
P.nullIndex Index
cl then a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
l Int
m else MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
cl)
          a
crx <- if Index -> Bool
P.nullIndex Index
cr then a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! Int -> Int -> a
initialProdDst Int
m Int
r else MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
cr)
          MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c') (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
clx a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
crx
          Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c'
      where
        -- shrink target interval to node interval
        ql :: Int
ql = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
ql_ Int
l
        qr :: Int
qr = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
qr_ Int
r
        len :: Int
len = Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ql

-- | \(O(\log L)\)
--
-- @since 1.2.1.0
{-# INLINEABLE maxRightM #-}
maxRightM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DynSegTree (PrimState m) a -> P.Index -> (a -> m Bool) -> m Int
maxRightM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DynSegTree (PrimState m) a -> Index -> (a -> m Bool) -> m Int
maxRightM dst :: DynSegTree (PrimState m) a
dst@DynSegTree {Bool
Int
MVector (PrimState m) a
MVector (PrimState m) Index
Pool (PrimState m) ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool (PrimState m) ()
lDst :: MVector (PrimState m) Index
rDst :: MVector (PrimState m) Index
xDst :: MVector (PrimState m) a
..} Index
root a -> m Bool
f = do
  (!Int
r, !a
_) <- Index -> Int -> Int -> a -> m (Int, a)
inner Index
root Int
l0Dst Int
r0Dst a
forall a. Monoid a => a
mempty
  Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
r
  where
    -- FIXME: it should not allocate new nodes
    inner :: Index -> Int -> Int -> a -> m (Int, a)
inner Index
c_ Int
l Int
r !a
x = do
      Index
c <- if Index -> Bool
P.nullIndex Index
c_ then ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ DynSegTree (PrimState m) a -> Int -> Int -> ST (PrimState m) Index
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Int -> Int -> ST s Index
newNodeInST DynSegTree (PrimState m) a
dst Int
l Int
r else Index -> m Index
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c_
      a
xWhole <- ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ (a
x <>) (a -> a) -> ST (PrimState m) a -> ST (PrimState m) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
      Bool
b <- a -> m Bool
f a
xWhole
      if Bool
b
        then do
          (Int, a) -> m (Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
r, a
xWhole)
        else do
          if Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
            then (Int, a) -> m (Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
l, a
x)
            else do
              let m :: Int
m = (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
              Index
cl <- ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
              (!Int
k, !a
xl) <- Index -> Int -> Int -> a -> m (Int, a)
inner Index
cl Int
l Int
m a
x
              if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m
                then (Int, a) -> m (Int, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
k, a
xl)
                else do
                  Index
cr <- ST (PrimState m) Index -> m Index
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Index -> m Index)
-> ST (PrimState m) Index -> m Index
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST (PrimState m))) Index
-> Int -> ST (PrimState m) Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) Index
MVector (PrimState (ST (PrimState m))) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
                  Index -> Int -> Int -> a -> m (Int, a)
inner Index
cr Int
m Int
r a
xl

-- -------------------------------------------------------------------------------------------------
-- Internals
-- -------------------------------------------------------------------------------------------------

-- | \(O(1)\) Optionally clones a node depending on the persistency setting.
{-# INLINEABLE cloneOnWriteST #-}
cloneOnWriteST :: (HasCallStack, Monoid a, VU.Unbox a) => DynSegTree s a -> P.Index -> ST s P.Index
cloneOnWriteST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DynSegTree s a -> Index -> ST s Index
cloneOnWriteST DynSegTree {Bool
Int
MVector s a
MVector s Index
Pool s ()
Int -> Int -> a
capacityDst :: forall s a. DynSegTree s a -> Int
isPersistentDst :: forall s a. DynSegTree s a -> Bool
l0Dst :: forall s a. DynSegTree s a -> Int
r0Dst :: forall s a. DynSegTree s a -> Int
initialProdDst :: forall s a. DynSegTree s a -> Int -> Int -> a
poolDst :: forall s a. DynSegTree s a -> Pool s ()
lDst :: forall s a. DynSegTree s a -> MVector s Index
rDst :: forall s a. DynSegTree s a -> MVector s Index
xDst :: forall s a. DynSegTree s a -> MVector s a
capacityDst :: Int
isPersistentDst :: Bool
l0Dst :: Int
r0Dst :: Int
initialProdDst :: Int -> Int -> a
poolDst :: Pool s ()
lDst :: MVector s Index
rDst :: MVector s Index
xDst :: MVector s a
..} Index
c
  | Bool -> Bool
not Bool
isPersistentDst Bool -> Bool -> Bool
|| Index -> Bool
P.nullIndex Index
c = Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
c
  | Bool
otherwise = do
      Index
i <- Pool (PrimState (ST s)) () -> () -> ST s Index
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Unbox a) =>
Pool (PrimState m) a -> a -> m Index
P.alloc Pool s ()
Pool (PrimState (ST s)) ()
poolDst ()
      MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) (Index -> ST s ()) -> ST s Index -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
lDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
      MVector (PrimState (ST s)) Index -> Int -> Index -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) (Index -> ST s ()) -> ST s Index -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) Index -> Int -> ST s Index
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s Index
MVector (PrimState (ST s)) Index
rDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
      MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
i) (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
xDst (Index -> Int
forall a b. Coercible a b => a -> b
coerce Index
c)
      Index -> ST s Index
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Index
i