{-|
== Mutable segment tree with lazy propagation

See LazySegTree. LazySegTreeMut is just that, but backed by mutable arrays.
When the arrays are unboxed, LazySegTreeMut is a few times faster than LazySegTree (see benchmarks).
However, this comes at the cost of purity.

-}

{-
Implementation notes:
* INLINE on setLSNM, applyLSNM and pushLSNM is critical!
* All functions are pretty much the same as SegTreeMut except for a call to pushLSNM. See if the
  code can be shared :thonk:
-}

module SegTreeLazyMut
    ( LazySegTreeMut
    , emptyLSTM
    , fromListLSTM
    , adjustLSTM
    , updateRangeLSTM
    , foldRangeLSTM
    , binSearchLSTM
    , foldrLSTM
    ) where

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

import Misc ( Action(..), bitLength, modifyArray' )

data LazySegTreeMut marru marra u a = LSTM !Int !Int !(marru Int u) !(marra Int a)

-- | Builds a segment tree on range (l, r) where each element is mempty. O(n).
emptyLSTM :: (Action u a, MArray marru u m, MArray marra a m)
          => (Int, Int) -> m (LazySegTreeMut marru marra u a)
emptyLSTM :: (Int, Int) -> m (LazySegTreeMut marru marra u a)
emptyLSTM (Int
l,Int
r) | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = [Char] -> m (LazySegTreeMut marru marra u a)
forall a. HasCallStack => [Char] -> a
error [Char]
"emptyLSTM: bad range"
emptyLSTM (Int
l,Int
r) = do
    let n :: Int
n = Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    marru Int u
ua <- (Int, Int) -> u -> m (marru Int u)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int -> Int
forall a. Bits a => Int -> a
bit (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))) u
forall a. Monoid a => a
mempty
    marra Int a
aa <- (Int, Int) -> a -> m (marra Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int -> Int
forall a. Bits a => Int -> a
bit (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))) a
forall a. Monoid a => a
mempty
    LazySegTreeMut marru marra u a
-> m (LazySegTreeMut marru marra u a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LazySegTreeMut marru marra u a
 -> m (LazySegTreeMut marru marra u a))
-> LazySegTreeMut marru marra u a
-> m (LazySegTreeMut marru marra u a)
forall a b. (a -> b) -> a -> b
$! Int
-> Int
-> marru Int u
-> marra Int a
-> LazySegTreeMut marru marra u a
forall (marru :: * -> * -> *) (marra :: * -> * -> *) u a.
Int
-> Int
-> marru Int u
-> marra Int a
-> LazySegTreeMut marru marra u a
LSTM Int
l Int
r marru Int u
ua marra Int a
aa

setLSNM :: (Monoid a, MArray marra a m) => marra Int a -> Int -> m ()
setLSNM :: marra Int a -> Int -> m ()
setLSNM marra Int a
aa 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
<$> marra Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marra Int a
aa (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
<*> marra Int a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marra Int a
aa (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
>>= (marra Int a -> Int -> a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray marra Int a
aa Int
i (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$!)
{-# INLINE setLSNM #-}

-- | 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).
fromListLSTM :: (Action u a, MArray marru u m, MArray marra a m)
             => (Int, Int) -> [a] -> m (LazySegTreeMut marru marra u a)
fromListLSTM :: (Int, Int) -> [a] -> m (LazySegTreeMut marru marra u a)
fromListLSTM (Int
l0,Int
r0) [a]
_ | Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = [Char] -> m (LazySegTreeMut marru marra u a)
forall a. HasCallStack => [Char] -> a
error [Char]
"fromListLSTM: bad range"
fromListLSTM (Int
l0,Int
r0) [a]
xs = do
    let n :: Int
n = Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    marru Int u
ua <- (Int, Int) -> u -> m (marru Int u)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int -> Int
forall a. Bits a => Int -> a
bit (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))) u
forall a. Monoid a => a
mempty
    marra Int a
aa <- (Int, Int) -> a -> m (marra Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
1, Int -> Int
forall a. Bits a => Int -> a
bit (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall b. FiniteBits b => b -> Int
bitLength (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))) a
forall a. Monoid a => a
mempty
    let pop :: StateT [a] m a
pop = ([a] -> m (a, [a])) -> StateT [a] m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT [a] -> m (a, [a])
forall (f :: * -> *) a.
(Applicative f, Monoid a) =>
[a] -> f (a, [a])
go' where
            go' :: [a] -> f (a, [a])
go' []     = (a, [a]) -> f (a, [a])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
forall a. Monoid a => a
mempty, [])
            go' (a
