content: {-# LANGUAGE RecordWildCards #-}
  
  -- original implementation:
  -- <https://github.com/maspypy/library/blob/main/ds/fastset.hpp>
  
  -- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).
  --
  -- ==== __Example__
  -- Create an `IntSet` with capacity \(10\):
  --
  -- >>> import AtCoder.Extra.IntSet qualified as IS
  -- >>> is <- IS.new @_ 10
  --
  -- `insert`, `delete`, `member`, `lookupGT` and other functions are available:
  --
  -- >>> IS.insert is 0
  -- >>> IS.insert is 9
  -- >>> IS.member is 9
  -- True
  --
  -- >>> IS.delete is 0
  -- True
  --
  -- >>> IS.size is
  -- 1
  --
  -- >>> IS.member is 1
  -- False
  --
  -- >>> IS.lookupGT is 5
  -- Just 9
  --
  -- @since 1.1.0.0
  module AtCoder.Extra.IntSet
    ( -- * IntSet
      IntSet,
  
      -- * Constructors
      new,
      build,
  
      -- * Metadata
      capacity,
      size,
      null,
  
      -- * Lookups
      member,
      notMember,
  
      -- ** Comparative lookups
      lookupGE,
      lookupGT,
      lookupLE,
      lookupLT,
  
      -- ** Max/Min lookups
      lookupMin,
      lookupMax,
  
      -- * Modifications
  
      -- ** Insertions
      insert,
  
      -- ** Deletions
      delete,
      delete_,
      deleteMin,
      deleteMax,
  
      -- * Conversions
      keys,
    )
  where
  
  import AtCoder.Internal.Assert qualified as ACIA
  import Control.Monad (unless, void)
  import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)
  import Control.Monad.ST (ST)
  import Data.Bifunctor (bimap)
  import Data.Bits
    ( Bits (clearBit, setBit, testBit),
      FiniteBits (countLeadingZeros, countTrailingZeros),
      (.<<.),
      (.>>.),
    )
  import Data.Maybe (fromJust)
  import Data.Vector qualified as V
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Generic.Mutable qualified as VGM
  import Data.Vector.Unboxed qualified as VU
  import Data.Vector.Unboxed.Mutable qualified as VUM
  import GHC.Stack (HasCallStack)
  import Prelude hiding (null)
  
  -- | \(O(1)\) Retrieves the most significant bit.
  --
  -- >>> msbOf 0
  -- -1
  --
  -- >>> msbOf maxBound
  -- 62
  --
  -- >>> msbOf $ 4 + 2 + 1
  -- 2
  {-# INLINE msbOf #-}
  msbOf :: Int -> Int
  msbOf !x = 63 - countLeadingZeros x
  
  -- | \(O(1)\) Retrieves the least significant bit.
  --
  -- >>> lsbOf 0
  -- -1
  --
  -- >>> lsbOf maxBound
  -- 0
  --
  -- >>> lsbOf $ 4 + 2 + 1
  -- 0
  {-# INLINE lsbOf #-}
  lsbOf :: Int -> Int
  lsbOf 0 = -1
  lsbOf x = countTrailingZeros x
  
  {-# INLINE wordSize #-}
  wordSize :: Int
  wordSize = 64
  
  -- | A dense, fast `Int` set implemented as a 64-ary tree that covers an interval \([0, n)\).
  --
  -- @since 1.1.0.0
  data IntSet s = IntSet
    { -- | Maximum number of elements.
      capacityIS :: {-# UNPACK #-} !Int,
      -- | The number of elements.
      sizeIS :: !(VUM.MVector s Int),
      -- | Segments.
      vecIS :: !(V.Vector (VUM.MVector s Int))
    }
  
  -- | \(O(n)\) Creates an `IntSet` for the interval \([0, n)\).
  --
  -- @since 1.1.0.0
  {-# INLINE new #-}
  new :: (PrimMonad m) => Int -> m (IntSet (PrimState m))
  new capacityIS = stToPrim $ newST capacityIS
  
  -- | \(O(n + m \log n)\) Creates an `IntSet` for the interval \([0, n)\) with initial values.
  --
  -- @since 1.1.0.0
  {-# INLINE build #-}
  build :: (PrimMonad m) => Int -> VU.Vector Int -> m (IntSet (PrimState m))
  build n vs = stToPrim $ buildST n vs
  
  -- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the set.
  --
  -- @since 1.1.0.0
  {-# INLINE capacity #-}
  capacity :: IntSet s -> Int
  capacity = capacityIS
  
  -- | \(O(1)\) Returns the number of elements in the set.
  --
  -- @since 1.1.0.0
  {-# INLINE size #-}
  size :: (PrimMonad m) => IntSet (PrimState m) -> m Int
  size = (`VUM.unsafeRead` 0) . sizeIS
  
  -- | \(O(1)\) Returns whether the set is empty.
  --
  -- @since 1.1.0.0
  {-# INLINE null #-}
  null :: (PrimMonad m) => IntSet (PrimState m) -> m Bool
  null = ((== 0) <$>) . size
  
  -- | \(O(\log n)\) Tests whether \(k\) is in the set.
  --
  -- @since 1.1.0.0
  {-# INLINE member #-}
  member :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
  member is k = stToPrim $ memberST is k
  
  -- | \(O(\log n)\) Tests whether \(k\) is not in the set.
  --
  -- @since 1.1.0.0
  {-# INLINE notMember #-}
  notMember :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
  notMember dis k = stToPrim $ not <$> memberST dis k
  
  -- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \ge k_0\).
  --
  -- @since 1.1.0.0
  {-# INLINE lookupGE #-}
  lookupGE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
  lookupGE is i0 = stToPrim $ lookupGEST is i0
  
  -- | \(O(\log n)\) Looks up the smallest key \(k\) such that \(k \gt k_0\).
  --
  -- @since 1.1.0.0
  {-# INLINE lookupGT #-}
  lookupGT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
  lookupGT is k = stToPrim $ lookupGTST is k
  
  -- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \le k_0\).
  --
  -- @since 1.1.0.0
  {-# INLINE lookupLE #-}
  lookupLE :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
  lookupLE is i0 = stToPrim $ lookupLEST is i0
  
  -- | \(O(\log n)\) Looks up the largest key \(k\) such that \(k \lt k_0\).
  --
  -- @since 1.1.0.0
  {-# INLINE lookupLT #-}
  lookupLT :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m (Maybe Int)
  lookupLT is k = stToPrim $ lookupLTST is k
  
  -- | \(O(\log n)\) Looks up the minimum key.
  --
  -- @since 1.1.0.0
  {-# INLINE lookupMin #-}
  lookupMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
  lookupMin is = stToPrim $ lookupMinST is
  
  -- | \(O(\log n)\) Looks up the maximum key.
  --
  -- @since 1.1.0.0
  {-# INLINE lookupMax #-}
  lookupMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
  lookupMax is = stToPrim $ lookupMaxST is
  
  -- | \(O(\log n)\) Inserts a key \(k\) into the set. If an entry with the same key already exists,
  -- it is overwritten.
  --
  -- @since 1.1.0.0
  {-# INLINE insert #-}
  insert :: (HasCallStack, PrimMonad m) => IntSet (PrimState m) -> Int -> m ()
  insert is k = stToPrim $ insertST is k
  
  -- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists. Returns
  -- whether the key existed.
  --
  -- @since 1.1.0.0
  {-# INLINE delete #-}
  delete :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool
  delete is k = stToPrim $ deleteST is k
  
  -- | \(O(\log n)\) Deletes a key \(k\) from the set. Does nothing if no such key exists.
  --
  -- @since 1.1.0.0
  {-# INLINE delete_ #-}
  delete_ :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m ()
  delete_ is k = stToPrim $ deleteST_ is k
  
  -- | \(O(\log n)\) Deletes the minimum key from the set. Returns `Nothing` if the set is empty.
  --
  -- @since 1.1.0.0
  {-# INLINE deleteMin #-}
  deleteMin :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
  deleteMin is = stToPrim $ deleteMinST is
  
  -- | \(O(\log n)\) Deletes the maximum key from the set. Returns `Nothing` if the set is empty.
  --
  -- @since 1.1.0.0
  {-# INLINE deleteMax #-}
  deleteMax :: (PrimMonad m) => IntSet (PrimState m) -> m (Maybe Int)
  deleteMax is = stToPrim $ deleteMaxST is
  
  -- | \(O(n \log n)\) Enumerates the keys in the map.
  --
  -- @since 1.1.0.0
  {-# INLINE keys #-}
  keys :: (PrimMonad m) => IntSet (PrimState m) -> m (VU.Vector Int)
  keys is = stToPrim $ keysST is
  
  -- -------------------------------------------------------------------------------------------------
  -- Internal
  -- -------------------------------------------------------------------------------------------------
  
  {-# INLINEABLE newST #-}
  newST :: Int -> ST s (IntSet s)
  newST capacityIS = do
    vecIS <-
      V.unfoldrExactNM
        (max 1 logSize)
        ( \len -> do
            let !len' = (len + wordSize - 1) `div` wordSize
            (,len') <$> VUM.replicate len' 0
        )
        capacityIS
    sizeIS <- VUM.replicate 1 (0 :: Int)
    pure IntSet {..}
    where
      (!_, !logSize) =
        until
          ((<= 1) . fst)
          (bimap ((`div` wordSize) . (+ (wordSize - 1))) (+ 1))
          (capacityIS, 0)
  
  {-# INLINEABLE buildST #-}
  buildST :: Int -> VU.Vector Int -> ST s (IntSet s)
  buildST n vs = do
    set <- newST n
    VU.forM_ vs (insertST set)
    pure set
  
  {-# INLINEABLE memberST #-}
  memberST :: IntSet s -> Int -> ST s Bool
  memberST IntSet {..} k
    | ACIA.testIndex k capacityIS = do
        let (!q, !r) = k `divMod` wordSize
        (`testBit` r) <$> VGM.unsafeRead (VG.unsafeHead vecIS) q
    | otherwise = pure False
  
  {-# INLINEABLE lookupGEST #-}
  lookupGEST :: IntSet s -> Int -> ST s (Maybe Int)
  lookupGEST IntSet {..} i0
    | i0 >= capacityIS = pure Nothing
    | otherwise = inner 0 $ max 0 i0 -- REMARK: it's very important to keep @i@ non-negative.
    where
      inner h i
        | h >= V.length vecIS = pure Nothing
        -- ?
        | q == VUM.length (VG.unsafeIndex vecIS h) = pure Nothing
        | otherwise = do
            d <- (.>>. r) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q
            if d == 0
              then inner (h + 1) (q + 1)
              else
                Just
                  <$> V.foldM'
                    ( \ !acc vec -> do
                        !dx <- lsbOf <$> VGM.unsafeRead vec acc
                        pure $ acc * wordSize + dx
                    )
                    (i + lsbOf d)
                    (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h))
        where
          (!q, !r) = i `divMod` wordSize
  
  {-# INLINEABLE lookupGTST #-}
  lookupGTST :: IntSet s -> Int -> ST s (Maybe Int)
  lookupGTST is k = lookupGEST is (k + 1)
  
  {-# INLINEABLE lookupLEST #-}
  lookupLEST :: IntSet s -> Int -> ST s (Maybe Int)
  lookupLEST IntSet {..} i0
    | i0 <= -1 = pure Nothing
    | otherwise = inner 0 $ min (capacityIS - 1) i0
    where
      inner h i
        | h >= V.length vecIS = pure Nothing
        | i == -1 = pure Nothing
        | otherwise = do
            d <- (.<<. (63 - r)) <$> VGM.unsafeRead (VG.unsafeIndex vecIS h) q
            if d == 0
              then inner (h + 1) (q - 1)
              else do
                Just
                  <$> V.foldM'
                    ( \ !acc vec -> do
                        !dx <- msbOf <$> VGM.unsafeRead vec acc
                        pure $ acc * wordSize + dx
                    )
                    (i - countLeadingZeros d)
                    (V.unsafeBackpermute vecIS (V.enumFromStepN (h - 1) (-1) h))
        where
          (!q, !r) = i `divMod` wordSize
  
  {-# INLINEABLE lookupLTST #-}
  lookupLTST :: IntSet s -> Int -> ST s (Maybe Int)
  lookupLTST is k = lookupLEST is (k - 1)
  
  {-# INLINEABLE lookupMinST #-}
  lookupMinST :: IntSet s -> ST s (Maybe Int)
  lookupMinST is = lookupGEST is 0
  
  {-# INLINEABLE lookupMaxST #-}
  lookupMaxST :: IntSet s -> ST s (Maybe Int)
  lookupMaxST is = lookupLEST is (capacityIS is - 1)
  
  {-# INLINEABLE insertST #-}
  insertST :: (HasCallStack) => IntSet s -> Int -> ST s ()
  insertST is@IntSet {..} k = do
    b <- memberST is k
    unless b $ do
      VUM.unsafeModify sizeIS (+ 1) 0
      V.foldM'_
        ( \i vec -> do
            let (!q, !r) = i `divMod` wordSize
            VGM.unsafeModify vec (`setBit` r) q
            pure q
        )
        k
        vecIS
    where
      !_ = ACIA.checkIndex "AtCoder.Extra.IntSet.insert" k capacityIS
  
  {-# INLINEABLE deleteST #-}
  deleteST :: IntSet s -> Int -> ST s Bool
  deleteST is@IntSet {..} k = do
    b_ <- memberST is k
    if b_
      then do
        VUM.unsafeModify sizeIS (subtract 1) 0
        V.foldM'_
          ( \(!b, !i) vec -> do
              let (!q, !r) = i `divMod` wordSize
              -- TODO: early return is possible
              unless b $ do
                VGM.unsafeModify vec (`clearBit` r) q
              -- `b` remembers if any other bit was on
              b' <- (/= 0) <$> VGM.unsafeRead vec q
              pure (b', q)
          )
          (False, k)
          vecIS
        pure True
      else pure False
  
  {-# INLINEABLE deleteST_ #-}
  deleteST_ :: IntSet s -> Int -> ST s ()
  deleteST_ is k = void $ deleteST is k
  
  {-# INLINEABLE deleteMinST #-}
  deleteMinST :: IntSet s -> ST s (Maybe Int)
  deleteMinST is = do
    lookupMinST is
      >>= mapM
        ( \key -> do
            deleteST_ is key
            pure key
        )
  
  {-# INLINEABLE deleteMaxST #-}
  deleteMaxST :: IntSet s -> ST s (Maybe Int)
  deleteMaxST is = do
    lookupMaxST is
      >>= mapM
        ( \key -> do
            deleteST_ is key
            pure key
        )
  
  {-# INLINEABLE keysST #-}
  keysST :: IntSet s -> ST s (VU.Vector Int)
  keysST is@IntSet {sizeIS} = do
    n <- VGM.unsafeRead sizeIS 0
    VU.unfoldrExactNM
      n
      ( \i -> do
          i' <- fromJust <$> lookupGTST is i
          pure (i', i')
      )
      (-1)
  