{-# LANGUAGE RecordWildCards #-}
module AtCoder.Extra.SegTree2d.Dense
(
DenseSegTree2d (..),
new,
build,
build',
read,
readMaybe,
write,
modify,
modifyM,
prod,
allProd,
)
where
import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bits
import Data.Foldable (for_)
import Data.Maybe (fromMaybe)
import Data.Vector qualified as V
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)
data DenseSegTree2d s a = DenseSegTree2d
{
forall s a. DenseSegTree2d s a -> Int
hDst :: {-# UNPACK #-} !Int,
forall s a. DenseSegTree2d s a -> Int
wDst :: {-# UNPACK #-} !Int,
forall s a. DenseSegTree2d s a -> MVector s a
dataDst :: !(VUM.MVector s a)
}
{-# INLINEABLE new #-}
new ::
(PrimMonad m, Monoid a, VU.Unbox a) =>
Int ->
Int ->
m (DenseSegTree2d (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Int -> Int -> m (DenseSegTree2d (PrimState m) a)
new Int
wDst Int
hDst = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ do
MVector (PrimState m) a
dataDst <- Int
-> a -> ST (PrimState m) (MVector (PrimState (ST (PrimState m))) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst) a
forall a. Monoid a => a
mempty
DenseSegTree2d (PrimState m) a
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
wDst :: Int
hDst :: Int
dataDst :: MVector (PrimState m) a
..}
{-# INLINE build #-}
build ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
Int ->
Int ->
VU.Vector a ->
m (DenseSegTree2d (PrimState m) a)
build :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Int -> Int -> Vector a -> m (DenseSegTree2d (PrimState m) a)
build Int
w Int
h Vector a
xs = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST (Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a))
-> Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Int
-> (Vector a -> (Vector a, Vector a))
-> Vector a
-> Vector (Vector a)
forall b a. Int -> (b -> (a, b)) -> b -> Vector a
V.unfoldrExactN Int
h (Int -> Vector a -> (Vector a, Vector a)
forall a. Unbox a => Int -> Vector a -> (Vector a, Vector a)
VU.splitAt Int
w) Vector a
xs
where
!()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) String
"AtCoder.Extra.SegTree2d.Dense.build: vector length mismatch"
{-# INLINE build' #-}
build' ::
(HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
V.Vector (VU.Vector a) ->
m (DenseSegTree2d (PrimState m) a)
build' :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector (Vector a) -> m (DenseSegTree2d (PrimState m) a)
build' Vector (Vector a)
xs = ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a))
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
-> m (DenseSegTree2d (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Vector (Vector a)
-> ST (PrimState m) (DenseSegTree2d (PrimState m) a)
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST Vector (Vector a)
xs
{-# INLINE read #-}
read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m a
read :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> m a
read DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y = do
let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkPoint2d String
"AtCoder.Extra.SegTree2d.Dense.read" Int
x Int
y Int
wDst Int
hDst
MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
dataDst (Int -> m a) -> Int -> m a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Int
idx Int
wDst (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst)
{-# INLINE readMaybe #-}
readMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> m (Maybe a)
readMaybe :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> m (Maybe a)
readMaybe DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y
| HasCallStack => Int -> Int -> Int -> Int -> Bool
Int -> Int -> Int -> Int -> Bool
ACIA.testPoint2d Int
x Int
y Int
wDst Int
hDst = do
a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst))
| Bool
otherwise = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> a -> m ()
write seg :: DenseSegTree2d (PrimState m) a
seg@DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x Int
y a
a = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
DenseSegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d (PrimState m) a
DenseSegTree2d (PrimState (ST (PrimState m))) a
seg (a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> (a -> a) -> a -> ST (PrimState m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a -> a
forall a b. a -> b -> a
const a
a) Int
x Int
y
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> (a -> a) -> Int -> Int -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> a) -> Int -> Int -> m ()
modify DenseSegTree2d (PrimState m) a
seg a -> a
f Int
x Int
y = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
DenseSegTree2d (PrimState (ST (PrimState m))) a
-> (a -> ST (PrimState m) a) -> Int -> Int -> ST (PrimState m) ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d (PrimState m) a
DenseSegTree2d (PrimState (ST (PrimState m))) a
seg (a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST (PrimState m) a) -> (a -> a) -> a -> ST (PrimState m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) Int
x Int
y
{-# INLINEABLE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> (a -> m a) -> Int -> Int -> m ()
modifyM DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} a -> m a
f Int
x0_ Int
y0_ = do
let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkPoint2d String
"AtCoder.Extra.SegTree2d.Dense.modifyM" Int
x0_ Int
y0_ Int
wDst Int
hDst
let y0 :: Int
y0 = Int
y0_ Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst
let x0 :: Int
x0 = Int
x0_ Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst
MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Int -> m ()
VGM.modifyM MVector (PrimState m) a
dataDst a -> m a
f (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 Int
x0)
ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ do
let updateCurrentRow :: Int -> f ()
updateCurrentRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
updateCurrentRow Int
x = do
a
xl <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0))
a
xr <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
MVector (PrimState f) a -> Int -> a -> f ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y0 Int
x) (a -> f ()) -> a -> f ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
Int -> f ()
updateCurrentRow (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateCurrentRow (Int
x0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
let updateOtherRow :: Int -> f ()
updateOtherRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
updateOtherRow Int
y = do
let updateRow :: Int -> f ()
updateRow Int
0 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
updateRow Int
x = do
a
xl <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
x)
a
xr <- MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x)
MVector (PrimState f) a -> Int -> a -> f ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState m) a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> f ()) -> a -> f ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
Int -> f ()
updateRow (Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
Int -> f ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateRow Int
x0
Int -> f ()
updateOtherRow (Int
y Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
updateOtherRow (Int
y0 Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
{-# INLINE prod #-}
prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a
prod :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> Int -> Int -> Int -> Int -> m a
prod seg :: DenseSegTree2d (PrimState m) a
seg@DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} Int
x1 Int
x2 Int
y1 Int
y2 = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ do
let !()
_ = HasCallStack => String -> Int -> Int -> Int -> Int -> ()
String -> Int -> Int -> Int -> Int -> ()
ACIA.checkRectShape String
"AtCoder.Extra.SegTree2d.Dense.prodST" Int
x1 Int
x2 Int
y1 Int
y2
DenseSegTree2d (PrimState m) a
-> Int -> Int -> Int -> Int -> ST (PrimState m) a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST DenseSegTree2d (PrimState m) a
seg (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
x1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
wDst Int
x2) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 Int
y1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
hDst Int
y2)
{-# INLINE allProd #-}
allProd :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => DenseSegTree2d (PrimState m) a -> m a
allProd :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
DenseSegTree2d (PrimState m) a -> m a
allProd DenseSegTree2d {Int
MVector (PrimState m) a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector (PrimState m) a
..} = ST (PrimState m) a -> m a
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) a -> m a) -> ST (PrimState m) a -> m a
forall a b. (a -> b) -> a -> b
$ do
a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
forall a. Monoid a => a
mempty (Maybe a -> a) -> ST (PrimState m) (Maybe a) -> ST (PrimState m) a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST (PrimState m))) a
-> Int -> ST (PrimState m) (Maybe a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (Maybe a)
VGM.readMaybe MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
1 Int
1)
{-# INLINE idx #-}
idx :: Int -> Int -> Int -> Int
idx :: Int -> Int -> Int -> Int
idx Int
w Int
y Int
x = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
{-# INLINEABLE buildST #-}
buildST :: (HasCallStack, Monoid a, VU.Unbox a) => V.Vector (VU.Vector a) -> ST s (DenseSegTree2d s a)
buildST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Vector (Vector a) -> ST s (DenseSegTree2d s a)
buildST Vector (Vector a)
vec = do
let hDst :: Int
hDst = Vector (Vector a) -> Int
forall a. Vector a -> Int
V.length Vector (Vector a)
vec
let wDst :: Int
wDst = Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length (Vector (Vector a) -> Vector a
forall a. Vector a -> a
V.head Vector (Vector a)
vec)
MVector s a
dataDst <- Int -> a -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst) a
forall a. Monoid a => a
mempty
Vector (Vector a) -> (Int -> Vector a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b.
Monad m =>
Vector a -> (Int -> a -> m b) -> m ()
V.iforM_ Vector (Vector a)
vec ((Int -> Vector a -> ST s ()) -> ST s ())
-> (Int -> Vector a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y Vector a
vs -> do
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
vs ((Int -> a -> ST s ()) -> ST s ())
-> (Int -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x a
v -> 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
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y) (Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)) a
v
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
hDst .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y -> do
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
a
xl <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0))
a
xr <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
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
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
0] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
y -> do
[Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
0 .. Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
x -> do
a
xl <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
x)
a
xr <- MVector (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState (ST s)) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
x)
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
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
x) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
xl a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
xr
DenseSegTree2d s a -> ST s (DenseSegTree2d s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DenseSegTree2d {Int
MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..}
{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> Int -> ST s a
prodST seg :: DenseSegTree2d s a
seg@DenseSegTree2d {Int
MVector s a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..} Int
x1 Int
x2 Int
y1 Int
y2 = do
a -> Int -> Int -> ST s a
inner a
forall a. Monoid a => a
mempty (Int
y1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst) (Int
y2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
hDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
where
inner :: a -> Int -> Int -> ST s a
inner !a
acc Int
yl Int
yr
| Int
yl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
yr = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
| Bool
otherwise = do
a
acc' <-
if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
yl Int
0
then (a
acc <>) (a -> a) -> ST s a -> ST s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d s a
seg Int
yl Int
x1 Int
x2
else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
a
acc'' <-
if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
yr Int
0
then (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc') (a -> a) -> ST s a -> ST s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d s a
seg Int
yr Int
x1 Int
x2
else a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
a -> Int -> Int -> ST s a
inner a
acc'' ((Int
yl 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
yr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1)
{-# INLINEABLE prodX #-}
prodX :: (HasCallStack, Monoid a, VU.Unbox a) => DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
DenseSegTree2d s a -> Int -> Int -> Int -> ST s a
prodX DenseSegTree2d {Int
MVector s a
hDst :: forall s a. DenseSegTree2d s a -> Int
wDst :: forall s a. DenseSegTree2d s a -> Int
dataDst :: forall s a. DenseSegTree2d s a -> MVector s a
hDst :: Int
wDst :: Int
dataDst :: MVector s a
..} Int
y Int
x1 Int
x2 = do
a -> Int -> Int -> ST s a
forall {f :: * -> *}.
(PrimState f ~ s, PrimMonad f) =>
a -> Int -> Int -> f a
inner a
forall a. Monoid a => a
mempty (Int
x1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst) (Int
x2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wDst Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
where
inner :: a -> Int -> Int -> f a
inner !a
acc Int
xl Int
xr
| Int
xl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
| Bool
otherwise =do
a
acc' <-
if Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
xl Int
0
then (a
acc <>) (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
xl)
else a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
a
acc'' <-
if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int
xr Int
0
then (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
acc') (a -> a) -> f a -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s a
MVector (PrimState f) a
dataDst (Int -> Int -> Int -> Int
idx Int
wDst Int
y Int
xr)
else a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc'
a -> Int -> Int -> f a
inner a
acc'' ((Int
xl 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
xr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
1)