y:[a]
ys) = (a, [a]) -> f (a, [a])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
y,      [a]
ys)
        go :: Int -> a -> a -> StateT [a] m ()
go Int
i a
l a
r | a
l a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
r = StateT [a] m a
pop StateT [a] m a -> (a -> StateT [a] m ()) -> StateT [a] m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m () -> StateT [a] m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateT [a] m ()) -> (a -> m ()) -> a -> StateT [a] m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (marra Int a -> Int -> a -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray marra Int a
aa Int
i (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$!)
        go Int
i a
l a
r = do
            let m :: a
m = (a
la -> a -> a
forall a. Num a => a -> a -> a
+a
r) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2
            Int -> a -> a -> StateT [a] m ()
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) a
l a
m
            Int -> a -> a -> StateT [a] m ()
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (a
ma -> a -> a
forall a. Num a => a -> a -> a
+a
1) a
r
            m () -> StateT [a] m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (marra Int a -> Int -> m ()
forall a (marra :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marra a m) =>
marra Int a -> Int -> m ()
setLSNM marra Int a
aa Int
i)
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ StateT [a] m () -> [a] -> m ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (Int -> Int -> Int -> StateT [a] m ()
forall a. Integral a => Int -> a -> a -> StateT [a] m ()
go Int
1 Int
l0 Int
r0) [a]
xs
    LazySegTreeMut marru marra u a
-> m (LazySegTreeMut marru marra u a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LazySegTreeMut marru marra u a
 -> m (LazySegTreeMut marru marra u a))
-> LazySegTreeMut marru marra u a
-> m (LazySegTreeMut marru marra u a)
forall a b. (a -> b) -> a -> b
$! Int
-> Int
-> marru Int u
-> marra Int a
-> LazySegTreeMut marru marra u a
forall (marru :: * -> * -> *) (marra :: * -> * -> *) u a.
Int
-> Int
-> marru Int u
-> marra Int a
-> LazySegTreeMut marru marra u a
LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa

applyLSNM :: (Action u a, MArray marru u m, MArray marra a m)
          => marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
applyLSNM :: marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
applyLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r u
u
    | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r    = marra 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' marra Int a
aa Int
i (a -> u -> a
forall u a. Action u a => a -> u -> a
`act` u
u)
    | Bool
