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

-- | Link/cut tree: dynamic forest with monoid values on vertices. If you need to have monoid values
-- on edges, treat the original edges as new vertices.
--
-- - __Most operations are unsafe__; user must ensure connectivities of \(u\) and \(v\) before
--   running each query.
-- - This specific implementation is not capable of applying monoid action to a subtree.
--
-- ==== __Example__
--
-- Create a link/cut tree of monoid @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)]
--
-- ===== `prodPath`, `prodSubtree`, `prodTree`
--
-- Monoid products can be calculated for paths or subtrees:
--
-- >>> Lct.prodPath lct 0 2
-- Sum {getSum = 3}
--
-- >>> -- If we create the LCT with `buildInv`, we can use `prodSubtree`:
-- >>> 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
--
-- ===== `lca`, `jump`, `lengthBetween`
--
-- 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.jump lct {- path -} 2 3 {- k -} 2
-- 3
--
-- >>> Lct.jumpMaybe lct {- path -} 2 3 {- k -} 1000
-- Nothing
--
-- >>> Lct.lengthBetween lct {- path -} 2 3
-- 2
--
-- ===== `parent`
--
-- >>> Lct.evert lct 0  -- set root `0`
-- >>> Lct.parent lct 0 -- under root `0`, parent of `0` is `Nothing`:
-- Nothing
--
-- >>> Lct.evert lct 0  -- set root `0`
-- >>> Lct.parent lct 1 -- under root `0`, parent of `1` is `0`:
-- Just 0
--
-- ===== `link` / `cut`
--
-- 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,

    -- * Monoid value access
    read,
    write,
    modify,
    modifyM,

    -- * Tree operations

    -- ** Link/cut
    link,
    cut,

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

    -- * Tree queries
    root,
    same,
    parent,
    jump,
    jumpMaybe,
    lengthBetween,
    lca,
    lcaMaybe,

    -- ** Monoid products
    prodPath,
    prodSubtree,
    prodTree,
  )
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.Maybe
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)
import Prelude hiding (read)

-- import GHC.Stack (HasCallStack)

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

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

