{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

-- | Link/cut tree: forest with monoid values.
--
-- ==== __Example__
--
-- Create a link/cut tree of @Sum Int@ with inverse operator `negate`:
--
-- >>> import AtCoder.Extra.Tree.Lct qualified as Lct
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> -- 0--1--2
-- >>> --    +--3
-- >>> lct <- Lct.buildInv negate (VU.generate 4 Sum) $ VU.fromList [(0, 1), (1, 2), (1, 3)]
--
-- Monoid products can be calculated for paths or subtrees:
--
-- >>> Lct.prodPath lct 0 2
-- Sum {getSum = 3}
--
-- >>> Lct.prodSubtree lct 1 {- parent -} 2
-- Sum {getSum = 4}
--
-- `root` returns the current root vertex of the underlying tree, which is not easy to predict:
--
-- >>> Lct.root lct 3
-- 2
--
-- Set (`evert`) the root of the underlying tree to \(0\) and get the `lca` of vertices \(2\) and
-- \(3\):
--
-- >>> Lct.evert lct 0
-- >>> Lct.lca lct 2 3
-- 1
--
-- Similar to @Hld@, `Lct` allows various tree queries:
--
-- >>> Lct.parent lct 3
-- Just 1
--
-- >>> Lct.jump lct 2 3 2
-- 3
--
-- Edges can be dynamically added (`link`) or removed (`cut`):
--
-- >>> -- 0  1  2
-- >>> --    +--3
-- >>> Lct.cut lct 0 1
-- >>> Lct.cut lct 1 2
-- >>> VU.generateM 4 (Lct.root lct)
-- [0,1,2,1]
--
-- >>> -- +-----+
-- >>> -- 0  1  2
-- >>> --    +--3
-- >>> Lct.link lct 0 2
-- >>> VU.generateM 4 (Lct.root lct)
-- [2,1,2,1]
--
-- @since 1.1.1.0
module AtCoder.Extra.Tree.Lct
  ( -- Link/cut tree
    Lct (..),
    Vertex,

    -- * Constructors
    new,
    newInv,
    build,
    buildInv,

    -- * Modifications

    -- ** Write
    write,
    modify,
    modifyM,

    -- ** Link/cut
    link,
    cut,

    -- ** Evert/expose
    evert,
    expose,
    expose_,

    -- * Tree queries

    -- ** Root, parent, jump, LCA
    root,
    parent,
    jump,
    lca,

    -- ** Products
    prodPath,
    prodSubtree,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import Control.Monad (unless, when)
import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
import Control.Monad.ST (ST)
import Data.Bit
import Data.Bits
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 GHC.Stack (HasCallStack)

-- | Alias of vertex type.
type Vertex = Int

{-# INLINE undefLct #-}
undefLct :: Vertex
undefLct :: Vertex
undefLct = -Vertex
1

{-# INLINE nullLct #-}
nullLct :: Vertex -> Bool
nullLct :: Vertex -> Bool
nullLct = (Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== -Vertex
1)

-- We could optimize the with options, but

-- | Link/cut tree.
--
-- @since 1.1.1.0
data Lct s a = Lct
  { -- | The number of vertices.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> Vertex
nLct :: {-# UNPACK #-} !Int,
    -- | Decomposed node data storage: left children.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s Vertex
lLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: right children.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s Vertex
rLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: parents.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s Vertex
pLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: subtree sizes.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s Vertex
sLct :: !(VUM.MVector s Int),
    -- | Decomposed node data storage: reverse flag.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s Bit
revLct :: !(VUM.MVector s Bit),
    -- | Decomposed node data storage: monoid values.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s a
vLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: monoid products.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s a
prodLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: dual monod product (right fold). This is required for
    -- non-commutative monoids only.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s a
dualProdLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: path-parent monoid product. This works for subtree product
    -- queries over commutative monoids only.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s a
midLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: monoid product of subtree. This works for subtree product
    -- queries over commutative monoids only.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> MVector s a
subtreeProdLct :: !(VUM.MVector s a),
    -- | Inverse operator of the monoid. This works for subtree product queries over commutative
    -- monoids only.
    --
    -- @since 1.1.1.0
    forall s a. Lct s a -> a -> a
invOpLct :: !(a -> a)
  }

-- | \(O(n)\) Creates a link/cut tree with \(n\) vertices and no edges.
--
-- @since 1.1.1.0
{-# INLINE new #-}
new :: (PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (Lct (PrimState m) a)
new :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
Vertex -> m (Lct (PrimState m) a)
new = (a -> a) -> Vertex -> m (Lct (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
(a -> a) -> Vertex -> m (Lct (PrimState m) a)
newInv a -> a
forall a. a -> a
id

-- | \(O(n + m \log n)\) Creates a link/cut tree with an inverse operator, initial monoid values and
-- no edges. This setup enables subtree queries (`prodSubtree`).
--
-- @since 1.1.1.0
{-# INLINE newInv #-}
newInv :: (PrimMonad m, Monoid a, VU.Unbox a) => (a -> a) -> Int -> m (Lct (PrimState m) a)
newInv :: forall (m :: * -> *) a.
(PrimMonad m, Monoid a, Unbox a) =>
(a -> a) -> Vertex -> m (Lct (PrimState m) a)
newInv !a -> a
invOpLct Vertex
nLct = (a -> a)
-> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
(a -> a)
-> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
buildInv a -> a
invOpLct (Vertex -> a -> Vector a
forall a. Unbox a => Vertex -> a -> Vector a
VU.replicate Vertex
nLct a
forall a. Monoid a => a
mempty) Vector (Vertex, Vertex)
forall a. Unbox a => Vector a
VU.empty

-- | \(O(n + m \log n)\) Creates a link/cut tree of initial monoid values and initial edges.
--
-- @since 1.1.1.0
{-# INLINE build #-}
build ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | Vertex monoid values
  VU.Vector a ->
  -- | Edges
  VU.Vector (Vertex, Vertex) ->
  -- | Link/cut tree
  m (Lct (PrimState m) a)
build :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
build = (a -> a)
-> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
(a -> a)
-> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
buildInv a -> a
forall a. a -> a
id

-- | \(O(n + m \log n)\) Creates a link/cut tree with an inverse operator, initial monoid values and
-- initial edges. This setup enables subtree queries (`prodSubtree`).
--
-- @since 1.1.1.0
{-# INLINE buildInv #-}
buildInv ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | Inverse operator
  (a -> a) ->
  -- | Vertex monoid values
  VU.Vector a ->
  -- | Edges
  VU.Vector (Vertex, Vertex) ->
  -- | Link/cut tree
  m (Lct (PrimState m) a)
buildInv :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
(a -> a)
-> Vector a -> Vector (Vertex, Vertex) -> m (Lct (PrimState m) a)
buildInv !a -> a
invOpLct Vector a
xs Vector (Vertex, Vertex)
es = do
  Lct (PrimState m) a
lct <- do
    let !nLct :: Vertex
nLct = Vector a -> Vertex
forall a. Unbox a => Vector a -> Vertex
VU.length Vector a
xs
    MVector (PrimState m) Vertex
lLct <- Vertex -> Vertex -> m (MVector (PrimState m) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct Vertex
undefLct
    MVector (PrimState m) Vertex
rLct <- Vertex -> Vertex -> m (MVector (PrimState m) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct Vertex
undefLct
    MVector (PrimState m) Vertex
pLct <- Vertex -> Vertex -> m (MVector (PrimState m) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct Vertex
undefLct
    MVector (PrimState m) Vertex
sLct <- Vertex -> Vertex -> m (MVector (PrimState m) Vertex)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct Vertex
0
    MVector (PrimState m) Bit
revLct <- Vertex -> Bit -> m (MVector (PrimState m) Bit)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct (Bool -> Bit
Bit Bool
False)
    MVector (PrimState m) a
vLct <- Vector a -> m (MVector (PrimState m) a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.thaw Vector a
xs
    MVector (PrimState m) a
prodLct <- Vertex -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct a
forall a. Monoid a => a
mempty
    MVector (PrimState m) a
dualProdLct <- Vertex -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct a
forall a. Monoid a => a
mempty
    MVector (PrimState m) a
midLct <- Vertex -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct a
forall a. Monoid a => a
mempty
    MVector (PrimState m) a
subtreeProdLct <- Vertex -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Vertex -> a -> m (MVector (PrimState m) a)
VUM.replicate Vertex
nLct a
forall a. Monoid a => a
mempty
    Lct (PrimState m) a -> m (Lct (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lct {Vertex
MVector (PrimState m) a
MVector (PrimState m) Vertex
MVector (PrimState m) Bit
a -> a
nLct :: Vertex
lLct :: MVector (PrimState m) Vertex
rLct :: MVector (PrimState m) Vertex
pLct :: MVector (PrimState m) Vertex
sLct :: MVector (PrimState m) Vertex
revLct :: MVector (PrimState m) Bit
vLct :: MVector (PrimState m) a
prodLct :: MVector (PrimState m) a
dualProdLct :: MVector (PrimState m) a
midLct :: MVector (PrimState m) a
subtreeProdLct :: MVector (PrimState m) a
invOpLct :: a -> a
invOpLct :: a -> a
nLct :: Vertex
lLct :: MVector (PrimState m) Vertex
rLct :: MVector (PrimState m) Vertex
pLct :: MVector (PrimState m) Vertex
sLct :: MVector (PrimState m) Vertex
revLct :: MVector (PrimState m) Bit
vLct :: MVector (PrimState m) a
prodLct :: MVector (PrimState m) a
dualProdLct :: MVector (PrimState m) a
midLct :: MVector (PrimState m) a
subtreeProdLct :: MVector (PrimState m) a
..}
  Vector (Vertex, Vertex) -> ((Vertex, Vertex) -> m ()) -> m ()
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
VU.forM_ Vector (Vertex, Vertex)
es (((Vertex, Vertex) -> m ()) -> m ())
-> ((Vertex, Vertex) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(!Vertex
u, !Vertex
v) -> do
    Lct (PrimState m) a -> Vertex -> Vertex -> m ()
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m ()
link Lct (PrimState m) a
lct Vertex
u Vertex
v
  Lct (PrimState m) a -> m (Lct (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lct (PrimState m) a
lct

-- -------------------------------------------------------------------------------------------------
-- Balancing
-- -------------------------------------------------------------------------------------------------

-- | \(O(1)\) Rotates up a non-root node.
{-# INLINEABLE rotateST #-}
rotateST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST :: forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST lct :: Lct s a
lct@Lct {MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct, MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct} Vertex
v = do
  Vertex
p <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
v
  Vertex
pp <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
p
  Vertex
pl <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
p

  Vertex
c <-
    if Vertex
pl Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
v
      then do
        -- rotate right:
        --   p      v  <-- reference from `pp` is updated later
        --  /        \
        -- v    ->    p
        --  \        /
        --   c      c
        Vertex
c <- MVector (PrimState (ST s)) Vertex
-> Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m a
VGM.unsafeExchange MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
v Vertex
p
        MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
p Vertex
c
        Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
c
      else do
        -- rotate left:
        -- p          v  <-- reference from `pp` is updated later
        --  \        /
        --   v  ->  p
        --  /        \
        -- c          c
        Vertex
c <- MVector (PrimState (ST s)) Vertex
-> Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m a
VGM.unsafeExchange MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
v Vertex
p
        MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
p Vertex
c
        Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
c

  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct s a
lct Vertex
p
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct s a
lct Vertex
v

  -- update the reference from `pp`:
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
pp) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    Vertex
ppl <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
pp
    if Vertex
ppl Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
p
      then MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
pp Vertex
v
      else do
        Vertex
ppr <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
pp
        if Vertex
ppr Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
p
          then MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
pp Vertex
v
          else do
            -- overwrite the light (path-parent) pointer:
            Lct s a -> Vertex -> Vertex -> Vertex -> ST s ()
forall s a. Lct s a -> Vertex -> Vertex -> Vertex -> ST s ()
changeLightST Lct s a
lct Vertex
pp Vertex
p Vertex
v

  -- update parent pointers to `pp`: pp <-- v <-- p <-- c
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
v Vertex
pp
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
p Vertex
v
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
c) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
c Vertex
p

-- | Amortized \(O(\log n)\). Moves a node up to the root, performing self-balancing heuristic
-- called rotations.
{-# INLINEABLE splayST #-}
splayST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()
splayST :: forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
splayST lct :: Lct s a
lct@Lct {MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct} Vertex
c = do
  Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
c
  let inner :: ST s ()
inner = do
        Bool
isRootC <- Lct s a -> Vertex -> ST s Bool
forall s a. Lct s a -> Vertex -> ST s Bool
isRootNodeST Lct s a
lct Vertex
c
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
isRootC (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          Vertex
p <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
c
          Vertex
pp <- if Vertex -> Bool
nullLct Vertex
p then Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
undefLct else MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
p
          NodePlaceLct
placeP <- Lct s a -> Vertex -> ST s NodePlaceLct
forall s a. Lct s a -> Vertex -> ST s NodePlaceLct
nodePlaceST Lct s a
lct Vertex
p
          if NodePlaceLct
placeP NodePlaceLct -> NodePlaceLct -> Bool
forall a. Eq a => a -> a -> Bool
== NodePlaceLct
RootNodeLct
            then do
              Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
p
              Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
c
              Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST Lct s a
lct Vertex
c
            else do
              NodePlaceLct
placeC <- Lct s a -> Vertex -> ST s NodePlaceLct
forall s a. Lct s a -> Vertex -> ST s NodePlaceLct
nodePlaceST Lct s a
lct Vertex
c
              Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
pp
              Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
p
              Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
c
              if NodePlaceLct
placeC NodePlaceLct -> NodePlaceLct -> Bool
forall a. Eq a => a -> a -> Bool
== NodePlaceLct
placeP
                then do
                  -- Rotate right twice:
                  --
                  --       pp       p         c
                  --      /        / \         \
                  --    p     ->  c   pp  ->    p
                  --   /                         \
                  -- c                            pp

                  -- Or rotate left twice:
                  --
                  --  pp             p            c
                  --   \            / \          /
                  --    p     ->  pp   c  ->    p
                  --     \                     /
                  --      c                   pp

                  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST Lct s a
lct Vertex
p
                  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST Lct s a
lct Vertex
c
                else do
                  --       pp         pp         c
                  --      /          /          | \
                  --    p     ->   c      ->   p   pp
                  --     \        /
                  --      c      p
                  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST Lct s a
lct Vertex
c
                  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
rotateST Lct s a
lct Vertex
c
          ST s ()
inner
  ST s ()
inner

-- * Node helpers

-- | \(O(1)\)
{-# INLINEABLE isRootNodeST #-}
isRootNodeST :: Lct s a -> Vertex -> ST s Bool
isRootNodeST :: forall s a. Lct s a -> Vertex -> ST s Bool
isRootNodeST Lct s a
lct Vertex
v = do
  (NodePlaceLct -> NodePlaceLct -> Bool
forall a. Eq a => a -> a -> Bool
== NodePlaceLct
RootNodeLct) (NodePlaceLct -> Bool) -> ST s NodePlaceLct -> ST s Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lct s a -> Vertex -> ST s NodePlaceLct
forall s a. Lct s a -> Vertex -> ST s NodePlaceLct
nodePlaceST Lct s a
lct Vertex
v

-- TODO: return heavy/light notion
data NodePlaceLct = RootNodeLct | LeftNodeLct | RightNodeLct
  deriving (NodePlaceLct -> NodePlaceLct -> Bool
(NodePlaceLct -> NodePlaceLct -> Bool)
-> (NodePlaceLct -> NodePlaceLct -> Bool) -> Eq NodePlaceLct
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NodePlaceLct -> NodePlaceLct -> Bool
== :: NodePlaceLct -> NodePlaceLct -> Bool
$c/= :: NodePlaceLct -> NodePlaceLct -> Bool
/= :: NodePlaceLct -> NodePlaceLct -> Bool
Eq)

-- | \(O(1)\)
{-# INLINEABLE nodePlaceST #-}
nodePlaceST :: Lct s a -> Vertex -> ST s NodePlaceLct
nodePlaceST :: forall s a. Lct s a -> Vertex -> ST s NodePlaceLct
nodePlaceST Lct {MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct, MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct} Vertex
v = do
  Vertex
p <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
v
  if Vertex -> Bool
nullLct Vertex
p
    then NodePlaceLct -> ST s NodePlaceLct
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodePlaceLct
RootNodeLct
    else do
      Vertex
pl <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
p
      if Vertex
pl Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
v
        then NodePlaceLct -> ST s NodePlaceLct
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodePlaceLct
LeftNodeLct
        else do
          Vertex
pr <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
p
          if Vertex
pr Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
v
            then NodePlaceLct -> ST s NodePlaceLct
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodePlaceLct
RightNodeLct
            else NodePlaceLct -> ST s NodePlaceLct
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NodePlaceLct
RootNodeLct

-- -------------------------------------------------------------------------------------------------
-- Node operations
-- -------------------------------------------------------------------------------------------------

-- | \(O(1)\) Propgates the lazily propagated values on a node.
{-# INLINEABLE pushNodeST #-}
pushNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()
pushNodeST :: forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST lct :: Lct s a
lct@Lct {MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct, MVector s Bit
revLct :: forall s a. Lct s a -> MVector s Bit
revLct :: MVector s Bit
revLct} Vertex
v = do
  Bit Bool
b <- MVector (PrimState (ST s)) Bit -> Vertex -> Bit -> ST s Bit
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m a
VGM.unsafeExchange MVector s Bit
MVector (PrimState (ST s)) Bit
revLct Vertex
v (Bool -> Bit
Bit Bool
False)
  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
    Vertex
l <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
v
    Vertex
r <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
v
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
l) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
reverseNodeST Lct s a
lct Vertex
l
    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
r) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
reverseNodeST Lct s a
lct Vertex
r

-- | \(O(1)\)
{-# INLINEABLE reverseNodeST #-}
reverseNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()
reverseNodeST :: forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
reverseNodeST lct :: Lct s a
lct@Lct {MVector s Bit
revLct :: forall s a. Lct s a -> MVector s Bit
revLct :: MVector s Bit
revLct} Vertex
i = do
  -- lazily propagate new reverse from the children, or cancel:
  MVector (PrimState (ST s)) Bit -> (Bit -> Bit) -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Vertex -> m ()
VGM.unsafeModify MVector s Bit
MVector (PrimState (ST s)) Bit
revLct (Bit -> Bit -> Bit
forall a. Bits a => a -> a -> a
xor (Bool -> Bit
Bit Bool
True)) Vertex
i
  -- swap
  Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
swapLrNodeST Lct s a
lct Vertex
i

-- | \(O(1)\) Reverses the left and the right children, lazily and recursively.
{-# INLINEABLE swapLrNodeST #-}
swapLrNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()
swapLrNodeST :: forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
swapLrNodeST Lct {MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct, MVector s a
prodLct :: forall s a. Lct s a -> MVector s a
prodLct :: MVector s a
prodLct, MVector s a
dualProdLct :: forall s a. Lct s a -> MVector s a
dualProdLct :: MVector s a
dualProdLct} Vertex
i = do
  -- swap chidlren
  MVector (PrimState (ST s)) Vertex
-> (Vertex -> ST s Vertex) -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Vertex -> m ()
VGM.unsafeModifyM MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct (MVector (PrimState (ST s)) Vertex
-> Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m a
VGM.unsafeExchange MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
i) Vertex
i
  -- swap prodLct[i] and dualProdLct[i]
  MVector (PrimState (ST s)) a -> (a -> ST s a) -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Vertex -> m ()
VGM.unsafeModifyM MVector s a
MVector (PrimState (ST s)) a
prodLct (MVector (PrimState (ST s)) a -> Vertex -> a -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m a
VGM.unsafeExchange MVector s a
MVector (PrimState (ST s)) a
dualProdLct Vertex
i) Vertex
i

-- | \(O(1)\) Recomputes the node size and the monoid product.
{-# INLINEABLE updateNodeST #-}
updateNodeST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST :: forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct {Vertex
MVector s a
MVector s Vertex
MVector s Bit
a -> a
nLct :: forall s a. Lct s a -> Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
sLct :: forall s a. Lct s a -> MVector s Vertex
revLct :: forall s a. Lct s a -> MVector s Bit
vLct :: forall s a. Lct s a -> MVector s a
prodLct :: forall s a. Lct s a -> MVector s a
dualProdLct :: forall s a. Lct s a -> MVector s a
midLct :: forall s a. Lct s a -> MVector s a
subtreeProdLct :: forall s a. Lct s a -> MVector s a
invOpLct :: forall s a. Lct s a -> a -> a
nLct :: Vertex
lLct :: MVector s Vertex
rLct :: MVector s Vertex
pLct :: MVector s Vertex
sLct :: MVector s Vertex
revLct :: MVector s Bit
vLct :: MVector s a
prodLct :: MVector s a
dualProdLct :: MVector s a
midLct :: MVector s a
subtreeProdLct :: MVector s a
invOpLct :: a -> a
..} Vertex
i = do
  Vertex
l <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
i
  Vertex
r <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
i
  a
v <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
vLct Vertex
i
  a
m <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
midLct Vertex
i

  (!Vertex
size', !a
prod', !a
dualProd', !a
subtreeProd') <-
    if Vertex -> Bool
nullLct Vertex
l
      then (Vertex, a, a, a) -> ST s (Vertex, a, a, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
1 :: Int, a
v, a
v, a
v a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
m)
      else do
        Vertex
lSize <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
sLct Vertex
l
        a
lProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
prodLct Vertex
l
        a
lDualProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
dualProdLct Vertex
l
        a
lSubtreeProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
subtreeProdLct Vertex
l
        (Vertex, a, a, a) -> ST s (Vertex, a, a, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
lSize Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1, a
lProd a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
v, a
v a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
lDualProd, a
lSubtreeProd a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
v a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
m)

  (!Vertex
size'', !a
prod'', !a
dualProd'', !a
subtreeProd'') <-
    if Vertex -> Bool
nullLct Vertex
r
      then (Vertex, a, a, a) -> ST s (Vertex, a, a, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
size', a
prod', a
dualProd', a
subtreeProd')
      else do
        Vertex
rSize <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
sLct Vertex
r
        a
rProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
prodLct Vertex
r
        a
rDualProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
dualProdLct Vertex
r
        a
rSubtreeProd <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
subtreeProdLct Vertex
r
        (Vertex, a, a, a) -> ST s (Vertex, a, a, a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vertex
size' Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
rSize, a
prod' a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
rProd, a
rDualProd a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
dualProd', a
subtreeProd' a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
rSubtreeProd)

  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
sLct Vertex
i Vertex
size''
  MVector (PrimState (ST s)) a -> Vertex -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s a
MVector (PrimState (ST s)) a
prodLct Vertex
i a
prod''
  MVector (PrimState (ST s)) a -> Vertex -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s a
MVector (PrimState (ST s)) a
dualProdLct Vertex
i a
dualProd''
  MVector (PrimState (ST s)) a -> Vertex -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s a
MVector (PrimState (ST s)) a
subtreeProdLct Vertex
i a
subtreeProd''

-- | \(O(1)\) Called on adding a path-parent edge. This is for subtree folding.
{-# INLINEABLE addLightST #-}
addLightST :: (Semigroup a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()
addLightST :: forall a s.
(Semigroup a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
addLightST Lct {MVector s a
subtreeProdLct :: forall s a. Lct s a -> MVector s a
subtreeProdLct :: MVector s a
subtreeProdLct, MVector s a
midLct :: forall s a. Lct s a -> MVector s a
midLct :: MVector s a
midLct} Vertex
p Vertex
c = do
  a
newChild <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
subtreeProdLct Vertex
c
  MVector (PrimState (ST s)) a -> (a -> a) -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Vertex -> m ()
VGM.unsafeModify MVector s a
MVector (PrimState (ST s)) a
midLct (a
newChild <>) Vertex
p

-- | \(O(1)\) Called on changing a path-parent edge. This is for subtree folding.
{-# INLINEABLE changeLightST #-}
changeLightST :: Lct s a -> Vertex -> Vertex -> Vertex -> ST s ()
changeLightST :: forall s a. Lct s a -> Vertex -> Vertex -> Vertex -> ST s ()
changeLightST Lct s a
_lct Vertex
_u Vertex
_v Vertex
_p = do
  () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | \(O(1)\) Called on erasing a path-parent edge. This is for subtree folding.
{-# INLINEABLE eraseLightST #-}
eraseLightST :: (Semigroup a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()
eraseLightST :: forall a s.
(Semigroup a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
eraseLightST Lct {MVector s a
subtreeProdLct :: forall s a. Lct s a -> MVector s a
subtreeProdLct :: MVector s a
subtreeProdLct, MVector s a
midLct :: forall s a. Lct s a -> MVector s a
midLct :: MVector s a
midLct, a -> a
invOpLct :: forall s a. Lct s a -> a -> a
invOpLct :: a -> a
invOpLct} Vertex
p Vertex
c = do
  a
sub <- MVector (PrimState (ST s)) a -> Vertex -> ST s a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s a
MVector (PrimState (ST s)) a
subtreeProdLct Vertex
c
  let !sub' :: a
sub' = a -> a
invOpLct a
sub
  MVector (PrimState (ST s)) a -> (a -> a) -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Vertex -> m ()
VGM.unsafeModify MVector s a
MVector (PrimState (ST s)) a
midLct (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
sub') Vertex
p

-- -------------------------------------------------------------------------------------------------
-- Write
-- -------------------------------------------------------------------------------------------------

-- TODO: read

-- | Amortized \(O(\log n)\). Writes the monoid value of a vertex.
--
-- @since 1.1.1.0
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> a -> m ()
write :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> a -> m ()
write Lct (PrimState m) a
lct Vertex
v 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
$ do
  -- make @v@ the new root of the underlying tree:
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
v
  MVector (PrimState (ST (PrimState m))) a
-> Vertex -> a -> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite (Lct (PrimState m) a -> MVector (PrimState m) a
forall s a. Lct s a -> MVector s a
vLct Lct (PrimState m) a
lct) Vertex
v a
x
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.write" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | Amortized \(O(\log n)\). Modifies the monoid value of a vertex with a pure function.
--
-- @since 1.1.1.0
{-# INLINE modify #-}
modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> (a -> a) -> Vertex -> m ()
modify :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> (a -> a) -> Vertex -> m ()
modify Lct (PrimState m) a
lct a -> a
f Vertex
v = 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
  -- make @v@ the new root of the underlying tree:
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
v
  MVector (PrimState (ST (PrimState m))) a
-> (a -> a) -> Vertex -> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Vertex -> m ()
VGM.unsafeModify (Lct (PrimState m) a -> MVector (PrimState m) a
forall s a. Lct s a -> MVector s a
vLct Lct (PrimState m) a
lct) a -> a
f Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.modify" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | Amortized \(O(\log n)\). Modifies the monoid value of a vertex with a monadic function.
--
-- @since 1.1.1.0
{-# INLINE modifyM #-}
modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> (a -> m a) -> Vertex -> m ()
modifyM :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> (a -> m a) -> Vertex -> m ()
modifyM Lct (PrimState m) a
lct a -> m a
f Vertex
v = do
  -- make @v@ the new root of the underlying tree:
  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
$ Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
v
  MVector (PrimState m) a -> (a -> m a) -> Vertex -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> m a) -> Vertex -> m ()
VGM.unsafeModifyM (Lct (PrimState m) a -> MVector (PrimState m) a
forall s a. Lct s a -> MVector s a
vLct Lct (PrimState m) a
lct) a -> m a
f Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.modifyM" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- -------------------------------------------------------------------------------------------------
-- Link/cut operations
-- -------------------------------------------------------------------------------------------------

-- | Amortized \(O(\log n)\).
{-# INLINEABLE linkST #-}
linkST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()
linkST :: forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
linkST lct :: Lct s a
lct@Lct {MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct} Vertex
c Vertex
p = do
  -- make @c@ the new root of the underlying tree
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct s a
lct Vertex
c
  -- remove right children of @p@.
  Vertex
_ <- Lct s a -> Vertex -> ST s Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct s a
lct Vertex
p
  Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
p

  -- dbgM $ do
  --   cp <- VGM.unsafeRead pLct c
  --   let !_ = ACIA.runtimeAssert (nullLct cp) $ "cp must be null: " ++ show (c, cp)
  --   pr <- VGM.unsafeRead rLct p
  --   let !_ = ACIA.runtimeAssert (nullLct pr) $ "pr must be null: " ++ show (p, pr)
  --   pure ()

  -- connect with a heavy edge:
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
c Vertex
p
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
p Vertex
c
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct s a
lct Vertex
p

-- | Amortized \(O(\log n)\). Creates an edge between \(c\) and \(p\). In the represented tree, the
-- parent of \(c\) will be \(p\) after this operation.
--
-- @since 1.1.1.0
{-# INLINE link #-}
link :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
link :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m ()
link Lct (PrimState m) a
lct Vertex
c Vertex
p = 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
$ Lct (PrimState m) a -> Vertex -> Vertex -> ST (PrimState m) ()
forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
linkST Lct (PrimState m) a
lct Vertex
c Vertex
p
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.link" Vertex
c (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.link" Vertex
p (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

{-# INLINEABLE cutST #-}
cutST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()
cutST :: forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
cutST lct :: Lct s a
lct@Lct {MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct, MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct} Vertex
u Vertex
v = do
  -- make @u@ the new root of the underlying tree
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct s a
lct Vertex
u
  -- make @v@ in the same preferred path as the root
  Vertex
_ <- Lct s a -> Vertex -> ST s Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct s a
lct Vertex
v

  -- dbgM $ do
  --   -- @v@ does not have any right children. because @u@ and @v@ are neighbors, @vl@ is @u@.
  --   vp <- VGM.unsafeRead pLct v
  --   let !_ = ACIA.runtimeAssert (nullLct vp) "vp must be null"
  --   vl <- VGM.unsafeRead lLct v
  --   let !_ = ACIA.runtimeAssert (vl == u) "vl must be `u`"
  --   pure ()

  -- do
  --   -- @v@ does not have any right children. because @u@ and @v@ are neighbors, @vl@ is @u@.
  --   vp <- VGM.unsafeRead pLct v
  --   vl <- VGM.unsafeRead lLct v
  --   let !_ = if nullLct vp then () else error "vp must be null"
  --   let !_ = if vl == u then () else error "vl must be `u`"
  --   pure ()

  -- delete the heavy edge.
  -- vl <- VGM.unsafeRead lLct v
  -- VGM.unsafeWrite pLct vl undefLct
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
u Vertex
undefLct
  MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
v Vertex
undefLct
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct s a
lct Vertex
v

-- | Amortized \(O(\log N)\). Deletes an edge between \(u\) and \(v\).
--
-- @since 1.1.1.0
{-# INLINE cut #-}
cut :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
cut :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m ()
cut Lct (PrimState m) a
lct Vertex
u Vertex
v = 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
$ Lct (PrimState m) a -> Vertex -> Vertex -> ST (PrimState m) ()
forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
cutST Lct (PrimState m) a
lct Vertex
u Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.cut" Vertex
u (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.cut" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | Amortized \(O(\log n)\). Makes \(v\) a new root of the underlying tree.
{-# INLINEABLE evertST #-}
evertST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s ()
evertST :: forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct s a
lct Vertex
v = do
  -- make @v@ be in the same preferred path as root. note that @v@ is at the root of the auxiliary tree.
  Vertex
_ <- Lct s a -> Vertex -> ST s Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct s a
lct Vertex
v
  -- reverse all the edges with respect to @v@: make @v@ a new root of the auxiliary tree.
  Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
reverseNodeST Lct s a
lct Vertex
v
  Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
v

-- | Amortized \(O(\log n)\). Makes \(v\) a new root of the underlying tree.
--
-- @since 1.1.1.0
{-# INLINE evert #-}
evert :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
evert :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m ()
evert Lct (PrimState m) a
lct Vertex
v = 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
$ Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.evert" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

{-# INLINEABLE exposeST #-}
exposeST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST :: forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST lct :: Lct s a
lct@Lct {MVector s Vertex
pLct :: forall s a. Lct s a -> MVector s Vertex
pLct :: MVector s Vertex
pLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct} Vertex
v0 = do
  let inner :: Vertex -> Vertex -> ST s Vertex
inner Vertex
v Vertex
lastRoot
        | Vertex -> Bool
nullLct Vertex
v = Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
lastRoot
        | Bool
otherwise = do
            -- go up to the top of the auxiliary tree:
            Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
splayST Lct s a
lct Vertex
v

            -- make @lastRoot@ the right child of @v@:
            --    v               v
            --   /|\        ->   /|\
            --    | r             | lastRoot  <-- @v0@ (in the @lastRoot@) will be connected to the root
            --    lastRoot        r
            Vertex
r <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
v
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
r) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Lct s a -> Vertex -> Vertex -> ST s ()
forall a s.
(Semigroup a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
addLightST Lct s a
lct Vertex
v Vertex
r
            Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Vertex -> Bool
nullLct Vertex
lastRoot) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Lct s a -> Vertex -> Vertex -> ST s ()
forall a s.
(Semigroup a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
eraseLightST Lct s a
lct Vertex
v Vertex
lastRoot
            MVector (PrimState (ST s)) Vertex -> Vertex -> Vertex -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> a -> m ()
VGM.unsafeWrite MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
v Vertex
lastRoot
            Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
updateNodeST Lct s a
lct Vertex
v

            -- go up to the next auxiliary tree:
            --    p
            --    |
            --    v
            --     \
            --      lastRoot
            Vertex
vp <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
pLct Vertex
v
            Vertex -> Vertex -> ST s Vertex
inner Vertex
vp Vertex
v

  Vertex
res <- Vertex -> Vertex -> ST s Vertex
inner Vertex
v0 Vertex
undefLct

  -- do
  --   -- FIXME: remove
  --   pRes <- VGM.unsafeRead pLct res
  --   unless (nullLct pRes) $ error $ "xxx must be null!!! " ++ show (res, pRes)

  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
splayST Lct s a
lct Vertex
v0

  -- do
  --   -- FIXME: remove
  --   p <- VGM.unsafeRead pLct v0
  --   unless (nullLct p) $ error $ "must be null!!! " ++ show (res, v0, p)

  Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
res

-- | Amortized \(O(\log n)\). Makes \(v\) and the root to be in the same preferred path (auxiliary
-- tree). After the opeartion, \(v\) will be the new root and all the children will be detached from
-- the preferred path.
--
-- @since 1.1.1.0
{-# INLINE expose #-}
expose :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m Vertex
expose :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m Vertex
expose Lct (PrimState m) a
lct Vertex
v = ST (PrimState m) Vertex -> m Vertex
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Vertex -> m Vertex)
-> ST (PrimState m) Vertex -> m Vertex
forall a b. (a -> b) -> a -> b
$ Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.expose_" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | Amortized \(O(\log n)\). `expose` with the return value discarded.
--
-- @since 1.1.1.0
{-# INLINE expose_ #-}
expose_ :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
expose_ :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m ()
expose_ Lct (PrimState m) a
lct Vertex
v0 = 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
  Vertex
_ <- Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
v0
  () -> ST (PrimState m) ()
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.expose_" Vertex
v0 (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- -------------------------------------------------------------------------------------------------
-- Jumo, LCA
-- -------------------------------------------------------------------------------------------------

-- | \(O(\log n)\) Returns the root of the underlying tree. Two vertices in the same connected
-- component have the same root vertex.
--
-- @since 1.1.1.0
{-# INLINE root #-}
root :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> m Vertex
root :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m Vertex
root lct :: Lct (PrimState m) a
lct@Lct {MVector (PrimState m) Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector (PrimState m) Vertex
lLct} Vertex
c0 = ST (PrimState m) Vertex -> m Vertex
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Vertex -> m Vertex)
-> ST (PrimState m) Vertex -> m Vertex
forall a b. (a -> b) -> a -> b
$ do
  Vertex
_ <- Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
c0
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct (PrimState m) a
lct Vertex
c0
  let inner :: Vertex -> ST (PrimState m) Vertex
inner Vertex
c = do
        Vertex
cl <- MVector (PrimState (ST (PrimState m))) Vertex
-> Vertex -> ST (PrimState m) Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) Vertex
MVector (PrimState (ST (PrimState m))) Vertex
lLct Vertex
c
        if Vertex -> Bool
nullLct Vertex
cl
          then Vertex -> ST (PrimState m) Vertex
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
c
          else do
            Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct (PrimState m) a
lct Vertex
cl
            Vertex -> ST (PrimState m) Vertex
inner Vertex
cl
  Vertex
c' <- Vertex -> ST (PrimState m) Vertex
inner Vertex
c0
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
splayST Lct (PrimState m) a
lct Vertex
c'
  Vertex -> ST (PrimState m) Vertex
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
c'
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.root" Vertex
c0 (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | \(O(\log n)\) Returns the parent vertex in the underlying tree.
--
-- @since 1.1.1.0
{-# INLINE parent #-}
parent :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> m (Maybe Vertex)
parent :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m (Maybe Vertex)
parent lct :: Lct (PrimState m) a
lct@Lct {MVector (PrimState m) Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector (PrimState m) Vertex
lLct, MVector (PrimState m) Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector (PrimState m) Vertex
rLct} Vertex
x = ST (PrimState m) (Maybe Vertex) -> m (Maybe Vertex)
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) (Maybe Vertex) -> m (Maybe Vertex))
-> ST (PrimState m) (Maybe Vertex) -> m (Maybe Vertex)
forall a b. (a -> b) -> a -> b
$ do
  Vertex
_ <- Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
x
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct (PrimState m) a
lct Vertex
x
  Vertex
xl <- MVector (PrimState (ST (PrimState m))) Vertex
-> Vertex -> ST (PrimState m) Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) Vertex
MVector (PrimState (ST (PrimState m))) Vertex
lLct Vertex
x
  if Vertex -> Bool
nullLct Vertex
xl
    then Maybe Vertex -> ST (PrimState m) (Maybe Vertex)
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Vertex
forall a. Maybe a
Nothing
    else do
      Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct (PrimState m) a
lct Vertex
xl
      let inner :: Vertex -> ST (PrimState m) Vertex
inner Vertex
y = do
            Vertex
yr <- MVector (PrimState (ST (PrimState m))) Vertex
-> Vertex -> ST (PrimState m) Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) Vertex
MVector (PrimState (ST (PrimState m))) Vertex
rLct Vertex
y
            if Vertex -> Bool
nullLct Vertex
yr
              then Vertex -> ST (PrimState m) Vertex
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
y
              else do
                Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct (PrimState m) a
lct Vertex
yr
                Vertex -> ST (PrimState m) Vertex
inner Vertex
yr
      Vertex -> Maybe Vertex
forall a. a -> Maybe a
Just (Vertex -> Maybe Vertex)
-> ST (PrimState m) Vertex -> ST (PrimState m) (Maybe Vertex)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vertex -> ST (PrimState m) Vertex
inner Vertex
xl
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.parent" Vertex
x (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

{-# INLINEABLE jumpST #-}
jumpST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> Int -> ST s Vertex
jumpST :: forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> Vertex -> ST s Vertex
jumpST lct :: Lct s a
lct@Lct {MVector s Vertex
lLct :: forall s a. Lct s a -> MVector s Vertex
lLct :: MVector s Vertex
lLct, MVector s Vertex
rLct :: forall s a. Lct s a -> MVector s Vertex
rLct :: MVector s Vertex
rLct, MVector s Vertex
sLct :: forall s a. Lct s a -> MVector s Vertex
sLct :: MVector s Vertex
sLct} Vertex
u0 Vertex
v0 Vertex
k0 = do
  -- make @v0@ a new root of the underlying tree
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct s a
lct Vertex
v0
  -- make @u0@ in the same preferred path as the root (@v0)
  Vertex
_ <- Lct s a -> Vertex -> ST s Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct s a
lct Vertex
u0

  do
    Vertex
size <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
sLct Vertex
u0
    let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vertex
0 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
<= Vertex
k0 Bool -> Bool -> Bool
&& Vertex
k0 Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< Vertex
size) String
"invalid jump"
    () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

  let inner :: Vertex -> Vertex -> ST s Vertex
inner Vertex
k Vertex
u = do
        Lct s a -> Vertex -> ST s ()
forall a s. Unbox a => Lct s a -> Vertex -> ST s ()
pushNodeST Lct s a
lct Vertex
u
        -- TODO: what is happening?
        Vertex
ur <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
rLct Vertex
u
        Vertex
urSize <- if Vertex -> Bool
nullLct Vertex
ur then Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
0 else MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
sLct Vertex
ur
        case Vertex -> Vertex -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Vertex
k Vertex
urSize of
          Ordering
LT -> Vertex -> Vertex -> ST s Vertex
inner Vertex
k Vertex
ur
          Ordering
EQ -> Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
u
          Ordering
GT -> do
            Vertex
ul <- MVector (PrimState (ST s)) Vertex -> Vertex -> ST s Vertex
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector s Vertex
MVector (PrimState (ST s)) Vertex
lLct Vertex
u
            Vertex -> Vertex -> ST s Vertex
inner (Vertex
k Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- (Vertex
urSize Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1)) Vertex
ul

  Vertex
res <- Vertex -> Vertex -> ST s Vertex
inner Vertex
k0 Vertex
u0
  Lct s a -> Vertex -> ST s ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
splayST Lct s a
lct Vertex
res
  Vertex -> ST s Vertex
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vertex
res

-- | \(O(\log n)\) Given a path between \(u\) and \(v\), returns the \(k\)-th vertex of the path.
--
-- ==== Constraints
-- - The \(k\)-th vertex must exist.
--
-- @since 1.1.1.0
{-# INLINE jump #-}
jump :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> Int -> m Vertex
jump :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> Vertex -> m Vertex
jump Lct (PrimState m) a
lct Vertex
u Vertex
v Vertex
k = ST (PrimState m) Vertex -> m Vertex
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Vertex -> m Vertex)
-> ST (PrimState m) Vertex -> m Vertex
forall a b. (a -> b) -> a -> b
$ Lct (PrimState m) a
-> Vertex -> Vertex -> Vertex -> ST (PrimState m) Vertex
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> Vertex -> ST s Vertex
jumpST Lct (PrimState m) a
lct Vertex
u Vertex
v Vertex
k

-- | \(O(\log n)\) Returns the LCA of \(u\) and \(v\). Because the root of the underlying changes
-- in almost every operation, one might want to use `evert` beforehand.
--
-- ==== Constraints
-- - \(u\) and \(v\) must be in the same connected component.
--
-- @since 1.1.1.0
{-# INLINE lca #-}
lca :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Int -> Int -> m Vertex
lca :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m Vertex
lca Lct (PrimState m) a
lct Vertex
u Vertex
v = ST (PrimState m) Vertex -> m Vertex
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) Vertex -> m Vertex)
-> ST (PrimState m) Vertex -> m Vertex
forall a b. (a -> b) -> a -> b
$ do
  Vertex
ru <- Lct (PrimState (ST (PrimState m))) a
-> Vertex -> ST (PrimState m) Vertex
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m Vertex
root Lct (PrimState m) a
Lct (PrimState (ST (PrimState m))) a
lct Vertex
u
  Vertex
rv <- Lct (PrimState (ST (PrimState m))) a
-> Vertex -> ST (PrimState m) Vertex
forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> m Vertex
root Lct (PrimState m) a
Lct (PrimState (ST (PrimState m))) a
lct Vertex
v
  let !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Vertex
ru Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
rv) (String -> ()) -> String -> ()
forall a b. (a -> b) -> a -> b
$ String
"AtCoder.Extra.Lct.lca: given two vertices in different connected components " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Vertex, Vertex) -> String
forall a. Show a => a -> String
show (Vertex
u, Vertex
v)
  Vertex
_ <- Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
u
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
v

-- -------------------------------------------------------------------------------------------------
-- Monoid produ\t
-- -------------------------------------------------------------------------------------------------

-- | Amortized \(O(\log n)\). Folds a path between \(u\) and \(v\) (inclusive).
--
-- @since 1.1.1.0
{-# INLINE prodPath #-}
prodPath :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m a
prodPath :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m a
prodPath lct :: Lct (PrimState m) a
lct@Lct {MVector (PrimState m) a
prodLct :: forall s a. Lct s a -> MVector s a
prodLct :: MVector (PrimState m) a
prodLct} Vertex
u Vertex
v = 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
  -- make @u@ the root of the underlying tree
  Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
u
  -- make @v@ in the same preferred path as @u@
  Vertex
_ <- Lct (PrimState m) a -> Vertex -> ST (PrimState m) Vertex
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s Vertex
exposeST Lct (PrimState m) a
lct Vertex
v
  -- now that @v@ is at the root of the auxiliary tree, its aggregation value is the path folding:
  MVector (PrimState (ST (PrimState m))) a
-> Vertex -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
prodLct Vertex
v
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.prodPath" Vertex
u (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.prodPath" Vertex
v (Lct (PrimState m) a -> Vertex
forall s a. Lct s a -> Vertex
nLct Lct (PrimState m) a
lct)

-- | Amortized \(O(\log n)\). Fold the subtree under \(v\), considering \(p\) as the root-side
-- vertex. Or, if \(p\) equals to \(v\), \(v\) will be the new root.
--
-- ==== Constraints
-- - The inverse operator has to be set on consturction (`newInv` or `buildInv`).
--
-- @since 1.1.1.0
{-# INLINE prodSubtree #-}
prodSubtree ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | Link/cut tree
  Lct (PrimState m) a ->
  -- | Vertex
  Vertex ->
  -- | Root or parent
  Vertex ->
  -- | Subtree's monoid product
  m a
prodSubtree :: forall (m :: * -> *) a.
(HasCallStack, PrimMonad m, Monoid a, Unbox a) =>
Lct (PrimState m) a -> Vertex -> Vertex -> m a
prodSubtree lct :: Lct (PrimState m) a
lct@Lct {Vertex
nLct :: forall s a. Lct s a -> Vertex
nLct :: Vertex
nLct, MVector (PrimState m) a
subtreeProdLct :: forall s a. Lct s a -> MVector s a
subtreeProdLct :: MVector (PrimState m) a
subtreeProdLct} Vertex
v Vertex
rootOrParent = 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
  if Vertex
v Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
rootOrParent
    then do
      -- `v` will be the root
      Lct (PrimState m) a -> Vertex -> ST (PrimState m) ()
forall a s. (Monoid a, Unbox a) => Lct s a -> Vertex -> ST s ()
evertST Lct (PrimState m) a
lct Vertex
v
      MVector (PrimState (ST (PrimState m))) a
-> Vertex -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
subtreeProdLct Vertex
v
    else do
      -- @rootOrParent@ can be far. retrieve the adjacent vertex:
      Vertex
parent_ <- Lct (PrimState m) a
-> Vertex -> Vertex -> Vertex -> ST (PrimState m) Vertex
forall a s.
(HasCallStack, Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> Vertex -> ST s Vertex
jumpST Lct (PrimState m) a
lct Vertex
v Vertex
rootOrParent Vertex
1
      -- detach @v@ from the parent. now that it's the root of the subtree vertices, the aggregation
      -- value is the aggregation of all the subtree vertices.
      Lct (PrimState m) a -> Vertex -> Vertex -> ST (PrimState m) ()
forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
cutST Lct (PrimState m) a
lct Vertex
v Vertex
parent_
      a
res <- MVector (PrimState (ST (PrimState m))) a
-> Vertex -> ST (PrimState m) a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Vertex -> m a
VGM.unsafeRead MVector (PrimState m) a
MVector (PrimState (ST (PrimState m))) a
subtreeProdLct Vertex
v
      -- attach again
      Lct (PrimState m) a -> Vertex -> Vertex -> ST (PrimState m) ()
forall a s.
(Monoid a, Unbox a) =>
Lct s a -> Vertex -> Vertex -> ST s ()
linkST Lct (PrimState m) a
lct Vertex
v Vertex
parent_
      a -> ST (PrimState m) a
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res
  where
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.prodSubtree" Vertex
v Vertex
nLct
    !()
_ = HasCallStack => String -> Vertex -> Vertex -> ()
String -> Vertex -> Vertex -> ()
ACIA.checkIndex String
"AtCoder.Extra.Lct.prodSubtree" Vertex
rootOrParent Vertex
nLct