{-# LANGUAGE RecordWildCards #-}

-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
-- commutative monoids.
--
-- - Point coordinates are fixed on `build`.
-- - Multiple points can exist at the same coordinate.
--
-- ==== __Examples__
-- >>> import AtCoder.Extra.LazyKdTree qualified as LKT
-- >>> import AtCoder.Extra.Monoid.Affine1 (Affine1)
-- >>> import AtCoder.Extra.Monoid.Affine1 qualified as Affine1
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let xyws = VU.fromList [(0, 0, Sum 1), (1, 1, Sum 2), (4, 2, Sum 3)]
-- >>> lkt <- LKT.build3 @_ @(Affine1 Int) @(Sum Int) xyws
--
-- >>> -- Get monoid product in [0, 2) x [0, 2)
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 3}
--
-- >>> LKT.applyIn lkt 0 2 0 2 $ Affine1.new 2 1
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 8}
--
-- >>> LKT.write lkt 0 $ Sum 10
-- >>> LKT.prod lkt 0 2 0 2
-- Sum {getSum = 15}
--
-- @since 1.2.2.0
module AtCoder.Extra.LazyKdTree
  ( -- * K-dimensional tree
    LazyKdTree (..),

    -- * Re-exports
    SegAct (..),

    -- * Constructors
    build,
    build2,
    build3,

    -- * Write
    write,
    modify,
    modifyM,

    -- * Monoid products
    prod,
    allProd,

    -- * Apply
    applyIn,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Bit qualified as ACIB
import AtCoder.LazySegTree (SegAct (..))
import Control.Monad (unless, when)
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.Ord (comparing)
import Data.Vector.Algorithms.Intro qualified as VAI
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)

