{- |
 = Segment Tree
 == Reference
   * <https://codeforces.com/blog/entry/18051>
-}
module Data.SegTree.Primal where

import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Function
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import Unsafe.Coerce

newtype SegTree mv s a = SegTree {forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree :: mv s a}

{- |
>>> import Data.Semigroup (Min)
>>> import qualified Data.Vector.Unboxed.Mutable as  UM
>>> newSegTree @(Min Int) @UM.MVector 123
-}
newSegTree ::
  (Monoid a, GM.MVector mv a, PrimMonad m) =>
  Int ->
  m (SegTree mv (PrimState m) a)
newSegTree :: forall a (mv :: * -> * -> *) (m :: * -> *).
(Monoid a, MVector mv a, PrimMonad m) =>
Int -> m (SegTree mv (PrimState m) a)
newSegTree Int
n = mv (PrimState m) a -> SegTree mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
mv s a -> SegTree mv s a
SegTree (mv (PrimState m) a -> SegTree mv (PrimState m) a)
-> m (mv (PrimState m) a) -> m (SegTree mv (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> m (mv (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
GM.replicate (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
extendToPowerOfTwo Int
n) a
forall a. Monoid a => a
mempty

-- | /O(n)/
buildSegTree ::
  (Monoid a, PrimMonad m, G.Vector v a) =>
  v a ->
  m (SegTree (G.Mutable v) (PrimState m) a)
buildSegTree :: forall a (m :: * -> *) (v :: * -> *).
(Monoid a, PrimMonad m, Vector v a) =>
v a -> m (SegTree (Mutable v) (PrimState m) a)
buildSegTree v a
vec = do
  let n :: Int
n = Int -> Int
extendToPowerOfTwo (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec
  Mutable v (PrimState m) a
tree <- Int -> a -> m (Mutable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
GM.replicate (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) a
forall a. Monoid a => a
mempty
  Mutable v (PrimState m) a -> v a -> m ()
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> v a -> m ()
G.unsafeCopy (Int
-> Int -> Mutable v (PrimState m) a -> Mutable v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
GM.unsafeSlice Int
n (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec) Mutable v (PrimState m) a
tree) v a
vec
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i -> Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    a -> a -> a
forall a. Monoid a => a -> a -> a
mappend
      (a -> a -> a) -> m a -> m (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead Mutable v (PrimState m) a
tree (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
1)
      m (a -> a) -> m a -> m a
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mutable v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead Mutable v (PrimState m) a
tree ((Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
1)
      m a -> (a -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite Mutable v (PrimState m) a
tree Int
i
    Int -> m ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  SegTree (Mutable v) (PrimState m) a
-> m (SegTree (Mutable v) (PrimState m) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SegTree (Mutable v) (PrimState m) a
 -> m (SegTree (Mutable v) (PrimState m) a))
-> SegTree (Mutable v) (PrimState m) a
-> m (SegTree (Mutable v) (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState m) a -> SegTree (Mutable v) (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
mv s a -> SegTree mv s a
SegTree Mutable v (PrimState m) a
tree

-- | /O(1)/
readSegTree ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  Int ->
  m a
readSegTree :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> m a
readSegTree SegTree mv (PrimState m) a
segtree Int
k = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
{-# INLINE readSegTree #-}

pullSegTree ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  Int ->
  m ()
pullSegTree :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> m ()
pullSegTree SegTree mv (PrimState m) a
seg Int
k = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
seg
  a -> a -> a
forall a. Monoid a => a -> a -> a
mappend
    (a -> a -> a) -> m a -> m (a -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
k Int
1)
    m (a -> a) -> m a -> m a
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
k Int
1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
1)
    m a -> (a -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= mv (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite mv (PrimState m) a
tree Int
k
{-# INLINE pullSegTree #-}

-- | /O(log n)/
writeSegTree ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  Int ->
  a ->
  m ()
writeSegTree :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> a -> m ()
writeSegTree SegTree mv (PrimState m) a
segtree Int
k a
v = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  mv (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite mv (PrimState m) a
tree (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) a
v
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
1) (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i ->
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      SegTree mv (PrimState m) a -> Int -> m ()
forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> m ()
pullSegTree SegTree mv (PrimState m) a
segtree Int
i
      Int -> m ()
loop (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
i Int
1
{-# INLINE writeSegTree #-}

-- | /O(log n)/
modifySegTree ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  (a -> a) ->
  Int ->
  m ()
modifySegTree :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> (a -> a) -> Int -> m ()
modifySegTree SegTree mv (PrimState m) a
segtree a -> a
f Int
k = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  mv (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
GM.unsafeModify mv (PrimState m) a
tree a -> a
f (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix (Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int
1) (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i ->
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      SegTree mv (PrimState m) a -> Int -> m ()
forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> m ()
pullSegTree SegTree mv (PrimState m) a
segtree Int
i
      Int -> m ()
loop (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
i Int
1
{-# INLINE modifySegTree #-}

{- |
mconcat[a[l],...,a[r-1]]

/O(log n)/
-}
mappendFromTo ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  Int ->
  Int ->
  m a
mappendFromTo :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> Int -> m a
mappendFromTo SegTree mv (PrimState m) a
segtree Int
l0 Int
r0 = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  ((a -> a -> Int -> Int -> m a) -> a -> a -> Int -> Int -> m a)
-> a -> a -> Int -> Int -> m a
forall a. (a -> a) -> a
fix
    ( \a -> a -> Int -> Int -> m a
loop !a
accL !a
accR !Int
l !Int
r -> do
        if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r
          then do
            a
accL' <-
              if Int
l Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                then a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
accL (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree Int
l
                else a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
accL
            a
accR' <-
              if Int
r Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                then (a -> a -> a) -> a -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
accR (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                else a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
accR
            a -> a -> Int -> Int -> m a
loop
              a
accL'
              a
accR'
              ((Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
1)
              (Int
r Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
1)
          else a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (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
    )
    a
forall a. Monoid a => a
mempty
    a
forall a. Monoid a => a
mempty
    (Int
l0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
    (Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
{-# INLINE mappendFromTo #-}

{- |
mconcat[a[0],...,a[k-1]]

/O(log n)/
-}
mappendTo ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  Int ->
  m a
mappendTo :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> m a
mappendTo SegTree mv (PrimState m) a
segtree = SegTree mv (PrimState m) a -> Int -> Int -> m a
forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> Int -> m a
mappendFromTo SegTree mv (PrimState m) a
segtree Int
0
{-# INLINE mappendTo #-}

{- |
mconcat[a[0],...,a[n-1]]

/O(1)/
-}
mappendAll ::
  (PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  m a
mappendAll :: forall (m :: * -> *) (mv :: * -> * -> *) a.
(PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> m a
mappendAll SegTree mv (PrimState m) a
segtree = mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead (SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree) Int
1
{-# INLINE mappendAll #-}

{- | max r s.t. f (mappendFromTo seg l r) == True

>>> import Data.Semigroup (Min)
>>> import qualified Data.Vector.Unboxed.Mutable as  UM
>>> seg <- newSegTree @(Min Int) @UM.MVector 10
>>> upperBoundFrom seg 0 (const True)
16
-}
upperBoundFrom ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  -- | left
  Int ->
  -- | predicate s.t. f memepty == True, monotone
  (a -> Bool) ->
  m Int
upperBoundFrom :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> (a -> Bool) -> m Int
upperBoundFrom SegTree mv (PrimState m) a
segtree Int
l a -> Bool
p = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let !n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  Maybe (a, Int)
violationNode <-
    ((a -> Int -> m (Maybe (a, Int)))
 -> a -> Int -> m (Maybe (a, Int)))
-> a -> Int -> m (Maybe (a, Int))
forall a. (a -> a) -> a
fix
      ( \a -> Int -> m (Maybe (a, Int))
loopUp !a
acc !Int
cur -> do
          let rightParent :: Int
rightParent = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
cur (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
cur)
          !a
acc' <- (a
acc <>) (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree Int
rightParent
          if a -> Bool
p a
acc'
            then do
              let !cur' :: Int
cur' = Int
rightParent Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
              if Int
cur' Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Num a => a -> a
negate Int
cur' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
cur'
                then a -> Int -> m (Maybe (a, Int))
loopUp a
acc' Int
cur'
                else Maybe (a, Int) -> m (Maybe (a, Int))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (a, Int)
forall a. Maybe a
Nothing
            else Maybe (a, Int) -> m (Maybe (a, Int))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (a, Int) -> m (Maybe (a, Int)))
-> Maybe (a, Int) -> m (Maybe (a, Int))
forall a b. (a -> b) -> a -> b
$ (a, Int) -> Maybe (a, Int)
forall a. a -> Maybe a
Just (a
acc, Int
rightParent)
      )
      a
forall a. Monoid a => a
mempty
      (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
  case Maybe (a, Int)
violationNode of
    Maybe (a, Int)
Nothing -> Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n
    Just (!a
acc0, !Int
cur0) -> do
      ((a -> Int -> m Int) -> a -> Int -> m Int) -> a -> Int -> m Int
forall a. (a -> a) -> a
fix
        ( \a -> Int -> m Int
loopDown !a
acc !Int
cur -> do
            if Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
              then do
                let !leftChild :: Int
leftChild = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cur
                !a
acc' <- (a
acc <>) (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree Int
leftChild
                if a -> Bool
p a
acc'
                  then a -> Int -> m Int
loopDown a
acc' (Int
leftChild Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                  else a -> Int -> m Int
loopDown a
acc Int
leftChild
              else Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Int
cur Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n
        )
        a
acc0
        Int
cur0
{-# INLINE upperBoundFrom #-}

-- | min l s.t. f (mappendFromTo seg l r) == True
lowerBoundTo ::
  (Monoid a, PrimMonad m, GM.MVector mv a) =>
  SegTree mv (PrimState m) a ->
  -- | right
  Int ->
  -- | predicate s.t. f memepty == True, monotone
  (a -> Bool) ->
  m Int
lowerBoundTo :: forall a (m :: * -> *) (mv :: * -> * -> *).
(Monoid a, PrimMonad m, MVector mv a) =>
SegTree mv (PrimState m) a -> Int -> (a -> Bool) -> m Int
lowerBoundTo SegTree mv (PrimState m) a
segtree Int
r a -> Bool
p = do
  let tree :: mv (PrimState m) a
tree = SegTree mv (PrimState m) a -> mv (PrimState m) a
forall {k} {k} (mv :: k -> k -> *) (s :: k) (a :: k).
SegTree mv s a -> mv s a
getSegTree SegTree mv (PrimState m) a
segtree
  let !n :: Int
n = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR (mv (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length mv (PrimState m) a
tree) Int
1
  Maybe (a, Int)
violationNode <-
    ((a -> Int -> m (Maybe (a, Int)))
 -> a -> Int -> m (Maybe (a, Int)))
-> a -> Int -> m (Maybe (a, Int))
forall a. (a -> a) -> a
fix
      ( \a -> Int -> m (Maybe (a, Int))
loopUp !a
acc !Int
cur -> do
          let leftParent :: Int
leftParent =
                case Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
cur (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Int -> Int
forall a. Bits a => a -> a
complement Int
cur)) of
                  Int
0 -> Int
1 -- cur: 2 ^ n
                  Int
v -> Int
v
          !a
acc' <- (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc) (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree Int
leftParent
          if a -> Bool
p a
acc'
            then do
              if Int
leftParent Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Num a => a -> a
negate Int
leftParent Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
leftParent
                then a -> Int -> m (Maybe (a, Int))
loopUp a
acc' (Int
leftParent Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                else Maybe (a, Int) -> m (Maybe (a, Int))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (a, Int)
forall a. Maybe a
Nothing
            else Maybe (a, Int) -> m (Maybe (a, Int))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (a, Int) -> m (Maybe (a, Int)))
-> Maybe (a, Int) -> m (Maybe (a, Int))
forall a b. (a -> b) -> a -> b
$ (a, Int) -> Maybe (a, Int)
forall a. a -> Maybe a
Just (a
acc, Int
leftParent)
      )
      a
forall a. Monoid a => a
mempty
      (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
  case Maybe (a, Int)
violationNode of
    Maybe (a, Int)
Nothing -> Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0
    Just (!a
acc0, !Int
cur0) ->
      ((a -> Int -> m Int) -> a -> Int -> m Int) -> a -> Int -> m Int
forall a. (a -> a) -> a
fix
        ( \a -> Int -> m Int
loopDown !a
acc !Int
cur -> do
            if Int
cur Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
              then do
                let !rightChild :: Int
rightChild = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cur Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                !a
acc' <- (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc) (a -> a) -> m a -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> mv (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead mv (PrimState m) a
tree Int
rightChild
                if a -> Bool
p a
acc'
                  then a -> Int -> m Int
loopDown a
acc' (Int
rightChild Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                  else a -> Int -> m Int
loopDown a
acc Int
rightChild
              else Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$! Int
cur Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n
        )
        a
acc0
        Int
cur0
{-# INLINE lowerBoundTo #-}

{- |
>>> extendToPowerOfTwo 0
1
-}
extendToPowerOfTwo :: Int -> Int
extendToPowerOfTwo :: Int -> Int
extendToPowerOfTwo Int
x
  | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 =
      forall a b. a -> b
unsafeCoerce @Word @Int (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$
        Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word -> Word
forall a. Bits a => a -> a
complement Word
forall a. Bits a => a
zeroBits) (Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1
  | Bool
otherwise = Int
1