{-# LANGUAGE BangPatterns #-}
{-|
== Mutable segment tree

See SegTree. SegTreeMut is just that, but backed by a mutable array.
With unboxed arrays, SegTreeMut is multiple times faster than SegTree (see benchmarks).
However, this comes at the cost of purity.

Sources:

* Al.Cash, "Efficient and easy segment trees"
  https://codeforces.com/blog/entry/18051
* AtCoder Library
  https://github.com/atcoder/ac-library/blob/master/atcoder/segtree.hpp

Let n = r - l + 1 where (l, r) is the range of the segment tree.
The complexities assume (<>) takes O(1) time.

-}

{-
Implementation notes:
* The implementation here is a bottom-up implementation, sometimes called "iterative" segment tree.
  A top-down implementation also works, but is generally slower.
* The values of the tree are stored in an Int-indexed array. Nodes are represented by Ints, which
  are the indices where the value of the node is stored. The value of the root is stored at index 1.
  The childen of node i are 2*i and 2*i+1. Each node covers a range the size of which is a power of
  2, and the two children of the node, if it's not a leaf, cover the two halves of that range.
* Operations on the tree proceed bottom-up, from the leaves towards the root. Due to the uniform
  structure of the tree it is easy and efficient to traverse the tree by manipulating indices.
* INLINE on setSNM is critical!
* The "unsafe" bitwise ops are safe in the given conditions. For instance,
  x `unsafeShiftR` (ctz (x+1)) is safe if x is not -1. These are used only where it will be run many
  many times and will make a difference.
-}

module SegTreeMut
    ( SegTreeMut
    , emptySTM
    , fromListSTM
    , adjustSTM
    , foldRangeSTM
    , binSearchSTM
    , foldrSTM
    ) where

import Control.Monad
import Data.Array.MArray
import Data.Bits
import Data.List

import Misc ( bitLength, modifyArray', unsafeBit )

data SegTreeMut marr a = STM !Int !Int !Int !(marr Int a)

-- | Builds a segment tree on (l, r) where each element is mempty. O(n).
emptySTM :: (Monoid a, MArray marr a m) => (Int, Int) -> m (SegTreeMut marr a)
emptySTM :: (Int, Int) -> m (SegTreeMut marr a)
emptySTM (Int
l,Int
r)
    | Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = [Char] -> m (SegTreeMut marr a)
forall a. HasCallStack => [Char] -> a
error [Char]
"emptySTM: bad range"
    | Bool
otherwise = do
        let n :: Int
n = Int -> Int
forall a. Bits a => Int -> a
bit (Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l))
        marr Int a
xa <- (Int, Int) -> a -> m (marr Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
forall a. Monoid a => a
mempty
        SegTreeMut marr a -> m (SegTreeMut marr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegTreeMut marr a -> m (SegTreeMut marr a))
-> SegTreeMut marr a -> m (SegTreeMut marr a)
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int -> marr Int a -> SegTreeMut marr a
forall (marr :: * -> * -> *) a.
Int -> Int -> Int -> marr Int a -> SegTreeMut marr a
STM Int
l Int
r Int
n marr Int a
xa

setSNM :: (Monoid a, MArray marr a m) => marr Int a -> Int -> m ()
setSNM :: marr Int a -> Int -> m ()
setSNM marr Int a
xa = \Int
i -> a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>) (a -> a -> a) -> m a -> m (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> marr Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) m (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> marr Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (marr Int a -> Int -> a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray marr Int a
xa Int
i (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$!)
{-# INLINE setSNM #-}

-- | Builds a segment tree on (l, r) where the elements are taken from a list. If the list is shorter
-- than the range, the remaining elements are mempty. O(n).
fromListSTM :: (Monoid a, MArray marr a m) => (Int, Int) -> [a] -> m (SegTreeMut marr a)
fromListSTM :: (Int, Int) -> [a] -> m (SegTreeMut marr a)
fromListSTM (Int
l,Int
r) [a]
xs
    | Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = [Char] -> m (SegTreeMut marr a)
forall a. HasCallStack => [Char] -> a
error [Char]
"fromListSTM: bad range"
    | Bool
otherwise = do
        let n :: Int
n = Int -> Int
forall a. Bits a => Int -> a
bit (Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l))
        marr Int a
xa <- (Int, Int) -> a -> m (marr Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
forall a. Monoid a => a
mempty
        [(Int, a)] -> ((Int, a) -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
n .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l] [a]
xs) (((Int, a) -> m ()) -> m ()) -> ((Int, a) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ (Int -> a -> m ()) -> (Int, a) -> m ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (marr Int a -> Int -> a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray marr Int a
xa)
        [Int] -> (Int -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2 .. Int
1] ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ marr Int a -> Int -> m ()
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
marr Int a -> Int -> m ()
setSNM marr Int a
xa
        SegTreeMut marr a -> m (SegTreeMut marr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SegTreeMut marr a -> m (SegTreeMut marr a))
-> SegTreeMut marr a -> m (SegTreeMut marr a)
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int -> marr Int a -> SegTreeMut marr a
forall (marr :: * -> * -> *) a.
Int -> Int -> Int -> marr Int a -> SegTreeMut marr a
STM Int
l Int
r Int
n marr Int a
xa

-- | Adjusts the element at index i. O(log n).
adjustSTM :: (Monoid a, MArray marr a m) => SegTreeMut marr a -> Int -> (a -> a) -> m ()
adjustSTM :: SegTreeMut marr a -> Int -> (a -> a) -> m ()
adjustSTM (STM Int
l Int
r Int
n marr Int a
xa) Int
qi a -> a
f
    | Int
qi 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
qi = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"adjustSTM: outside range"
    | Bool
otherwise = do
        let qi' :: Int
qi' = Int
qi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n
        marr Int a -> Int -> (a -> a) -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' marr Int a
xa Int
qi' a -> a
f
        (Int -> m ()) -> [Int] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (marr Int a -> Int -> m ()
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
marr Int a -> Int -> m ()
setSNM marr Int a
xa) ([Int] -> m ()) -> [Int] -> m ()
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate' (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2) (Int
qi' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)

-- | Folds the elements in the range (ql, qr). Elements outside (l, r) are considered to be mempty.
-- O(log n), or more precisely O(log(min(r, qr) - max(l, ql) + 1)).
foldRangeSTM :: (Monoid a, MArray marr a m) => SegTreeMut marr a -> Int -> Int -> m a
foldRangeSTM :: SegTreeMut marr a -> Int -> Int -> m a
foldRangeSTM (STM Int
l0 Int
r0 Int
n marr Int a
xa) Int
ql Int
qr
    | Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql = [Char] -> m a
forall a. HasCallStack => [Char] -> a
error [Char]
"foldRangeSTM: bad range"
    | Int
qr' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql'   = a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
    | Bool
otherwise = do
        a
accL <- Int -> Int -> Int -> a -> m a
forall (f :: * -> *).
MArray marr a f =>
Int -> Int -> Int -> a -> f a
goUpRt (Int
ql' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
ql' Int
0 a
forall a. Monoid a => a
mempty
        a
accR <- Int -> Int -> Int -> a -> m a
forall (f :: * -> *).
MArray marr a f =>
Int -> Int -> Int -> a -> f a
goUpLt (Int
qr' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
qr' Int
0 a
forall a. Monoid a => a
mempty
        a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$! a
accL a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
accR
  where
    ql' :: Int
ql' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
l0 Int
ql
    qr' :: Int
qr' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
r0 Int
qr
    goUpRt :: Int -> Int -> Int -> a -> f a
goUpRt !Int
i1 !Int
l !Int
d1 !a
acc = do
        let tz :: Int
tz = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
i1
            i :: Int
i = Int
i1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
tz
            d :: Int
d = Int
d1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
tz
        if Int
qr' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        then a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
        else marr Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa Int
i f a -> (a -> f a) -> f a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> Int -> a -> f a
goUpRt (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d) Int
d (a -> f a) -> (a -> a) -> a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<>)
    goUpLt :: Int -> Int -> Int -> a -> f a
goUpLt !Int
j1 !Int
r !Int
d1 !a
acc = do
        let to :: Int
to = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Int
j1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
            j :: Int
j = Int
j1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
to
            d :: Int
d = Int
d1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
to
        if Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql'
        then a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
        else marr Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa Int
j f a -> (a -> f a) -> f a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> Int -> a -> f a
goUpLt (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d) Int
d (a -> f a) -> (a -> a) -> a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc)

-- | Binary search in the intersection of (l, r) and (ql, qr) for the shortest prefix whose fold
-- satisfies the given monotonic predicate. Returns the end index and the fold. O(log n), or more
-- precisely O(log(min(r, qr) - max(l, ql) + 1)).
binSearchSTM :: (Monoid a, MArray marr a m)
             => SegTreeMut marr a -> Int -> Int -> (a -> Bool) -> m (Maybe (Int, a))
binSearchSTM :: SegTreeMut marr a
-> Int -> Int -> (a -> Bool) -> m (Maybe (Int, a))
binSearchSTM (STM Int
l0 Int
r0 Int
n marr Int a
xa) Int
ql Int
qr a -> Bool
p
    | Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql = [Char] -> m (Maybe (Int, a))
forall a. HasCallStack => [Char] -> a
error [Char]
"binSearchSTM: bad range"
    | Int
qr' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql'   = Maybe (Int, a) -> m (Maybe (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, a)
forall a. Maybe a
Nothing
    | Bool
otherwise   = Int -> Int -> Int -> a -> m (Maybe (Int, a))
forall (m :: * -> *).
MArray marr a m =>
Int -> Int -> Int -> a -> m (Maybe (Int, a))
goUpRt (Int
ql' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
ql' Int
0 a
forall a. Monoid a => a
mempty
  where
    ql' :: Int
ql' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
l0 Int
ql
    qr' :: Int
qr' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
r0 Int
qr
    goUpRt :: Int -> Int -> Int -> a -> m (Maybe (Int, a))
goUpRt !Int
i1 !Int
l !Int
d1 !a
acc = do
        let tz :: Int
tz = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
i1
            i :: Int
i = Int
i1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
tz
            d :: Int
d = Int
d1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
tz
        !a
acc' <- (a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<>) (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> marr Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa Int
i
        case () of
            ()
_ | a -> Bool
p a
acc'                 -> Int -> Int -> Int -> a -> m (Maybe (Int, a))
forall (m :: * -> *).
MArray marr a m =>
Int -> Int -> Int -> a -> m (Maybe (Int, a))
goDn Int
i Int
l Int
d a
acc
              | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
qr' -> Int -> Int -> Int -> a -> m (Maybe (Int, a))
goUpRt (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit Int
d) Int
d a
acc'
              | Bool
otherwise              -> Maybe (Int, a) -> m (Maybe (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, a)
forall a. Maybe a
Nothing
    goDn :: Int -> Int -> Int -> a -> f (Maybe (Int, a))
goDn !Int
i !Int
l !Int
d !a
acc
        | Int
qr' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = Maybe (Int, a) -> f (Maybe (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, a)
forall a. Maybe a
Nothing
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = do
            !a
acc' <- (a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<>) (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> marr Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i)
            if a -> Bool
p a
acc'
            then Int -> Int -> Int -> a -> f (Maybe (Int, a))
goDn (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
acc
            else Int -> Int -> Int -> a -> f (Maybe (Int, a))
goDn (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a. (Bits a, Num a) => Int -> a
unsafeBit (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
acc'
        | Bool
otherwise = do
            !a
acc' <- (a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<>) (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> marr Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa Int
i
            Maybe (Int, a) -> f (Maybe (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just (Int
l, a
acc'))

-- | Right fold over the elements of the segment tree. O(n).
foldrSTM :: (Monoid a, MArray marr a m) => SegTreeMut marr a -> (a -> b -> b) -> b -> m b
foldrSTM :: SegTreeMut marr a -> (a -> b -> b) -> b -> m b
foldrSTM (STM Int
l Int
r Int
n marr Int a
xa) a -> b -> b
f b
z = (Int -> m b -> m b) -> m b -> [Int] -> m b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((a -> b -> b) -> m a -> m b -> m b
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 a -> b -> b
f (m a -> m b -> m b) -> (Int -> m a) -> Int -> m b -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. marr Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marr Int a
xa) (b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
z) [Int
n .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l]

--------------------------------------------------------------------------------
-- For tests

-- Allows specialization across modules
{-# INLINABLE fromListSTM #-}
{-# INLINABLE adjustSTM #-}
{-# INLINABLE foldRangeSTM #-}
{-# INLINABLE binSearchSTM #-}