{-# INLINE nullLct #-}
nullLct :: Vertex -> Bool
nullLct = (== -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
    nLct :: {-# UNPACK #-} !Int,
    -- | Decomposed node data storage: left children.
    --
    -- @since 1.1.1.0
    lLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: right children.
    --
    -- @since 1.1.1.0
    rLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: parents.
    --
    -- @since 1.1.1.0
    pLct :: !(VUM.MVector s Vertex),
    -- | Decomposed node data storage: subtree sizes.
    --
    -- @since 1.1.1.0
    sLct :: !(VUM.MVector s Int),
    -- | Decomposed node data storage: reverse flags.
    --
    -- @since 1.1.1.0
    revLct :: !(VUM.MVector s Bit),
    -- | Decomposed node data storage: monoid values.
    --
    -- @since 1.1.1.0
    vLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: monoid products.
    --
    -- @since 1.1.1.0
    prodLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: dual monoid products (right foldings). This is required for
    -- non-commutative monoids only.
    --
    -- @since 1.1.1.0
    dualProdLct :: !(VUM.MVector s a),
    -- | Decomposed node data storage: path-parent monoid products. This works for subtree product
    -- queries over commutative monoids only.
    --
    -- @since 1.1.1.0
    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
    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
    invOpLct :: !(a -> a)
  }

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

-- | \(O(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 !invOpLct nLct = buildInv invOpLct (VU.replicate nLct mempty) VU.empty

-- | \(O(n + m \log n)\) Creates a link/cut tree of initial monoid values and initial edges. This
-- setup disables subtree queries (`prodSubtree`).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
--
-- @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 xs es = stToPrim $ buildInv id xs es

-- | \(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`).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
--
-- @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 invOpLct xs es = stToPrim $ buildST invOpLct xs es

-- -------------------------------------------------------------------------------------------------
-- Monoid value access
-- -------------------------------------------------------------------------------------------------

-- | \(O(1)\). Reads the monoid value on a vertex \(v\).
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.5.1.0
{-# INLINE read #-}
read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m a
read lct v = stToPrim $ do
  VGM.read (vLct lct) v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.read" v (nLct lct)

-- | Amortized \(O(\log n)\). Writes to the monoid value of a vertex \(v\).
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE write #-}
write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> a -> m ()
write lct v x = stToPrim $ do
  -- make @v@ a new root of the underlying tree:
  evertST lct v
  VGM.unsafeWrite (vLct lct) v x
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.write" v (nLct lct)

-- | Amortized \(O(\log n)\). Given a user function \(f\), modifies the monoid value of a vertex
-- \(v\).
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @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 lct f v = stToPrim $ do
  -- make @v@ a new root of the underlying tree:
  evertST lct v
  VGM.unsafeModify (vLct lct) f v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.modify" v (nLct lct)

-- | Amortized \(O(\log n)\). Given a user function \(f\), modifies the monoid value of a vertex
-- \(v\).
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @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 lct f v = do
  -- make @v@ a new root of the underlying tree:
  stToPrim $ evertST lct v
  VGM.unsafeModifyM (vLct lct) f v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.modifyM" v (nLct lct)

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

-- | Amortized \(O(\log n)\). Creates an edge between \(c\) and \(p\). In the represented tree, the
-- \(p\) will be the parent of \(c\).
--
-- ==== Constraints
-- - \(0 \le c, p \lt n\)
-- - \(u \ne v\)
-- - \(c\) and \(p\) are in different trees, otherwise the behavior is undefined.
--
-- @since 1.1.1.0
{-# INLINE link #-}
link :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
link lct c p = stToPrim $ linkST lct c p
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.link" c (nLct lct)
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.link" p (nLct lct)

-- | Amortized \(O(\log n)\). Deletes an edge between \(u\) and \(v\).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(u \ne v\)
-- - There's an edge between \(u\) and \(v\), otherwise the behavior is undefined.
--
-- @since 1.1.1.0
{-# INLINE cut #-}
cut :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m ()
cut lct u v = stToPrim $ cutST lct u v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.cut" u (nLct lct)
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.cut" v (nLct lct)

-- | Amortized \(O(\log n)\). Makes \(v\) a new root of the underlying tree.
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE evert #-}
evert :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
evert lct v = stToPrim $ evertST lct v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.evert" v (nLct lct)

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

-- | Amortized \(O(\log n)\). `expose` with the return value discarded.
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE expose_ #-}
expose_ :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m ()
expose_ lct v0 = stToPrim $ do
  _ <- exposeST lct v0
  pure ()

-- -------------------------------------------------------------------------------------------------
-- Jump, LCA
-- -------------------------------------------------------------------------------------------------

-- | \(O(\log n)\) Returns the root of the underlying tree.
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE root #-}
root :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m Vertex
root lct c0 = stToPrim $ rootST lct c0

-- | \(O(\log n)\) Returns whether the vertices \(u\) and \(v\) are in the same connected component
-- (have the same `root`).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
--
-- @since 1.5.1.0
{-# INLINE same #-}
same :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m Bool
same lct u v = stToPrim $ do
  r1 <- rootST lct u
  r2 <- rootST lct v
  pure $ r1 == r2

-- | \(O(\log n)\) Returns the parent vertex in the underlying tree.
--
-- ==== Constraints
-- - \(0 \le v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE parent #-}
parent :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> m (Maybe Vertex)
parent lct v = stToPrim $ parentST lct v

-- | \(O(\log n)\) Given a path between \(u\) and \(v\), returns the \(k\)-th vertex from \(u\) in
-- the path.
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(0 \le k \lt \mathrm{|path|}\)
-- - \(u\) and \(v\) must be in the same connected component, otherwise the behavior is undefined.
--
-- @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 lct u v k = stToPrim $ do
  fromMaybe (error "AtCoder.Extra.Tree.Lct.jump: invalid jump") <$> jumpMaybeST lct u v k

-- | \(O(\log n)\) Given a path between \(u\) and \(v\), returns the \(k\)-th vertex from \(u\) in
-- the path.
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(u\) and \(v\) must be in the same connected component, otherwise the behavior is undefined.
--
-- @since 1.5.1.0
{-# INLINE jumpMaybe #-}
jumpMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> Int -> m (Maybe Vertex)
jumpMaybe lct u v k = stToPrim $ jumpMaybeST lct u v k

-- | \(O(\log n)\) Returns the length of path between \(u\) and \(v\).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(u\) and \(v\) must be in the same connected component.
--
-- @since 1.5.1.0
{-# INLINE lengthBetween #-}
lengthBetween :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m Vertex
lengthBetween lct u v = stToPrim $ do
  lengthBetweenST lct u v

-- | \(O(\log n)\) Returns the LCA of \(u\) and \(v\). Because the root of the underlying tree changes
-- in almost every operation, one might want to use `evert` beforehand.
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(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 -> Vertex -> Vertex -> m Vertex
lca lct u v = stToPrim $ do
  fromMaybe (error ("AtCoder.Extra.Tree.Lct.lca: given two vertices in different connected components " ++ show (u, v)))
    <$> lcaMaybe lct u v

-- | \(O(\log n)\) Returns the LCA of \(u\) and \(v\). Because the root of the underlying tree changes
-- in almost every operation, one might want to use `evert` beforehand.
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
--
-- @since 1.5.1.0
{-# INLINE lcaMaybe #-}
lcaMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Lct (PrimState m) a -> Vertex -> Vertex -> m (Maybe Vertex)
lcaMaybe lct u v = stToPrim $ do
  ru <- rootST lct u
  rv <- rootST lct v
  if ru == rv
    then do
      _ <- exposeST lct u
      Just <$> exposeST lct v
    else pure Nothing

-- -------------------------------------------------------------------------------------------------
-- Monoid product
-- -------------------------------------------------------------------------------------------------

-- | Amortized \(O(\log n)\). Folds a path between \(u\) and \(v\) (inclusive).
--
-- ==== Constraints
-- - \(0 \le u, v \lt n\)
-- - \(u\) and \(v\) must be in the same connected component, otherwise the return value is
--   nonsense.
--
-- @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 lct@Lct {prodLct} u v = stToPrim $ do
  -- make @u@ the root of the underlying tree
  evertST lct u
  -- make @v@ in the same preferred path as @u@
  _ <- exposeST lct v
  -- now that @v@ is at the root of the auxiliary tree, its aggregation value is the path folding:
  VGM.unsafeRead prodLct v
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.prodPath" u (nLct lct)
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.prodPath" v (nLct lct)

-- | Amortized \(O(\log n)\). Fold the subtree under \(v\), considering \(p\) as the root-side
-- vertex. Or, if \(p\) equals \(v\), \(v\) will be the new root.
--
-- ==== Constraints
-- - The inverse operator must be set on construction (`newInv` or `buildInv`).
-- - \(0 \le u, v \lt n\)
--
-- @since 1.1.1.0
{-# INLINE prodSubtree #-}
prodSubtree ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | Link/cut tree
  Lct (PrimState m) a ->
  -- | Vertex \(v\)
  Vertex ->
  -- | Parent \(p\) (not need be adjacent to \(v\)), or same as \(v\), making it a new root.
  Vertex ->
  -- | Subtree's monoid product
  m a
prodSubtree lct v rootOrParent = stToPrim $ prodSubtreeST lct v rootOrParent

-- | Amortized \(O(\log n)\). Fold a tree that contains \(v\).
--
-- ==== Constraints
-- - The inverse operator must be set on construction (`newInv` or `buildInv`).
-- - \(0 \le v \lt n\)
--
-- @since 1.5.1.0
{-# INLINE prodTree #-}
prodTree ::
  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>
  -- | Link/cut tree
  Lct (PrimState m) a ->
  -- | Vertex \(v\)
  Vertex ->
  -- | Subtree's monoid product
  m a
prodTree lct v = stToPrim $ prodSubtreeST lct v v

-- -------------------------------------------------------------------------------------------------
-- Internal
-- -------------------------------------------------------------------------------------------------

{-# INLINEABLE buildST #-}
buildST ::
  (HasCallStack, Monoid a, VU.Unbox a) =>
  -- | Inverse operator
  (a -> a) ->
  -- | Vertex monoid values
  VU.Vector a ->
  -- | Edges
  VU.Vector (Vertex, Vertex) ->
  -- | Link/cut tree
  ST s (Lct s a)
buildST invOpLct xs es = do
  lct <- do
    let !nLct = VU.length xs
    lLct <- VUM.replicate nLct undefLct
    rLct <- VUM.replicate nLct undefLct
    pLct <- VUM.replicate nLct undefLct
    sLct <- VUM.replicate nLct 0
    revLct <- VUM.replicate nLct (Bit False)
    vLct <- VU.thaw xs
    prodLct <- VUM.replicate nLct mempty
    dualProdLct <- VUM.replicate nLct mempty
    midLct <- VUM.replicate nLct mempty
    subtreeProdLct <- VUM.replicate nLct mempty
    pure Lct {..}
  VU.forM_ es $ \(!u, !v) -> do
    link lct u v
  pure 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 lct@Lct {pLct, lLct, rLct} v = do
  p <- VGM.unsafeRead pLct v
  pp <- VGM.unsafeRead pLct p
  pl <- VGM.unsafeRead lLct p

  c <-
    if pl == v
      then do
        -- rotate right:
        --   p      v  <-- reference from `pp` is updated later
        --  /        \
        -- v    ->    p
        --  \        /
        --   c      c
        c <- VGM.unsafeExchange rLct v p
        VGM.unsafeWrite lLct p c
        pure c
      else do
        -- rotate left:
        -- p          v  <-- reference from `pp` is updated later
        --  \        /
        --   v  ->  p
        --  /        \
        -- c          c
        c <- VGM.unsafeExchange lLct v p
        VGM.unsafeWrite rLct p c
        pure c

  updateNodeST lct p
  updateNodeST lct v

  -- update the reference from `pp`:
  unless (nullLct pp) $ do
    ppl <- VGM.unsafeRead lLct pp
    if ppl == p
      then VGM.unsafeWrite lLct pp v
      else do
        ppr <- VGM.unsafeRead rLct pp
        if ppr == p
          then VGM.unsafeWrite rLct pp v
          else do
            -- overwrite the light (path-parent) pointer:
            changeLightST lct pp p v

  -- update parent pointers to `pp`: pp <-- v <-- p <-- c
  VGM.unsafeWrite pLct v pp
  VGM.unsafeWrite pLct p v
  unless (nullLct c) $ do
    VGM.unsafeWrite pLct c 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 lct@Lct {pLct} c = do
  pushNodeST lct c
  let inner = do
        isRootC <- isRootNodeST lct c
        unless isRootC $ do
          p <- VGM.unsafeRead pLct c
          pp <- if nullLct p then pure undefLct else VGM.unsafeRead pLct p
          placeP <- nodePlaceST lct p
          if placeP == RootNodeLct
            then do
              pushNodeST lct p
              pushNodeST lct c
              rotateST lct c
            else do
              placeC <- nodePlaceST lct c
              pushNodeST lct pp
              pushNodeST lct p
              pushNodeST lct c
              if placeC == 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

                  rotateST lct p
                  rotateST lct c
                else do
                  --       pp         pp         c
                  --      /          /          | \
                  --    p     ->   c      ->   p   pp
                  --     \        /
                  --      c      p
                  rotateST lct c
                  rotateST lct c
          inner
  inner

-- * Node helpers

-- | \(O(1)\)
{-# INLINE isRootNodeST #-}
isRootNodeST :: Lct s a -> Vertex -> ST s Bool
isRootNodeST lct v = do
  (== RootNodeLct) <$> nodePlaceST lct v

-- TODO: return heavy/light notion
data NodePlaceLct = RootNodeLct | LeftNodeLct | RightNodeLct
  deriving (Eq)

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

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

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

-- | \(O(1)\)
{-# INLINEABLE reverseNodeST #-}
reverseNodeST :: (VU.Unbox a) => Lct s a -> Vertex -> ST s ()
reverseNodeST lct@Lct {revLct} i = do
  -- lazily propagate new reverse from the children, or cancel:
  VGM.unsafeModify revLct (xor (Bit True)) i
  -- swap
  swapLrNodeST lct 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 Lct {lLct, rLct, prodLct, dualProdLct} i = do
  -- swap children
  VGM.unsafeModifyM lLct (VGM.unsafeExchange rLct i) i
  -- swap prodLct[i] and dualProdLct[i]
  VGM.unsafeModifyM prodLct (VGM.unsafeExchange dualProdLct i) 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 Lct {..} i = do
  l <- VGM.unsafeRead lLct i
  r <- VGM.unsafeRead rLct i
  v <- VGM.unsafeRead vLct i
  m <- VGM.unsafeRead midLct i

  (!size', !prod', !dualProd', !subtreeProd') <-
    if nullLct l
      then pure (1 :: Int, v, v, v <> m)
      else do
        lSize <- VGM.unsafeRead sLct l
        lProd <- VGM.unsafeRead prodLct l
        lDualProd <- VGM.unsafeRead dualProdLct l
        lSubtreeProd <- VGM.unsafeRead subtreeProdLct l
        pure (lSize + 1, lProd <> v, v <> lDualProd, lSubtreeProd <> v <> m)

  (!size'', !prod'', !dualProd'', !subtreeProd'') <-
    if nullLct r
      then pure (size', prod', dualProd', subtreeProd')
      else do
        rSize <- VGM.unsafeRead sLct r
        rProd <- VGM.unsafeRead prodLct r
        rDualProd <- VGM.unsafeRead dualProdLct r
        rSubtreeProd <- VGM.unsafeRead subtreeProdLct r
        pure (size' + rSize, prod' <> rProd, rDualProd <> dualProd', subtreeProd' <> rSubtreeProd)

  VGM.unsafeWrite sLct i size''
  VGM.unsafeWrite prodLct i prod''
  VGM.unsafeWrite dualProdLct i dualProd''
  VGM.unsafeWrite subtreeProdLct i 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 Lct {subtreeProdLct, midLct} p c = do
  newChild <- VGM.unsafeRead subtreeProdLct c
  VGM.unsafeModify midLct (newChild <>) 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 _lct _u _v _p = do
  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 Lct {subtreeProdLct, midLct, invOpLct} p c = do
  sub <- VGM.unsafeRead subtreeProdLct c
  let !sub' = invOpLct sub
  VGM.unsafeModify midLct (<> sub') p

-- * Link/cut

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

  -- no assertion and very unsafe..
  -- 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:
  VGM.unsafeWrite pLct c p
  VGM.unsafeWrite rLct p c
  updateNodeST lct p

-- no assertion and very unsafe..
{-# INLINEABLE cutST #-}
cutST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s ()
cutST lct@Lct {pLct, lLct} u v = do
  -- make @u@ the new root of the underlying tree
  evertST lct u
  -- make @v@ in the same preferred path as the root
  _ <- exposeST lct 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
  VGM.unsafeWrite pLct u undefLct
  VGM.unsafeWrite lLct v undefLct
  updateNodeST lct v

-- | 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 lct v = do
  -- make @v@ be in the same preferred path as root. note that @v@ is at the root of the auxiliary tree.
  _ <- exposeST lct v
  -- reverse all the edges with respect to @v@: make @v@ a new root of the auxiliary tree.
  reverseNodeST lct v
  pushNodeST lct v

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

            -- make @lastRoot@ the right child of @v@:
            --    v               v
            --   /|\        ->   /|\
            --    | r             | lastRoot  <-- @v0@ (in the @lastRoot@) will be connected to the root
            --    lastRoot        r
            r <- VGM.unsafeRead rLct v
            unless (nullLct r) $ addLightST lct v r
            unless (nullLct lastRoot) $ eraseLightST lct v lastRoot
            VGM.unsafeWrite rLct v lastRoot
            updateNodeST lct v

            -- go up to the next auxiliary tree:
            --    p
            --    |
            --    v
            --     \
            --      lastRoot
            vp <- VGM.unsafeRead pLct v
            inner vp v

  res <- inner v0 undefLct

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

  splayST lct v0

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

  pure res

-- * Jump, LCA

{-# INLINEABLE rootST #-}
rootST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Int -> ST s Vertex
rootST lct@Lct {lLct} c0 = do
  _ <- exposeST lct c0
  pushNodeST lct c0
  let inner c = do
        cl <- VGM.unsafeRead lLct c
        if nullLct cl
          then pure c
          else do
            pushNodeST lct cl
            inner cl
  c' <- inner c0
  splayST lct c'
  pure c'
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.rootST" c0 (nLct lct)

{-# INLINEABLE parentST #-}
parentST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Int -> ST s (Maybe Vertex)
parentST lct@Lct {lLct, rLct} x = do
  _ <- exposeST lct x
  pushNodeST lct x
  xl <- VGM.unsafeRead lLct x
  if nullLct xl
    then pure Nothing
    else do
      pushNodeST lct xl
      let inner y = do
            yr <- VGM.unsafeRead rLct y
            if nullLct yr
              then pure y
              else do
                pushNodeST lct yr
                inner yr
      Just <$> inner xl
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.parentST" x (nLct lct)

{-# INLINEABLE jumpMaybeST #-}
jumpMaybeST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> Int -> ST s (Maybe Vertex)
jumpMaybeST lct@Lct {lLct, rLct, sLct} u0 v0 k0 = do
  -- make @v0@ a new root of the underlying tree
  evertST lct v0
  -- make @u0@ in the same preferred path as the root (@v0)
  _ <- exposeST lct u0
  size <- VGM.unsafeRead sLct u0
  if 0 <= k0 && k0 < size
    then do
      let inner k u = do
            pushNodeST lct u
            -- TODO: what is happening?
            ur <- VGM.unsafeRead rLct u
            urSize <- if nullLct ur then pure 0 else VGM.unsafeRead sLct ur
            case compare k urSize of
              LT -> inner k ur
              EQ -> pure u
              GT -> do
                ul <- VGM.unsafeRead lLct u
                inner (k - (urSize + 1)) ul

      res <- inner k0 u0
      splayST lct res
      pure $ Just res
    else do
      pure Nothing

{-# INLINEABLE prodSubtreeST #-}
prodSubtreeST ::
  (HasCallStack, Monoid a, VU.Unbox a) =>
  -- | Link/cut tree
  Lct s a ->
  -- | Vertex
  Vertex ->
  -- | Parent \(p\) (not need be adjacent to \(v\)), or same as \(v\), making it a new root.
  Vertex ->
  -- | Subtree's monoid product
  ST s a
prodSubtreeST lct@Lct {nLct, subtreeProdLct} v rootOrParent = do
  if v == rootOrParent
    then do
      -- `v` will be the root
      evertST lct v
      VGM.unsafeRead subtreeProdLct v
    else do
      -- retrieve the adjacent parent:
      parent_ <- fromJust <$> jumpMaybeST lct v rootOrParent 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.
      cutST lct v parent_
      res <- VGM.unsafeRead subtreeProdLct v
      -- attach again
      linkST lct v parent_
      pure res
  where
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.prodSubtreeST" v nLct
    !_ = ACIA.checkIndex "AtCoder.Extra.Tree.Lct.prodSubtreeST" rootOrParent nLct

{-# INLINEABLE lengthBetweenST #-}
lengthBetweenST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s Int
lengthBetweenST lct@Lct {sLct} u0 v0 = do
  -- make @v0@ a new root of the underlying tree
  evertST lct v0
  -- make @u0@ in the same preferred path as the root (@v0)
  _ <- exposeST lct u0
  size <- VGM.unsafeRead sLct u0
  pure $ size - 1