{-# LANGUAGE RecordWildCards #-}
module AtCoder.Extra.Tree.TreeMonoid
(
TreeMonoid,
Vertex,
VertexHld,
Commutativity (..),
fromVerts,
fromEdges,
prod,
prodSubtree,
read,
write,
exchange,
modify,
modifyM,
)
where
import AtCoder.Extra.Tree.Hld qualified as Hld
import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.SegTree qualified as ST
import Control.Monad
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Monoid (Dual (..))
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)
type Vertex = Int
type VertexHld = Vertex
data TreeMonoid a s = TreeMonoid
{
forall a s. TreeMonoid a s -> Hld
hldTM :: !Hld.Hld,
forall a s. TreeMonoid a s -> Commutativity
commuteTM :: !Commutativity,
forall a s. TreeMonoid a s -> WeightPolicy
weightPolicyTM :: !Hld.WeightPolicy,
forall a s. TreeMonoid a s -> SegTree s a
segFTM :: !(ST.SegTree s a),
forall a s. TreeMonoid a s -> SegTree s (Dual a)
segBTM :: !(ST.SegTree s (Dual a))
}
data Commutativity
=
Commute
|
NonCommute
deriving
(
Commutativity -> Commutativity -> Bool
(Commutativity -> Commutativity -> Bool)
-> (Commutativity -> Commutativity -> Bool) -> Eq Commutativity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Commutativity -> Commutativity -> Bool
== :: Commutativity -> Commutativity -> Bool
$c/= :: Commutativity -> Commutativity -> Bool
/= :: Commutativity -> Commutativity -> Bool
Eq,
Int -> Commutativity -> ShowS
[Commutativity] -> ShowS
Commutativity -> String
(Int -> Commutativity -> ShowS)
-> (Commutativity -> String)
-> ([Commutativity] -> ShowS)
-> Show Commutativity
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Commutativity -> ShowS
showsPrec :: Int -> Commutativity -> ShowS
$cshow :: Commutativity -> String
show :: Commutativity -> String
$cshowList :: [Commutativity] -> ShowS
showList :: [Commutativity] -> ShowS
Show
)
buildImpl ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
Hld.Hld ->
Commutativity ->
Hld.WeightPolicy ->
VU.Vector a ->
m (TreeMonoid a (PrimState m))
buildImpl :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> m (TreeMonoid a (PrimState m))
buildImpl Hld
hldTM Commutativity
commuteTM WeightPolicy
weightPolicyTM Vector a
weights = do
SegTree (PrimState m) a
segFTM <- Vector a -> m (SegTree (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build Vector a
weights
SegTree (PrimState m) (Dual a)
segBTM <-
case Commutativity
commuteTM of
Commutativity
Commute -> Vector (Dual a) -> m (SegTree (PrimState m) (Dual a))
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build Vector (Dual a)
forall a. Unbox a => Vector a
VU.empty
Commutativity
NonCommute -> Vector (Dual a) -> m (SegTree (PrimState m) (Dual a))
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vector a -> m (SegTree (PrimState m) a)
ST.build (Vector (Dual a) -> m (SegTree (PrimState m) (Dual a)))
-> Vector (Dual a) -> m (SegTree (PrimState m) (Dual a))
forall a b. (a -> b) -> a -> b
$ (a -> Dual a) -> Vector a -> Vector (Dual a)
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map a -> Dual a
forall a. a -> Dual a
Dual Vector a
weights
TreeMonoid a (PrimState m) -> m (TreeMonoid a (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..}
fromVerts ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
Hld.Hld ->
Commutativity ->
VU.Vector a ->
m (TreeMonoid a (PrimState m))
fromVerts :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld -> Commutativity -> Vector a -> m (TreeMonoid a (PrimState m))
fromVerts hld :: Hld
hld@Hld.Hld {Vector Int
indexHld :: Vector Int
indexHld :: Hld -> Vector Int
indexHld} Commutativity
commuteTM Vector a
xs_ = do
let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld 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 -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Tree.TreeMonoid.fromVerts: vertex number mismatch (`" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"` and `" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs_) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`)"
let !xs :: Vector a
xs = (forall s. ST s (MVector s a)) -> Vector a
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s a)) -> Vector a)
-> (forall s. ST s (MVector s a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
MVector s a
vec <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> ST s (MVector (PrimState (ST s)) a))
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs_
Vector a -> (Int -> a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (Int -> a -> m b) -> m ()
VU.iforM_ Vector a
xs_ ((Int -> a -> ST s ()) -> ST s ())
-> (Int -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i a
x -> do
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
vec (Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i) a
x
MVector s a -> ST s (MVector s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s a
vec
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> m (TreeMonoid a (PrimState m))
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> m (TreeMonoid a (PrimState m))
buildImpl Hld
hld Commutativity
commuteTM WeightPolicy
Hld.WeightsAreOnVertices Vector a
xs
fromEdges ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
Hld.Hld ->
Commutativity ->
VU.Vector (Vertex, Vertex, a) ->
m (TreeMonoid a (PrimState m))
fromEdges :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> Vector (Int, Int, a)
-> m (TreeMonoid a (PrimState m))
fromEdges hld :: Hld
hld@Hld.Hld {Vector Int
indexHld :: Hld -> Vector Int
indexHld :: Vector Int
indexHld} Commutativity
commuteTM Vector (Int, Int, a)
edges = do
let !xs :: Vector a
xs = (forall s. ST s (MVector s a)) -> Vector a
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s a)) -> Vector a)
-> (forall s. ST s (MVector s a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
MVector s a
vec <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> ST s (MVector (PrimState (ST s)) a))
-> Int -> ST s (MVector (PrimState (ST s)) a)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
indexHld
Vector (Int, Int, a) -> ((Int, Int, a) -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Int, Int, a)
edges (((Int, Int, a) -> ST s ()) -> ST s ())
-> ((Int, Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(!Int
u, !Int
v, !a
w) -> do
let u' :: Int
u' = Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
u
let v' :: Int
v' = Vector Int
indexHld Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
v
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
vec (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
u' Int
v') a
w
MVector s a -> ST s (MVector s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s a
vec
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> m (TreeMonoid a (PrimState m))
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Hld
-> Commutativity
-> WeightPolicy
-> Vector a
-> m (TreeMonoid a (PrimState m))
buildImpl Hld
hld Commutativity
commuteTM WeightPolicy
Hld.WeightsAreOnEdges Vector a
xs
prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> Vertex -> m a
prod :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> Int -> m a
prod TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} Int
u Int
v = do
case Commutativity
commuteTM of
Commutativity
Commute -> WeightPolicy
-> Hld
-> (Int -> Int -> m a)
-> (Int -> Int -> m a)
-> Int
-> Int
-> m a
forall mono (m :: * -> *).
(HasCallStack, Monoid mono, Monad m) =>
WeightPolicy
-> Hld
-> (Int -> Int -> m mono)
-> (Int -> Int -> m mono)
-> Int
-> Int
-> m mono
Hld.prod WeightPolicy
weightPolicyTM Hld
hldTM (SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) a
segFTM) (SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) a
segFTM) Int
u Int
v
Commutativity
NonCommute -> WeightPolicy
-> Hld
-> (Int -> Int -> m a)
-> (Int -> Int -> m a)
-> Int
-> Int
-> m a
forall mono (m :: * -> *).
(HasCallStack, Monoid mono, Monad m) =>
WeightPolicy
-> Hld
-> (Int -> Int -> m mono)
-> (Int -> Int -> m mono)
-> Int
-> Int
-> m mono
Hld.prod WeightPolicy
weightPolicyTM Hld
hldTM (SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) a
segFTM) (\Int
l Int
r -> Dual a -> a
forall a. Dual a -> a
getDual (Dual a -> a) -> m (Dual a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SegTree (PrimState m) (Dual a) -> Int -> Int -> m (Dual a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) (Dual a)
segBTM Int
l Int
r) Int
u Int
v
prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a
prodSubtree :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> m a
prodSubtree TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} Int
subtreeRoot = do
let (!Int
l, !Int
r) = HasCallStack => Hld -> Int -> (Int, Int)
Hld -> Int -> (Int, Int)
Hld.subtreeSegmentInclusive Hld
hldTM Int
subtreeRoot
case WeightPolicy
weightPolicyTM of
WeightPolicy
Hld.WeightsAreOnVertices -> SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) a
segFTM Int
l (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
WeightPolicy
Hld.WeightsAreOnEdges -> do
if Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r
then a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
else SegTree (PrimState m) a -> Int -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> Int -> m a
ST.prod SegTree (PrimState m) a
segFTM (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> m a
read TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} Int
i_ = do
let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
SegTree (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> m a
ST.read SegTree (PrimState m) a
segFTM Int
i
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> a -> m ()
write TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} Int
i_ a
x = do
let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
SegTree (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree (PrimState m) a
segFTM Int
i a
x
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
SegTree (PrimState m) (Dual a) -> Int -> Dual a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree (PrimState m) (Dual a)
segBTM Int
i (Dual a -> m ()) -> Dual a -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Dual a
forall a. a -> Dual a
Dual a
x
exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m a
exchange :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> Int -> a -> m a
exchange TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} Int
i_ a
x = do
let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
!a
res <- SegTree (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m a
ST.exchange SegTree (PrimState m) a
segFTM Int
i a
x
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
SegTree (PrimState m) (Dual a) -> Int -> Dual a -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> Int -> a -> m ()
ST.write SegTree (PrimState m) (Dual a)
segBTM Int
i (Dual a -> m ()) -> Dual a -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Dual a
forall a. a -> Dual a
Dual a
x
a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()
modify TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} a -> a
f Int
i_ = do
let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
ST.modify SegTree (PrimState m) a
segFTM a -> a
f Int
i
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
SegTree (PrimState m) (Dual a) -> (Dual a -> Dual a) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> a) -> Int -> m ()
ST.modify SegTree (PrimState m) (Dual a)
segBTM (a -> Dual a
forall a. a -> Dual a
Dual (a -> Dual a) -> (Dual a -> a) -> Dual a -> Dual a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f (a -> a) -> (Dual a -> a) -> Dual a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual a -> a
forall a. Dual a -> a
getDual) Int
i
modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
TreeMonoid a (PrimState m) -> (a -> m a) -> Int -> m ()
modifyM TreeMonoid {WeightPolicy
Hld
SegTree (PrimState m) a
SegTree (PrimState m) (Dual a)
Commutativity
hldTM :: forall a s. TreeMonoid a s -> Hld
commuteTM :: forall a s. TreeMonoid a s -> Commutativity
weightPolicyTM :: forall a s. TreeMonoid a s -> WeightPolicy
segFTM :: forall a s. TreeMonoid a s -> SegTree s a
segBTM :: forall a s. TreeMonoid a s -> SegTree s (Dual a)
hldTM :: Hld
commuteTM :: Commutativity
weightPolicyTM :: WeightPolicy
segFTM :: SegTree (PrimState m) a
segBTM :: SegTree (PrimState m) (Dual a)
..} a -> m a
f Int
i_ = do
let !i :: Int
i = Hld -> Vector Int
Hld.indexHld Hld
hldTM Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i_
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
ST.modifyM SegTree (PrimState m) a
segFTM a -> m a
f Int
i
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Commutativity
commuteTM Commutativity -> Commutativity -> Bool
forall a. Eq a => a -> a -> Bool
== Commutativity
NonCommute) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
SegTree (PrimState m) (Dual a)
-> (Dual a -> m (Dual a)) -> Int -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
SegTree (PrimState m) a -> (a -> m a) -> Int -> m ()
ST.modifyM SegTree (PrimState m) (Dual a)
segBTM ((a -> Dual a
forall a. a -> Dual a
Dual <$>) (m a -> m (Dual a)) -> (Dual a -> m a) -> Dual a -> m (Dual a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m a
f (a -> m a) -> (Dual a -> a) -> Dual a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual a -> a
forall a. Dual a -> a
getDual) Int
i