content: {-# LANGUAGE RecordWildCards #-}
  
  -- | It is the data structure for [monoids](https://en.wikipedia.org/wiki/Monoid)
  -- \((S, \cdot: S \times S \to S, e \in S)\), i.e., the algebraic structure that satisfies the
  -- following properties.
  --
  -- - associativity: \((a \cdot b) \cdot c = a \cdot (b \cdot c)\) for all \(a, b, c \in S\)
  -- - existence of the identity element: \(a \cdot e = e \cdot a = a\) for all \(a \in S\)
  --
  -- Given an array \(S\) of length \(N\), it processes the following queries in \(O(\log N)\) time
  -- (see [Appendix](./appendix.html) for further details).
  --
  -- - Updating an element
  -- - Calculating the product of the elements of an interval
  --
  -- For simplicity, in this document, we assume that the oracles @op@ and @e@ work in constant time.
  -- If these oracles work in \(O(T)\) time, each time complexity appearing in this document is
  -- multiplied by \(O(T)\).
  --
  -- ==== __Example__
  -- Create a `SegTree` of @'Sum' Int@:
  --
  -- >>> import AtCoder.SegTree qualified as ST
  -- >>> import Data.Vector.Unboxed qualified as VU
  -- >>> import Data.Monoid (Sum(..))
  -- >>> seg <- ST.new @_ @(Sum Int) 4
  --
  -- Modify the vertex values:
  --
  -- >>> ST.write seg 1 $ Sum 1
  -- >>> ST.modify seg (+ Sum 2) 2
  -- >>> ST.write seg 3 $ Sum 3 -- [0, 1, 2, 3]
  -- >>> ST.read seg 1
  -- Sum {getSum = 1}
  --
  -- Get product of the monoids:
  --
  -- >>> ST.prod seg 0 3
  -- Sum {getSum = 3}
  --
  -- >>> ST.allProd seg
  -- Sum {getSum = 6}
  --
  -- Binary searches:
  --
  -- >>> ST.maxRight seg 0 (< (Sum 5)) -- sum [0, 3) = 2 < 5
  -- 3
  --
  -- >>> ST.minLeft seg 4 (< (Sum 5)) -- sum [3, 4) = 3 < 5
  -- 3
  --
  -- Inspect all the values in \(O(n)\) with `freeze` or in \(O(1)\) with `unsafeFreeze`:
  --
  -- >>> VU.map getSum <$> ST.freeze seg
  -- [0,1,2,3]
  --
  -- ==== Tips
  --
  -- - `prod` returns \(a_l \cdot a_{l + 1} \cdot .. \cdot a_{r - 1}\). If you need \(a_{r - 1} \cdot a_{r - 2} \cdot .. \cdot a_{l}\),
  -- wrap your monoid in `Data.Monoid.Dual`.
  -- - If you ever need to store boxed types to `LazySegTree`, wrap it in @Data.Vector.Unboxed.DoNotUnboxStrict@
  -- or the like.
  --
  -- ==== Major changes from the original @ac-library@
  -- - The implementation is `Monoid` based, not function objects.
  -- - @get@ and @set@ are renamed to `read` and `write`.
  -- - `modify`, `modifyM`, `exchange`, `freeze` and `unsafeFreeze` are added.
  --
  -- @since 1.0.0.0
  module AtCoder.SegTree
    ( -- * Segment tree
      SegTree (nSt, sizeSt, logSt),
  
      -- * Constructors
      new,
      build,
  
      -- * Accessing elements
      write,
      modify,
      modifyM,
      exchange,
      read,
  
      -- * Products
      prod,
      prodMaybe,
      allProd,
  
      -- * Binary searches
  
      -- ** Left binary searches
      minLeft,
      minLeftM,
  
      -- ** Right binary searches
      maxRight,
      maxRightM,
  
      -- * Conversions
      freeze,
      unsafeFreeze,
    )
  where
  
  import AtCoder.Internal.Assert qualified as ACIA
  import AtCoder.Internal.Bit qualified as ACIBIT
  import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
  import Control.Monad.ST (ST)
  import Data.Bits (countTrailingZeros, testBit, (.&.), (.>>.))
  import Data.Foldable (for_)
  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)
  
  -- | Segment tree.
  --
  -- @since 1.0.0.0
  data SegTree s a = SegTree
    { -- | The number of vertices.
      --
      -- @since 1.0.0.0
      nSt :: {-# UNPACK #-} !Int,
      -- | \(\lceil \log_2 \mathrm{nSt} \rceil\).
      --
      -- @since 1.0.0.0
      sizeSt :: {-# UNPACK #-} !Int,
      -- | \(\log_2 \mathrm{sizeSt}\).
      --
      -- @since 1.0.0.0
      logSt :: {-# UNPACK #-} !Int,
      -- | Data storage of length @2 * sizeSt@.
      dSt :: !(VUM.MVector s a)
    }
  
  -- | Creates an array \(a\) of length \(n\). All the elements are initialized to `mempty`.
  --
  -- ==== Constraints
  -- - \(0 \leq n\)
  --
  -- ==== Complexity
  -- - \(O(n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE new #-}
  new :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (SegTree (PrimState m) a)
  new n = stToPrim $ newST n
  
  -- | Creates an array with initial values.
  --
  -- ==== Complexity
  -- - \(O(n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE build #-}
  build :: (PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector a -> m (SegTree (PrimState m) a)
  build vs = stToPrim $ buildST vs
  
  -- | Writes \(p\)-th value of the array to \(x\).
  --
  -- ==== Constraints
  -- - \(0 \leq p \lt n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE write #-}
  write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m ()
  write self p x = stToPrim $ writeST self p x
  
  -- | (Extra API) Given a user function \(f\), modifies \(p\)-th value with it.
  --
  -- ==== Constraints
  -- - \(0 \leq p \lt n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE modify #-}
  modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
  modify self f p = stToPrim $ modifyST self f p
  
  -- | (Extra API) Given a user function \(f\), modifies \(p\)-th value with it.
  --
  -- ==== Constraints
  -- - \(0 \leq p \lt n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE modifyM #-}
  modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
  modifyM self@SegTree {..} f p = do
    let !_ = ACIA.checkIndex "AtCoder.SegTree.modifyM" p nSt
    VGM.modifyM dSt f (p + sizeSt)
    stToPrim $ for_ [1 .. logSt] $ \i -> do
      updateST self ((p + sizeSt) .>>. i)
  
  -- | (Extra API) Writes \(p\)-th value of the array to \(x\) and returns the old value.
  --
  -- ==== Constraints
  -- - \(0 \leq p \lt n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.1.0.0
  {-# INLINE exchange #-}
  exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m a
  exchange self p x = stToPrim $ exchangeST self p x
  
  -- | Returns \(p\)-th value of the array.
  --
  -- ==== Constraints
  -- - \(0 \leq p \lt n\)
  --
  -- ==== Complexity
  -- - \(O(1)\)
  --
  -- @since 1.0.0.0
  {-# INLINE read #-}
  read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> m a
  read SegTree {..} p = do
    let !_ = ACIA.checkIndex "AtCoder.SegTree.read" p nSt
    VGM.read dSt $ p + sizeSt
  
  -- | Returns \(a[l] \cdot ... \cdot a[r - 1]\), assuming the properties of the monoid. It
  -- returns `mempty` if \(l = r\).
  --
  -- ==== Constraints
  -- - \(0 \leq l \leq r \leq n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE prod #-}
  prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m a
  prod self@SegTree {nSt} l0 r0
    | ACIA.testInterval l0 r0 nSt = stToPrim $ unsafeProdST self l0 r0
    | otherwise = ACIA.errorInterval "AtCoder.SegTree.prod" l0 r0 nSt
  
  -- | Total variant of `prod`. Returns \(a[l] \cdot ... \cdot a[r - 1]\), assuming the properties of
  -- the monoid. It returns `Just` `mempty` if \(l = r\). Returns `Nothing` if the interval is
  -- invalid.
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE prodMaybe #-}
  prodMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m (Maybe a)
  prodMaybe self@SegTree {nSt} l0 r0
    | ACIA.testInterval l0 r0 nSt = stToPrim $ Just <$> unsafeProdST self l0 r0
    -- l0 == r0 = pure (Just mempty)
    | otherwise = pure Nothing
  
  -- | Returns @a[0] <> ... <> a[n - 1]@, assuming the properties of the monoid. It returns `mempty`
  -- if \(n = 0\).
  --
  -- ==== Complexity
  -- - \(O(1)\)
  --
  -- @since 1.0.0.0
  {-# INLINE allProd #-}
  allProd :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> m a
  allProd SegTree {..} = VGM.read dSt 1
  
  -- | Applies a binary search on the segment tree. It returns an index \(l\) that satisfies both of
  -- the following.
  --
  -- - \(l = r\) or \(f(a[l] \cdot a[l + 1] \cdot ... \cdot a[r - 1])\) returns `True`.
  -- - \(l = 0\) or \(f(a[l - 1] \cdot a[l] \cdot ... \cdot a[r - 1])\) returns `False`.
  --
  -- If \(f\) is monotone, this is the minimum \(l\) that satisfies
  -- \(f(a[l] \cdot a[l + 1] \cdot ... \cdot a[r - 1])\).
  --
  -- ==== Constraints
  --
  -- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side
  --   effect.
  -- - @f mempty == True@.
  -- - \(0 \leq r \leq n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE minLeft #-}
  minLeft ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | The segment tree
    SegTree (PrimState m) a ->
    -- | \(r\)
    Int ->
    -- | \(p\): user predicate
    (a -> Bool) ->
    -- | \(l\): \(p\) holds for \([l, r)\)
    m Int
  minLeft seg r0 f = minLeftM seg r0 (pure . f)
  
  -- | Monadic variant of `minLeft`.
  --
  -- ==== Constraints
  --
  -- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side
  --   effect.
  -- - @f mempty == True@.
  -- - \(0 \leq r \leq n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINEABLE minLeftM #-}
  minLeftM ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | The segment tree
    SegTree (PrimState m) a ->
    -- | \(r\)
    Int ->
    -- | \(p\): user predicate
    (a -> m Bool) ->
    -- | \(l\): \(p\) holds for \([l, r)\)
    m Int
  minLeftM SegTree {..} r0 f = do
    b <- f mempty
    let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.minLeftM: `f empty` returned `False`"
    if r0 == 0
      then pure 0
      else inner (r0 + sizeSt) mempty
    where
      -- NOTE: Not ordinary bounds check!
      !_ = ACIA.runtimeAssert (0 <= r0 && r0 <= nSt) $ "AtCoder.SegTree.minLeftM: given invalid `right` index `" ++ show r0 ++ "` over length `" ++ show nSt ++ "`"
      inner r !sm = do
        let r' = chooseBit $ r - 1
        !sm' <- stToPrim $ (<> sm) <$> VGM.read dSt r'
        b <- f sm'
        if not b
          then do
            inner2 r' sm
          else do
            if (r' .&. (-r')) /= r'
              then inner r' sm'
              else pure 0
      chooseBit r
        | r > 1 && odd r = chooseBit $ r .>>. 1
        | otherwise = r
      inner2 r sm
        | r < sizeSt = do
            let r' = 2 * r + 1
            !sm' <- stToPrim $ (<> sm) <$> VGM.read dSt r'
            b <- f sm'
            if b
              then inner2 (r' - 1) sm'
              else inner2 r' sm
        | otherwise = pure $ r + 1 - sizeSt
  
  -- | Applies a binary search on the segment tree. It returns an index \(r\) that satisfies both of the
  -- following.
  --
  -- - \(r = l\) or \(f(a[l] \cdot a[l + 1] \cdot ... \cdot a[r - 1])\) returns `True`.
  -- - \(r = n\) or \(f(a[l] \cdot a[l + 1] \cdot ... \cdot a[r]))\) returns `False`.
  --
  -- If \(f\) is monotone, this is the maximum \(r\) that satisfies
  -- \(f(a[l] \cdot a[l + 1] \cdot ... \cdot a[r - 1])\).
  --
  -- ==== Constraints
  -- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect.
  -- - @f mempty == True@.
  -- - \(0 \leq l \leq n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINE maxRight #-}
  maxRight ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | The segment tree
    SegTree (PrimState m) a ->
    -- | \(l\)
    Int ->
    -- | \(p\): user predicate
    (a -> Bool) ->
    -- | \(r\): \(p\) holds for \([l, r)\)
    m Int
  maxRight seg l0 f = maxRightM seg l0 (pure . f)
  
  -- | Monadic variant of `maxRight`.
  --
  -- ==== Constraints
  -- - if \(f\) is called with the same argument, it returns the same value, i.e., \(f\) has no side effect.
  -- - @f mempty == True@.
  -- - \(0 \leq l \leq n\)
  --
  -- ==== Complexity
  -- - \(O(\log n)\)
  --
  -- @since 1.0.0.0
  {-# INLINEABLE maxRightM #-}
  maxRightM ::
    (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
    -- | The segment tree
    SegTree (PrimState m) a ->
    -- | \(l\)
    Int ->
    -- | \(p\): user predicate
    (a -> m Bool) ->
    -- | \(r\): \(p\) holds for \([l, r)\)
    m Int
  maxRightM SegTree {..} l0 f = do
    b <- f mempty
    let !_ = ACIA.runtimeAssert b "AtCoder.SegTree.maxRightM: `f mempty` must return `True`"
    if l0 == nSt
      then pure nSt
      else inner (l0 + sizeSt) mempty
    where
      -- NOTE: Not ordinary bounds check!
      !_ = ACIA.runtimeAssert (0 <= l0 && l0 <= nSt) $ "AtCoder.SegTree.maxRightM: given invalid `left` index `" ++ show l0 ++ "` over length `" ++ show nSt ++ "`"
      inner l !sm = do
        let l' = chooseBit l
        !sm' <- stToPrim $ (sm <>) <$> VGM.read dSt l'
        b <- f sm'
        if not b
          then do
            inner2 l' sm
          else do
            let l'' = l' + 1
            if (l'' .&. (-l'')) /= l''
              then inner l'' sm'
              else pure nSt
      chooseBit :: Int -> Int
      chooseBit l
        | even l = chooseBit $ l .>>. 1
        | otherwise = l
      inner2 l !sm
        | l < sizeSt = do
            let l' = 2 * l
            !sm' <- stToPrim $ (sm <>) <$> VGM.read dSt l'
            b <- f sm'
            if b
              then inner2 (l' + 1) sm'
              else inner2 l' sm
        | otherwise = pure $ l - sizeSt
  
  -- TODO: add `minLeft`
  
  -- | \(O(n)\) Yields an immutable copy of the mutable vector.
  --
  -- @since 1.0.0.0
  {-# INLINE freeze #-}
  freeze :: (PrimMonad m, VU.Unbox a) => SegTree (PrimState m) a -> m (VU.Vector a)
  freeze SegTree {..} = do
    VU.freeze . VUM.take nSt $ VUM.drop sizeSt dSt
  
  -- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable
  -- vector may not be used after this operation.
  --
  -- @since 1.0.0.0
  {-# INLINE unsafeFreeze #-}
  unsafeFreeze :: (PrimMonad m, VU.Unbox a) => SegTree (PrimState m) a -> m (VU.Vector a)
  unsafeFreeze SegTree {..} = do
    VU.unsafeFreeze . VUM.take nSt $ VUM.drop sizeSt dSt
  
  -- -------------------------------------------------------------------------------------------------
  -- Internal
  -- -------------------------------------------------------------------------------------------------
  
  {-# INLINEABLE newST #-}
  newST :: (HasCallStack, Monoid a, VU.Unbox a) => Int -> ST s (SegTree s a)
  newST nSt
    | nSt >= 0 = build $ VU.replicate nSt mempty
    | otherwise = error $ "AtCoder.SegTree.newST: given negative size (`" ++ show nSt ++ "`)"
  
  {-# INLINEABLE buildST #-}
  buildST :: (Monoid a, VU.Unbox a) => VU.Vector a -> ST s (SegTree s a)
  buildST vs = do
    let nSt = VU.length vs
    let sizeSt = ACIBIT.bitCeil nSt
    let logSt = countTrailingZeros sizeSt
    dSt <- VUM.replicate (2 * sizeSt) mempty
    VU.iforM_ vs $ \i v -> do
      VGM.write dSt (sizeSt + i) v
    let segtree = SegTree {..}
    for_ [sizeSt - 1, sizeSt - 2 .. 1] $ \i -> do
      updateST segtree i
    pure segtree
  
  {-# INLINEABLE writeST #-}
  writeST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> Int -> a -> ST s ()
  writeST self@SegTree {..} p x = do
    let !_ = ACIA.checkIndex "AtCoder.SegTree.writeST" p nSt
    VGM.write dSt (p + sizeSt) x
    for_ [1 .. logSt] $ \i -> do
      updateST self ((p + sizeSt) .>>. i)
  
  {-# INLINEABLE modifyST #-}
  modifyST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> (a -> a) -> Int -> ST s ()
  modifyST self@SegTree {..} f p = do
    let !_ = ACIA.checkIndex "AtCoder.SegTree.modifyST" p nSt
    VGM.modify dSt f (p + sizeSt)
    for_ [1 .. logSt] $ \i -> do
      updateST self ((p + sizeSt) .>>. i)
  
  {-# INLINEABLE exchangeST #-}
  exchangeST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> Int -> a -> ST s a
  exchangeST self@SegTree {..} p x = do
    let !_ = ACIA.checkIndex "AtCoder.SegTree.exchangeST" p nSt
    ret <- VGM.exchange dSt (p + sizeSt) x
    VGM.write dSt (p + sizeSt) x
    for_ [1 .. logSt] $ \i -> do
      updateST self ((p + sizeSt) .>>. i)
    pure ret
  
  {-# INLINEABLE unsafeProdST #-}
  unsafeProdST :: (Monoid a, VU.Unbox a) => SegTree s a -> Int -> Int -> ST s a
  unsafeProdST SegTree {..} l0 r0 = inner (l0 + sizeSt) (r0 + sizeSt - 1) mempty mempty
    where
      -- NOTE: we're using inclusive range [l, r] for simplicity
      inner l r !smL !smR
        | l > r = pure $! smL <> smR
        | otherwise = do
            !smL' <-
              if testBit l 0
                then (smL <>) <$> VGM.read dSt l
                else pure smL
            !smR' <-
              if not $ testBit r 0
                then (<> smR) <$> VGM.read dSt r
                else pure smR
            inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) smL' smR'
  
  {-# INLINE updateST #-}
  updateST :: (Monoid a, VU.Unbox a) => SegTree s a -> Int -> ST s ()
  updateST SegTree {..} k = do
    opL <- VGM.read dSt $ 2 * k
    opR <- VGM.read dSt $ 2 * k + 1
    VGM.write dSt k $! opL <> opR
  