-- | Static, \(k\)-dimensional tree \((k = 2)\) with lazily propagated monoid actions and
-- commutative monoids.
--
-- @since 1.2.2.0
data LazyKdTree s f a = LazyKdTree
  { -- | The number of points in the \(k\)-d tree.
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> Int
nLkt :: {-# UNPACK #-} !Int,
    -- | \(\lceil \log_2 (n + 1) \rceil\)
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> Int
logLkt :: {-# UNPACK #-} !Int,
    -- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\).
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
incRectsLkt :: !(VU.Vector (Int, Int, Int, Int)),
    -- | Rectangle information: monoid values.
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> MVector s a
dataLkt :: !(VUM.MVector s a),
    -- | Rectangle information: lazily propagated monoid actions for children.
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> MVector s f
lazyLkt :: !(VUM.MVector s f),
    -- | Rectangle information: the number of vertices in the rectangle.
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> Vector Int
sizeLkt :: !(VU.Vector Int),
    -- | Maps original vertices into the belonging rectangle index.
    --
    -- @since 1.2.2.0
    forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: !(VU.Vector Int)
  }

-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xs@, @ys@ and @ws@ vectors.
--
-- ==== Constraints
-- - \(|\mathrm{xs}| = |\mathrm{ys}| = |\mathrm{vs}|\).
--
-- @since 1.2.2.0
{-# INLINE build #-}
build ::
  (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | \(x\) coordnates
  VU.Vector Int ->
  -- | \(y\) coordnates
  VU.Vector Int ->
  -- | monoid \(v\)alues
  VU.Vector a ->
  -- | `LazyKdTree`
  m (LazyKdTree (PrimState m) f a)
build :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, Monoid f, Unbox f, Semigroup a,
 Unbox a) =>
Vector Int
-> Vector Int -> Vector a -> m (LazyKdTree (PrimState m) f a)
build Vector Int
xs Vector Int
ys Vector a
vs = ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (LazyKdTree (PrimState m) f a)
 -> m (LazyKdTree (PrimState m) f a))
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall a b. (a -> b) -> a -> b
$ Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
forall s f a.
(HasCallStack, Monoid f, Unbox f, Semigroup a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (LazyKdTree s f a)
buildST Vector Int
xs Vector Int
ys Vector a
vs

-- | \(O(n \log n)\) Creates a `LazyKdTree` from @xys@ and @ws@ vectors.
--
-- ==== Constraints
-- - \(|\mathrm{xys}| = |\mathrm{vs}|\).
--
-- @since 1.2.2.0
{-# INLINE build2 #-}
build2 ::
  (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | \((x, y)\) coordinates
  VU.Vector (Int, Int) ->
  -- | Monoid \(v\)alues
  VU.Vector a ->
  -- | `LazyKdTree`
  m (LazyKdTree (PrimState m) f a)
build2 :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, Monoid f, Unbox f, Semigroup a,
 Unbox a) =>
Vector (Int, Int) -> Vector a -> m (LazyKdTree (PrimState m) f a)
build2 Vector (Int, Int)
xys Vector a
ws = ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (LazyKdTree (PrimState m) f a)
 -> m (LazyKdTree (PrimState m) f a))
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall a b. (a -> b) -> a -> b
$ Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
forall s f a.
(HasCallStack, Monoid f, Unbox f, Semigroup a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (LazyKdTree s f a)
buildST Vector Int
xs Vector Int
ys Vector a
ws
  where
    (!Vector Int
xs, !Vector Int
ys) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
VU.unzip Vector (Int, Int)
xys

-- | \(O(n \log n)\) Creates a `LazyKdTree` from a @xyws@ vector.
--
-- @since 1.2.2.0
{-# INLINE build3 #-}
build3 ::
  (HasCallStack, PrimMonad m, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | \((x, y, v)\) tuples
  VU.Vector (Int, Int, a) ->
  -- | `LazyKdTree`
  m (LazyKdTree (PrimState m) f a)
build3 :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, Monoid f, Unbox f, Semigroup a,
 Unbox a) =>
Vector (Int, Int, a) -> m (LazyKdTree (PrimState m) f a)
build3 Vector (Int, Int, a)
xyws = ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (LazyKdTree (PrimState m) f a)
 -> m (LazyKdTree (PrimState m) f a))
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
-> m (LazyKdTree (PrimState m) f a)
forall a b. (a -> b) -> a -> b
$ Vector Int
-> Vector Int
-> Vector a
-> ST (PrimState m) (LazyKdTree (PrimState m) f a)
forall s f a.
(HasCallStack, Monoid f, Unbox f, Semigroup a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (LazyKdTree s f a)
buildST Vector Int
xs Vector Int
ys Vector a
ws
  where
    (!Vector Int
xs, !Vector Int
ys, !Vector a
ws) = Vector (Int, Int, a) -> (Vector Int, Vector Int, Vector a)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector (a, b, c) -> (Vector a, Vector b, Vector c)
VU.unzip3 Vector (Int, Int, a)
xyws

-- | \(O(\log n)\) Writes to the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINE write #-}
write ::
  (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | Original vertex index.
  Int ->
  -- | Monoid value
  a ->
  -- | Monadic tuple
  m ()
write :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, SegAct f a, Eq f, Unbox f, Semigroup a,
 Unbox a) =>
LazyKdTree (PrimState m) f a -> Int -> a -> m ()
write LazyKdTree (PrimState m) f a
kt Int
i a
x = 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
$ LazyKdTree (PrimState (ST (PrimState m))) f a
-> (a -> ST (PrimState m) a) -> Int -> ST (PrimState m) ()
forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, SegAct f a, Eq f, Unbox f, Semigroup a,
 Unbox a) =>
LazyKdTree (PrimState m) f a -> (a -> m a) -> Int -> m ()
modifyM LazyKdTree (PrimState m) f a
LazyKdTree (PrimState (ST (PrimState m))) f a
kt (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
x) Int
i

-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINE modify #-}
modify ::
  (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | Creates a new monoid value from the old one.
  (a -> a) ->
  -- | Original vertex index.
  Int ->
  -- | Monadic tuple
  m ()
modify :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, SegAct f a, Eq f, Unbox f, Semigroup a,
 Unbox a) =>
LazyKdTree (PrimState m) f a -> (a -> a) -> Int -> m ()
modify LazyKdTree (PrimState m) f a
kt a -> a
f Int
i = 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
$ LazyKdTree (PrimState (ST (PrimState m))) f a
-> (a -> ST (PrimState m) a) -> Int -> ST (PrimState m) ()
forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, SegAct f a, Eq f, Unbox f, Semigroup a,
 Unbox a) =>
LazyKdTree (PrimState m) f a -> (a -> m a) -> Int -> m ()
modifyM LazyKdTree (PrimState m) f a
LazyKdTree (PrimState (ST (PrimState m))) f a
kt (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
i

-- | \(O(\log n)\) Modifies the \(k\)-th point's monoid value.
--
-- @since 1.2.2.0
{-# INLINEABLE modifyM #-}
modifyM ::
  (HasCallStack, PrimMonad m, SegAct f a, Eq f, VU.Unbox f, Semigroup a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | Creates a new monoid value from the old one.
  (a -> m a) ->
  -- | Original vertex index.
  Int ->
  -- | Monadic tuple
  m ()
modifyM :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, SegAct f a, Eq f, Unbox f, Semigroup a,
 Unbox a) =>
LazyKdTree (PrimState m) f a -> (a -> m a) -> Int -> m ()
modifyM kt :: LazyKdTree (PrimState m) f a
kt@LazyKdTree {Int
MVector (PrimState m) f
MVector (PrimState m) a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: forall s f a. LazyKdTree s f a -> Int
logLkt :: forall s f a. LazyKdTree s f a -> Int
incRectsLkt :: forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
dataLkt :: forall s f a. LazyKdTree s f a -> MVector s a
lazyLkt :: forall s f a. LazyKdTree s f a -> MVector s f
sizeLkt :: forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: forall s f a. LazyKdTree s f a -> Vector Int
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector (PrimState m) a
lazyLkt :: MVector (PrimState m) f
sizeLkt :: Vector Int
posLkt :: Vector Int
..} a -> m a
f Int
i0 = do
  let i_ :: Int
i_ = Vector Int
posLkt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i0
  -- propagate lazily propagated monoid actions from the root:
  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
$ [Int] -> (Int -> ST (PrimState m) ()) -> ST (PrimState m) ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Int
logLkt, Int
logLkt Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 .. Int
1] ((Int -> ST (PrimState m) ()) -> ST (PrimState m) ())
-> (Int -> ST (PrimState m) ()) -> ST (PrimState m) ()
forall a b. (a -> b) -> a -> b
$ \Int
k -> do
    LazyKdTree (PrimState m) f a -> Int -> ST (PrimState m) ()
forall f a s.
(SegAct f a, Eq f, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> ST s ()
pushST LazyKdTree (PrimState m) f a
kt (Int
i_ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
.>>. Int
k)
  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
dataLkt a -> m a
f Int
i_
  -- update parents:
  let inner :: Int -> f ()
inner Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = () -> f ()
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            let i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i' 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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i' 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
dataLkt Int
i' (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 ()
inner Int
i'
  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
$ Int -> ST (PrimState m) ()
forall {f :: * -> *}.
(PrimState f ~ PrimState m, PrimMonad f) =>
Int -> f ()
inner Int
i_

-- | \(O(\log n)\) Returns monoid product in \([x_l, x_r) \times [y_l, y_r)\).
--
-- @since 1.2.2.0
{-# INLINE prod #-}
prod ::
  (HasCallStack, PrimMonad m, Eq f, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | \(x_l\)
  Int ->
  -- | \(x_r\)
  Int ->
  -- | \(y_l\)
  Int ->
  -- | \(y_r\)
  Int ->
  -- | Monoid product in \([x_l, x_r) \times [y_l, y_r)\)
  m a
prod :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, Eq f, SegAct f a, Eq f, Unbox f,
 Monoid a, Unbox a) =>
LazyKdTree (PrimState m) f a -> Int -> Int -> Int -> Int -> m a
prod LazyKdTree (PrimState m) f a
kt 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
$ LazyKdTree (PrimState m) f a
-> Int -> Int -> Int -> Int -> ST (PrimState m) a
forall f a s.
(HasCallStack, SegAct f a, Eq f, Unbox f, Monoid a, Unbox a) =>
LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a
prodST LazyKdTree (PrimState m) f a
kt Int
x1 Int
x2 Int
y1 Int
y2

-- | \(O(1)\) Returns monoid product of all the points.
--
-- @since 1.2.2.0
{-# INLINE allProd #-}
allProd ::
  (PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | Monoid product in the whole space.
  m a
allProd :: forall (m :: * -> *) a f.
(PrimMonad m, Monoid a, Unbox a) =>
LazyKdTree (PrimState m) f a -> m a
allProd LazyKdTree (PrimState m) f a
kt = do
  -- In case of zero vertices, use `Maybe`:
  a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
forall a. Monoid a => a
mempty (Maybe a -> a) -> m (Maybe a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) a -> Int -> m (Maybe a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (Maybe a)
VGM.readMaybe (LazyKdTree (PrimState m) f a -> MVector (PrimState m) a
forall s f a. LazyKdTree s f a -> MVector s a
dataLkt LazyKdTree (PrimState m) f a
kt) Int
1

-- | \(O(\log n)\) Applies a monoid action to points in \([x_l, x_r) \times [y_l, y_r)\).
--
-- @since 1.2.2.0
{-# INLINE applyIn #-}
applyIn ::
  (HasCallStack, PrimMonad m, Eq f, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) =>
  -- | `LazyKdTree`
  LazyKdTree (PrimState m) f a ->
  -- | \(x_l\)
  Int ->
  -- | \(x_r\)
  Int ->
  -- | \(y_l\)
  Int ->
  -- | \(y_r\)
  Int ->
  -- | \(f\)
  f ->
  -- | Monadic tuple
  m ()
applyIn :: forall (m :: * -> *) f a.
(HasCallStack, PrimMonad m, Eq f, SegAct f a, Unbox f, Monoid a,
 Unbox a) =>
LazyKdTree (PrimState m) f a
-> Int -> Int -> Int -> Int -> f -> m ()
applyIn LazyKdTree (PrimState m) f a
kt Int
x1 Int
x2 Int
y1 Int
y2 f
f = 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
$ LazyKdTree (PrimState m) f a
-> Int -> Int -> Int -> Int -> Int -> f -> ST (PrimState m) ()
forall f a s.
(HasCallStack, SegAct f a, Eq f, Unbox f, Monoid a, Unbox a) =>
LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s ()
applyInST LazyKdTree (PrimState m) f a
kt Int
1 Int
x1 Int
x2 Int
y1 Int
y2 f
f

-- -------------------------------------------------------------------------------------------------
-- Private
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE buildST #-}
buildST :: forall s f a. (HasCallStack, Monoid f, VU.Unbox f, Semigroup a, VU.Unbox a) => VU.Vector Int -> VU.Vector Int -> VU.Vector a -> ST s (LazyKdTree s f a)
buildST :: forall s f a.
(HasCallStack, Monoid f, Unbox f, Semigroup a, Unbox a) =>
Vector Int -> Vector Int -> Vector a -> ST s (LazyKdTree s f a)
buildST Vector Int
xs0 Vector Int
ys0 Vector a
vs0 = do
  let nLkt :: Int
nLkt = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
xs0
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
nLkt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ys0 Bool -> Bool -> Bool
&& Int
nLkt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
vs0) String
"AtCoder.Extra.LazyKdTree.buildST: the length of `xs`, `ys` and `vs` must be equal"
  if Int
nLkt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      let logLkt :: Int
logLkt = Int
0
      MVector s a
dataLkt <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
0
      MVector s f
lazyLkt <- Int -> ST s (MVector (PrimState (ST s)) f)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
0
      let incRectsLkt :: Vector (Int, Int, Int, Int)
incRectsLkt = Vector (Int, Int, Int, Int)
forall a. Unbox a => Vector a
VU.empty
      let sizeLkt :: Vector Int
sizeLkt = Vector Int
forall a. Unbox a => Vector a
VU.empty
      let posLkt :: Vector Int
posLkt = Vector Int
forall a. Unbox a => Vector a
VU.empty
      LazyKdTree s f a -> ST s (LazyKdTree s f a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
nLkt :: Int
logLkt :: Int
dataLkt :: MVector s a
lazyLkt :: MVector s f
incRectsLkt :: Vector (Int, Int, Int, Int)
sizeLkt :: Vector Int
posLkt :: Vector Int
..}
    else do
      let logLkt :: Int
logLkt = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
ACIB.bitCeil (Int
nLkt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      MVector s a
dataLkt <- Int -> ST s (MVector (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> Int
forall a. Bits a => Int -> a
bit (Int
logLkt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      MVector s f
lazyLkt <- Int -> f -> ST s (MVector (PrimState (ST s)) f)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int -> Int
forall a. Bits a => Int -> a
bit Int
logLkt) f
forall a. Monoid a => a
mempty
      MVector s (Int, Int, Int, Int)
incRectsVec <- Int
-> (Int, Int, Int, Int)
-> ST s (MVector (PrimState (ST s)) (Int, Int, Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int -> Int
forall a. Bits a => Int -> a
bit (Int
logLkt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound, Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound)
      MVector s Int
size <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> Int
forall a. Bits a => Int -> a
bit (Int
logLkt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
      MVector s Int
pos <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew Int
nLkt
      let VUM.MV_4 Int
_ MVector s Int
xMins MVector s Int
xMaxes MVector s Int
yMins MVector s Int
yMaxes = MVector s (Int, Int, Int, Int)
incRectsVec

      -- - idx: rectangle index (one-based)
      -- - xs, ys, vs: point information (x, y and monoid value)
      -- - ids: maps sorted vertices to the original vertex indices
      -- - divX: represents hyperplane direction for point partition
      let buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector a -> VU.Vector Int -> Bool -> ST s ()
          buildSubtree :: Int
-> Vector Int
-> Vector Int
-> Vector a
-> Vector Int
-> Bool
-> ST s ()
buildSubtree Int
idx Vector Int
xs Vector Int
ys Vector a
vs Vector Int
ids Bool
divX = do
            let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
xs
            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
size Int
idx Int
n

            -- retrieve the bounds:
            let (!Int
xMin, !Int
xMax, !Int
yMin, !Int
yMax) =
                  ((Int, Int, Int, Int) -> (Int, Int) -> (Int, Int, Int, Int))
-> (Int, Int, Int, Int)
-> Vector (Int, Int)
-> (Int, Int, Int, Int)
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl'
                    (\(!Int
a, !Int
b, !Int
c, !Int
d) (!Int
x, !Int
y) -> (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
a Int
x, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
b Int
x, Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
c Int
y, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
d Int
y))
                    (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound, Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound)
                    (Vector (Int, Int) -> (Int, Int, Int, Int))
-> Vector (Int, Int) -> (Int, Int, Int, Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
xs Vector Int
ys
            MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
xMins (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
xMin) Int
idx
            MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
xMaxes (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
xMax) Int
idx
            MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
yMins (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
yMin) Int
idx
            MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s Int
MVector (PrimState (ST s)) Int
yMaxes (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
yMax) Int
idx

            if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
              then do
                -- it's a terminal. note that it's not always a leaf; the case is handled carefully in
                -- other methods
                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
dataLkt Int
idx (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vector a
vs Vector a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0
                -- record original vertex index -> rectangle index
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
pos (Vector Int
ids Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0) Int
idx
              else do
                -- partition the vertices into two:
                let m :: Int
m = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                let is :: Vector Int
is = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
                      MVector s Int
vec <- Int -> (Int -> Int) -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
VUM.generate Int
n Int -> Int
forall a. a -> a
id
                      if Bool
divX
                        then Comparison Int -> MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
VAI.selectBy ((Int -> Int) -> Comparison Int
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Vector Int
xs VG.!)) MVector s Int
MVector (PrimState (ST s)) Int
vec Int
m
                        else Comparison Int -> MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
VAI.selectBy ((Int -> Int) -> Comparison Int
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Vector Int
ys VG.!)) MVector s Int
MVector (PrimState (ST s)) Int
vec Int
m
                      MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec

                -- TODO: permute in-place?
                let (!Vector Int
xsL, !Vector Int
xsR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
xs Vector Int
is
                let (!Vector Int
ysL, !Vector Int
ysR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
ys Vector Int
is
                let (!Vector a
vsL, !Vector a
vsR) = Int -> Vector a -> (Vector a, Vector a)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector a -> (Vector a, Vector a))
-> Vector a -> (Vector a, Vector a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector Int -> Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector a
vs Vector Int
is
                let (!Vector Int
idsL, !Vector Int
idsR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
ids Vector Int
is

                -- build the subtree:
                Int
-> Vector Int
-> Vector Int
-> Vector a
-> Vector Int
-> Bool
-> ST s ()
buildSubtree (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Vector Int
xsL Vector Int
ysL Vector a
vsL Vector Int
idsL (Bool -> Bool
not Bool
divX)
                Int
-> Vector Int
-> Vector Int
-> Vector a
-> Vector Int
-> Bool
-> ST s ()
buildSubtree (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector Int
xsR Vector Int
ysR Vector a
vsR Vector Int
idsR (Bool -> Bool
not Bool
divX)
                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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx 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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx 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
dataLkt Int
idx (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
-> Vector Int
-> Vector Int
-> Vector a
-> Vector Int
-> Bool
-> ST s ()
buildSubtree Int
1 Vector Int
xs0 Vector Int
ys0 Vector a
vs0 (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nLkt Int -> Int
forall a. a -> a
id) Bool
True
      Vector Int
sizeLkt <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
size
      Vector Int
posLkt <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s Int
MVector (PrimState (ST s)) Int
pos
      Vector (Int, Int, Int, Int)
incRectsLkt <- MVector (PrimState (ST s)) (Int, Int, Int, Int)
-> ST s (Vector (Int, Int, Int, Int))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector s (Int, Int, Int, Int)
MVector (PrimState (ST s)) (Int, Int, Int, Int)
incRectsVec
      LazyKdTree s f a -> ST s (LazyKdTree s f a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
nLkt :: Int
logLkt :: Int
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
incRectsLkt :: Vector (Int, Int, Int, Int)
..}

{-# INLINE applyAtST #-}
applyAtST :: (SegAct f a, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST :: forall f a s.
(SegAct f a, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: forall s f a. LazyKdTree s f a -> Int
logLkt :: forall s f a. LazyKdTree s f a -> Int
incRectsLkt :: forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
dataLkt :: forall s f a. LazyKdTree s f a -> MVector s a
lazyLkt :: forall s f a. LazyKdTree s f a -> MVector s f
sizeLkt :: forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: forall s f a. LazyKdTree s f a -> Vector Int
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
..} Int
i f
f = do
  -- NOTE: Here we're asssuming each monoid value has length one. If you need a monoid of length
  -- zero, e.g., if you're just reserving new point insertion, you must not rely on
  -- `segActWithLength`. You might want to use `V2` instead of `Sum`.
  let len :: Int
len = Vector Int
sizeLkt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
  MVector (PrimState (ST s)) a -> (a -> a) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s a
MVector (PrimState (ST s)) a
dataLkt (Int -> f -> a -> a
forall f a. SegAct f a => Int -> f -> a -> a
segActWithLength Int
len f
f) Int
i
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int
forall a. Bits a => Int -> a
bit Int
logLkt) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState (ST s)) f -> (f -> f) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector s f
MVector (PrimState (ST s)) f
lazyLkt (f
f <>) Int
i

-- TODO: consider `INLINE`?
{-# INLINE pushST #-}
pushST :: (SegAct f a, Eq f, VU.Unbox f, VU.Unbox a) => LazyKdTree s f a -> Int -> ST s ()
pushST :: forall f a s.
(SegAct f a, Eq f, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> ST s ()
pushST kt :: LazyKdTree s f a
kt@LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: forall s f a. LazyKdTree s f a -> Int
logLkt :: forall s f a. LazyKdTree s f a -> Int
incRectsLkt :: forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
dataLkt :: forall s f a. LazyKdTree s f a -> MVector s a
lazyLkt :: forall s f a. LazyKdTree s f a -> MVector s f
sizeLkt :: forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: forall s f a. LazyKdTree s f a -> Vector Int
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
..} Int
i = do
  f
lazy <- MVector (PrimState (ST s)) f -> Int -> ST s f
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read MVector s f
MVector (PrimState (ST s)) f
lazyLkt Int
i
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (f
lazy f -> f -> Bool
forall a. Eq a => a -> a -> Bool
== f
forall a. Monoid a => a
mempty) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    LazyKdTree s f a -> Int -> f -> ST s ()
forall f a s.
(SegAct f a, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST LazyKdTree s f a
kt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) f
lazy
    LazyKdTree s f a -> Int -> f -> ST s ()
forall f a s.
(SegAct f a, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST LazyKdTree s f a
kt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) f
lazy
    MVector (PrimState (ST s)) f -> Int -> f -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s f
MVector (PrimState (ST s)) f
lazyLkt Int
i f
forall a. Monoid a => a
mempty

{-# INLINEABLE prodST #-}
prodST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a
prodST :: forall f a s.
(HasCallStack, SegAct f a, Eq f, Unbox f, Monoid a, Unbox a) =>
LazyKdTree s f a -> Int -> Int -> Int -> Int -> ST s a
prodST kt :: LazyKdTree s f a
kt@LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: forall s f a. LazyKdTree s f a -> Int
logLkt :: forall s f a. LazyKdTree s f a -> Int
incRectsLkt :: forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
dataLkt :: forall s f a. LazyKdTree s f a -> MVector s a
lazyLkt :: forall s f a. LazyKdTree s f a -> MVector s f
sizeLkt :: forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: forall s f a. LazyKdTree s f a -> Vector Int
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
..} Int
x1 Int
x2 Int
y1 Int
y2
  | Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x2 Bool -> Bool -> Bool
|| Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y2 = a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
  | Bool
otherwise = Int -> ST s a
inner Int
1
  where
    inner :: Int -> ST s a
inner Int
i = case Vector (Int, Int, Int, Int)
incRectsLkt Vector (Int, Int, Int, Int) -> Int -> Maybe (Int, Int, Int, Int)
forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
VG.!? Int
i of
      Maybe (Int, Int, Int, Int)
Nothing -> a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
      Just (!Int
xl, !Int
xr, !Int
yl, !Int
yr)
        -- TODO: what is this?
        | Int
xl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr -> a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
        -- not intersecting
        | Int
x2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
xl Bool -> Bool -> Bool
|| Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr Bool -> Bool -> Bool
|| Int
y2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yl Bool -> Bool -> Bool
|| Int
y1 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
forall a. Monoid a => a
mempty
        -- the rectangle is fully contained by the query:
        | Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
xl Bool -> Bool -> Bool
&& Int
xr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x2 Bool -> Bool -> Bool
&& Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yl Bool -> Bool -> Bool
&& Int
yr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y2 -> do
            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
dataLkt Int
i
        | Bool
otherwise -> do
            LazyKdTree s f a -> Int -> ST s ()
forall f a s.
(SegAct f a, Eq f, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> ST s ()
pushST LazyKdTree s f a
kt Int
i
            a
l <- Int -> ST s a
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0)
            a
r <- Int -> ST s a
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            a -> ST s a
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> ST s a) -> a -> ST s a
forall a b. (a -> b) -> a -> b
$! a
l a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
r

{-# INLINEABLE applyInST #-}
applyInST :: (HasCallStack, SegAct f a, Eq f, VU.Unbox f, Monoid a, VU.Unbox a) => LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s ()
applyInST :: forall f a s.
(HasCallStack, SegAct f a, Eq f, Unbox f, Monoid a, Unbox a) =>
LazyKdTree s f a -> Int -> Int -> Int -> Int -> Int -> f -> ST s ()
applyInST kt :: LazyKdTree s f a
kt@LazyKdTree {Int
MVector s f
MVector s a
Vector Int
Vector (Int, Int, Int, Int)
nLkt :: forall s f a. LazyKdTree s f a -> Int
logLkt :: forall s f a. LazyKdTree s f a -> Int
incRectsLkt :: forall s f a. LazyKdTree s f a -> Vector (Int, Int, Int, Int)
dataLkt :: forall s f a. LazyKdTree s f a -> MVector s a
lazyLkt :: forall s f a. LazyKdTree s f a -> MVector s f
sizeLkt :: forall s f a. LazyKdTree s f a -> Vector Int
posLkt :: forall s f a. LazyKdTree s f a -> Vector Int
nLkt :: Int
logLkt :: Int
incRectsLkt :: Vector (Int, Int, Int, Int)
dataLkt :: MVector s a
lazyLkt :: MVector s f
sizeLkt :: Vector Int
posLkt :: Vector Int
..} Int
i0 Int
x1 Int
x2 Int
y1 Int
y2 f
f
  | Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x2 Bool -> Bool -> Bool
|| Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y2 = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  | Bool
otherwise = Int -> ST s ()
inner Int
i0
  where
    inner :: Int -> ST s ()
inner Int
i = case Vector (Int, Int, Int, Int)
incRectsLkt Vector (Int, Int, Int, Int) -> Int -> Maybe (Int, Int, Int, Int)
forall (v :: * -> *) a. Vector v a => v a -> Int -> Maybe a
VG.!? Int
i of
      Maybe (Int, Int, Int, Int)
Nothing -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall a. Monoid a => a
mempty
      Just (!Int
xl, !Int
xr, !Int
yl, !Int
yr)
        | Int
xl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        -- not intersecting
        | Int
x2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
xl Bool -> Bool -> Bool
|| Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
xr Bool -> Bool -> Bool
|| Int
y2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yl Bool -> Bool -> Bool
|| Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
yr -> () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        -- the rectangle is fully contained by the query:
        | Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
xl Bool -> Bool -> Bool
&& Int
xr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x2 Bool -> Bool -> Bool
&& Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yl Bool -> Bool -> Bool
&& Int
yr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y2 -> do
            LazyKdTree s f a -> Int -> f -> ST s ()
forall f a s.
(SegAct f a, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> f -> ST s ()
applyAtST LazyKdTree s f a
kt Int
i f
f
        | Bool
otherwise -> do
            LazyKdTree s f a -> Int -> ST s ()
forall f a s.
(SegAct f a, Eq f, Unbox f, Unbox a) =>
LazyKdTree s f a -> Int -> ST s ()
pushST LazyKdTree s f a
kt Int
i
            Int -> ST s ()
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0)
            Int -> ST s ()
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            a
l <- 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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0)
            a
r <- 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
dataLkt (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i 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
dataLkt Int
i (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
l a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
r