otherwise = marra 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' marra Int a
aa Int
i (a -> u -> a
forall u a. Action u a => a -> u -> a
`act` u
u) m () -> m () -> m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> marru Int u -> Int -> (u -> u) -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' marru Int u
ua Int
i (u -> u -> u
forall a. Semigroup a => a -> a -> a
<> u
u)
{-# INLINE applyLSNM #-}

pushLSNM :: (Action u a, MArray marru u m, MArray marra a m)
         => marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM :: marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r = do
    u
u <- marru Int u -> Int -> m u
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marru Int u
ua Int
i
    marru Int u -> Int -> u -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray marru Int u
ua Int
i (u -> m ()) -> u -> m ()
forall a b. (a -> b) -> a -> b
$! u
forall a. Monoid a => a
mempty
    let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
    marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
applyLSNM marru Int u
ua marra Int a
aa (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m u
u
    marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
applyLSNM marru Int u
ua marra Int a
aa (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r u
u
{-# INLINE pushLSNM #-}

-- | Adjusts the element at index i. O(log n).
adjustLSTM :: (Action u a, MArray marru u m, MArray marra a m)
           => LazySegTreeMut marru marra u a -> Int -> (a -> a) -> m ()
adjustLSTM :: LazySegTreeMut marru marra u a -> Int -> (a -> a) -> m ()
adjustLSTM (LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa) Int
qi a -> a
f
    | Int
qi Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l0 Bool -> Bool -> Bool
|| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
qi = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"adjustLSTM: outside range"
    | Bool
otherwise          = Int -> Int -> Int -> m ()
forall (f :: * -> *).
(MArray marra a f, MArray marru u f) =>
Int -> Int -> Int -> f ()
go Int
1 Int
l0 Int
r0
  where
    go :: Int -> Int -> Int -> f ()
go Int
i Int
l Int
r
        | 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 = () -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r           = marra Int a -> Int -> (a -> a) -> f ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' marra Int a
aa Int
i a -> a
f
        | Bool
otherwise = do
            marru Int u -> marra Int a -> Int -> Int -> Int -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r
            let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            Int -> Int -> Int -> f ()
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m
            Int -> Int -> Int -> f ()
go (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r
            marra Int a -> Int -> f ()
forall a (marra :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marra a m) =>
marra Int a -> Int -> m ()
setLSNM marra Int a
aa Int
i

-- | Applies an update on elements in the range (ql, qr). O(log n).
updateRangeLSTM :: (Action u a, MArray marru u m, MArray marra a m)
                => LazySegTreeMut marru marra u a -> Int -> Int -> u -> m ()
updateRangeLSTM :: LazySegTreeMut marru marra u a -> Int -> Int -> u -> m ()
updateRangeLSTM (LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa) Int
ql Int
qr u
qu
    | Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1        = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"updateRangeLSTM: bad range"
    | Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l0 Bool -> Bool -> Bool
|| Int
r0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
qr = [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"updateRangeLSTM: outside range"
    | Bool
otherwise          = Int -> Int -> Int -> m ()
forall (f :: * -> *).
(MArray marru u f, MArray marra a f) =>
Int -> Int -> Int -> f ()
go Int
1 Int
l0 Int
r0
  where
    go :: Int -> Int -> Int -> f ()
go Int
i Int
l Int
r
        | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql Bool -> Bool -> Bool
|| Int
qr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l   = () -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | 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 = marru Int u -> marra Int a -> Int -> Int -> Int -> u -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> u -> m ()
applyLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r u
qu
        | Bool
otherwise = do
            marru Int u -> marra Int a -> Int -> Int -> Int -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r
            let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            Int -> Int -> Int -> f ()
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m
            Int -> Int -> Int -> f ()
go (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r
            marra Int a -> Int -> f ()
forall a (marra :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marra a m) =>
marra Int a -> Int -> m ()
setLSNM marra Int a
aa Int
i

-- | Folds the elements in the range (ql, qr). Elements outside (l, r) are considered to be mempty.
-- O(log n).
foldRangeLSTM :: (Action u a, MArray marru u m, MArray marra a m)
              => LazySegTreeMut marru marra u a -> Int -> Int -> m a
foldRangeLSTM :: LazySegTreeMut marru marra u a -> Int -> Int -> m a
foldRangeLSTM (LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa) Int
ql Int
qr
    | Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = [Char] -> m a
forall a. HasCallStack => [Char] -> a
error [Char]
"foldRangeLSTM: bad range"
    | Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r0     = a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
    | Bool
otherwise   = Int -> Int -> Int -> a -> m a
forall (f :: * -> *).
(MArray marra a f, MArray marru u f) =>
Int -> Int -> Int -> a -> f a
go Int
1 Int
l0 Int
r0 a
forall a. Monoid a => a
mempty
  where
    go :: Int -> Int -> Int -> a -> f a
go Int
i Int
l Int
r a
acc
        | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql Bool -> Bool -> Bool
|| Int
qr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l   = a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
        | 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 = (a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<>) (a -> a) -> f a -> f a
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> marra Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marra Int a
aa Int
i
        | Bool
otherwise = do
            marru Int u -> marra Int a -> Int -> Int -> Int -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r
            let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            Int -> Int -> Int -> a -> f a
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m a
acc 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
go (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r

-- | 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).
binSearchLSTM :: (Action u a, MArray marru u m, MArray marra a m)
              => LazySegTreeMut marru marra u a -> Int -> Int -> (a -> Bool) -> m (Maybe (Int, a))
binSearchLSTM :: LazySegTreeMut marru marra u a
-> Int -> Int -> (a -> Bool) -> m (Maybe (Int, a))
binSearchLSTM (LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa) Int
ql Int
qr a -> Bool
p
    | Int
ql Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
qr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = [Char] -> m (Maybe (Int, a))
forall a. HasCallStack => [Char] -> a
error [Char]
"binSearchLSTM: bad range"
    | Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r0     = 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   = (a -> Maybe (Int, a))
-> ((Int, a) -> Maybe (Int, a))
-> Either a (Int, a)
-> Maybe (Int, a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (Int, a) -> a -> Maybe (Int, a)
forall a b. a -> b -> a
const Maybe (Int, a)
forall a. Maybe a
Nothing) (Int, a) -> Maybe (Int, a)
forall a. a -> Maybe a
Just (Either a (Int, a) -> Maybe (Int, a))
-> m (Either a (Int, a)) -> m (Maybe (Int, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Int -> a -> m (Either a (Int, a))
forall (f :: * -> *).
(MArray marra a f, MArray marru u f) =>
Int -> Int -> Int -> a -> f (Either a (Int, a))
go Int
1 Int
l0 Int
r0 a
forall a. Monoid a => a
mempty
  where
    go :: Int -> Int -> Int -> a -> f (Either a (Int, a))
go Int
i Int
l Int
r a
acc
        | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ql Bool -> Bool -> Bool
|| Int
qr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l = Either a (Int, a) -> f (Either a (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Either a (Int, a)
forall a b. a -> Either a b
Left a
acc)
        | 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 = do
            a
a <- marra Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marra Int a
aa Int
i
            let acc' :: a
acc' = a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
a
            case () of
                ()
_ | Bool -> Bool
not (a -> Bool
p a
acc') -> Either a (Int, a) -> f (Either a (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Either a (Int, a)
forall a b. a -> Either a b
Left a
acc')
                  | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r       -> Either a (Int, a) -> f (Either a (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int, a) -> Either a (Int, a)
forall a b. b -> Either a b
Right (Int
l, a
acc'))
                  | Bool
otherwise    -> Int -> Int -> Int -> a -> f (Either a (Int, a))
goLR Int
i Int
l Int
r a
acc
        | Bool
otherwise = Int -> Int -> Int -> a -> f (Either a (Int, a))
goLR Int
i Int
l Int
r a
acc
    goLR :: Int -> Int -> Int -> a -> f (Either a (Int, a))
goLR Int
i Int
l Int
r a
acc = do
        marru Int u -> marra Int a -> Int -> Int -> Int -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r
        let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        Either a (Int, a)
lres <- Int -> Int -> Int -> a -> f (Either a (Int, a))
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m a
acc
        case Either a (Int, a)
lres of
            Left a
acc' -> Int -> Int -> Int -> a -> f (Either a (Int, a))
go (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r a
acc'
            Either a (Int, a)
_         -> Either a (Int, a) -> f (Either a (Int, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Either a (Int, a)
lres

-- | Right fold over the elements of the segment tree. O(n).
foldrLSTM :: (Action u a, MArray marru u m, MArray marra a m)
          => LazySegTreeMut marru marra u a -> (a -> b -> b) -> b -> m b
foldrLSTM :: LazySegTreeMut marru marra u a -> (a -> b -> b) -> b -> m b
foldrLSTM (LSTM Int
l0 Int
r0 marru Int u
ua marra Int a
aa) a -> b -> b
f b
z0
    | Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
r0   = b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
z0
    | Bool
otherwise = Int -> Int -> Int -> b -> m b
forall (f :: * -> *).
(MArray marra a f, MArray marru u f) =>
Int -> Int -> Int -> b -> f b
go Int
1 Int
l0 Int
r0 b
z0
  where
    go :: Int -> Int -> Int -> b -> f b
go Int
i Int
l Int
r b
z
        | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r = (a -> b -> b
`f` b
z) (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> marra Int a -> Int -> f a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray marra Int a
aa Int
i
        | Bool
otherwise = do
            marru Int u -> marra Int a -> Int -> Int -> Int -> f ()
forall u a (marru :: * -> * -> *) (m :: * -> *)
       (marra :: * -> * -> *).
(Action u a, MArray marru u m, MArray marra a m) =>
marru Int u -> marra Int a -> Int -> Int -> Int -> m ()
pushLSNM marru Int u
ua marra Int a
aa Int
i Int
l Int
r
            let m :: Int
m = (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
r) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            Int -> Int -> Int -> b -> f b
go (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
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r b
z f b -> (b -> f b) -> f b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> Int -> b -> f b
go (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int
l Int
m

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

-- Allows specialization across modules
{-# INLINABLE fromListLSTM #-}
{-# INLINABLE adjustLSTM #-}
{-# INLINABLE updateRangeLSTM #-}
{-# INLINABLE foldRangeLSTM #-}
{-# INLINABLE binSearchLSTM #-}