packages feed

ac-library-hs 1.2.2.0 → 1.2.2.1

raw patch · 37 files changed

+2245/−1493 lines, 37 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

CHANGELOG.md view
@@ -1,5 +1,9 @@ # Revision history for acl-hs +## 1.2.2.1 -- March 2025++- Reduced build time with `ST` monad and `INLINEABLE` pragmas.+ ## 1.2.2.0 -- Feb 2025  - Added `Extra.KdTree` and `Extra.LazyKdTree`.
ac-library-hs.cabal view
@@ -4,7 +4,7 @@ -- PVP summary:  +-+------- breaking API changes --               | | +----- non-breaking API additions --               | | | +--- code changes with no API change-version:         1.2.2.0+version:         1.2.2.1 synopsis:        Data structures and algorithms description:   Haskell port of [ac-library](https://github.com/atcoder/ac-library), a library for competitive
src/AtCoder/Convolution.hs view
@@ -78,7 +78,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0-{-# INLINE convolution #-}+{-# INLINEABLE convolution #-} convolution ::   forall p.   (HasCallStack, AM.Modulus p) =>@@ -109,7 +109,7 @@ -- - \(O(n\log{n} + \log{\mathrm{mod}})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0-{-# INLINE convolutionRaw #-}+{-# INLINEABLE convolutionRaw #-} convolutionRaw ::   forall p a.   (HasCallStack, AM.Modulus p, Integral a, VU.Unbox a) =>@@ -139,7 +139,7 @@ -- - \(O(n\log{n})\), where \(n = |a| + |b|\). -- -- @since 1.0.0.0-{-# INLINE convolution64 #-}+{-# INLINEABLE convolution64 #-} convolution64 ::   (HasCallStack) =>   VU.Vector Int ->@@ -166,7 +166,6 @@           -- !_ = ACIA.runtimeAssert (mod2 `mod` bit maxAbBit == 1) $ "AtCoder.Convolution.convolution64: `mod2` isn't enough to support an array of length `2^25`."           -- !_ = ACIA.runtimeAssert (mod3 `mod` bit maxAbBit == 1) $ "AtCoder.Convolution.convolution64: `mod3` isn't enough to support an array of length `2^26`."           !_ = ACIA.runtimeAssert (n + m - 1 <= bit maxAbBit) "AtCoder.Convolution.convolution64: given too long vector as input"-          -- TODO: convolution vs convolutionRaw for the speed. I think the former is faster.           c1 = convolution {- mod1 -} (VU.map (AM.new @754974721) a) (VU.map (AM.new @754974721) b)           c2 = convolution {- mod2 -} (VU.map (AM.new @167772161) a) (VU.map (AM.new @167772161) b)           c3 = convolution {- mod3 -} (VU.map (AM.new @469762049) a) (VU.map (AM.new @469762049) b)
src/AtCoder/Dsu.hs view
@@ -65,8 +65,8 @@  import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (when)-import Control.Monad.ST (ST) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Vector qualified as V import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM@@ -98,11 +98,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (PrimMonad m) => Int -> m (Dsu (PrimState m))-new nDsu-  | nDsu >= 0 = do-      parentOrSizeDsu <- VUM.replicate nDsu (-1)-      pure Dsu {..}-  | otherwise = error $ "new: given negative size (`" ++ show nDsu ++ "`)"+new = stToPrim . newST  -- | Adds an edge \((a, b)\). If the vertices \(a\) and \(b\) are in the same connected component, it -- returns the representative (`leader`) of this connected component. Otherwise, it returns the@@ -118,22 +114,7 @@ -- @since 1.0.0.0 {-# INLINE merge #-} merge :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> Int -> m Int-merge dsu@Dsu {..} a b = stToPrim $ do-  let !_ = ACIA.checkVertex "AtCoder.Dsu.merge" a nDsu-  let !_ = ACIA.checkVertex "AtCoder.Dsu.merge" b nDsu-  x <- leaderST dsu a-  y <- leaderST dsu b-  if x == y-    then do-      pure x-    else do-      px <- VGM.read parentOrSizeDsu x-      py <- VGM.read parentOrSizeDsu y-      when (-px < -py) $ do-        VGM.swap parentOrSizeDsu x y-      sizeY <- VGM.exchange parentOrSizeDsu y x-      VGM.modify parentOrSizeDsu (+ sizeY) x-      pure x+merge dsu a b = stToPrim $ mergeST dsu a b  -- | `merge` with the return value discarded. --@@ -163,23 +144,7 @@ -- @since 1.0.0.0 {-# INLINE same #-} same :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> Int -> m Bool-same dsu@Dsu {..} a b = do-  let !_ = ACIA.checkVertex "AtCoder.Dsu.same" a nDsu-  let !_ = ACIA.checkVertex "AtCoder.Dsu.same" b nDsu-  la <- leader dsu a-  lb <- leader dsu b-  pure $ la == lb--{-# INLINE leaderST #-}-leaderST :: Dsu s -> Int -> ST s Int-leaderST dsu@Dsu {..} a = do-  pa <- VGM.read parentOrSizeDsu a-  if pa < 0-    then pure a-    else do-      lpa <- leaderST dsu pa-      VGM.write parentOrSizeDsu a lpa-      pure lpa+same dsu a b = stToPrim $ sameST dsu a b  -- | Returns the representative of the connected component that contains the vertex \(a\). --@@ -205,11 +170,7 @@ -- @since 1.0.0.0 {-# INLINE size #-} size :: (HasCallStack, PrimMonad m) => Dsu (PrimState m) -> Int -> m Int-size dsu@Dsu {..} a = stToPrim $ do-  let !_ = ACIA.checkVertex "AtCoder.Dsu.size" a nDsu-  la <- leaderST dsu a-  sizeLa <- VGM.read parentOrSizeDsu la-  pure (-sizeLa)+size dsu a = stToPrim $ sizeST dsu a  -- | Divides the graph into connected components and returns the vector of them. --@@ -222,7 +183,71 @@ -- @since 1.0.0.0 {-# INLINE groups #-} groups :: (PrimMonad m) => Dsu (PrimState m) -> m (V.Vector (VU.Vector Int))-groups dsu@Dsu {..} = stToPrim $ do+groups = stToPrim . groupsST++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: Int -> ST s (Dsu s)+newST nDsu+  | nDsu >= 0 = do+      parentOrSizeDsu <- VUM.replicate nDsu (-1)+      pure Dsu {..}+  | otherwise = error $ "AtCoder.Dsu.newST: given negative size (`" ++ show nDsu ++ "`)"++{-# INLINEABLE mergeST #-}+mergeST :: (HasCallStack) => Dsu s -> Int -> Int -> ST s Int+mergeST dsu@Dsu {..} a b = do+  let !_ = ACIA.checkVertex "AtCoder.Dsu.mergeST" a nDsu+  let !_ = ACIA.checkVertex "AtCoder.Dsu.mergeST" b nDsu+  x <- leaderST dsu a+  y <- leaderST dsu b+  if x == y+    then do+      pure x+    else do+      px <- VGM.read parentOrSizeDsu x+      py <- VGM.read parentOrSizeDsu y+      when (-px < -py) $ do+        VGM.swap parentOrSizeDsu x y+      sizeY <- VGM.exchange parentOrSizeDsu y x+      VGM.modify parentOrSizeDsu (+ sizeY) x+      pure x++{-# INLINEABLE sameST #-}+sameST :: (HasCallStack) => Dsu s -> Int -> Int -> ST s Bool+sameST dsu@Dsu {..} a b = do+  let !_ = ACIA.checkVertex "AtCoder.Dsu.sameST" a nDsu+  let !_ = ACIA.checkVertex "AtCoder.Dsu.sameST" b nDsu+  la <- leaderST dsu a+  lb <- leaderST dsu b+  pure $ la == lb++-- TODO: INLINE?+{-# INLINEABLE leaderST #-}+leaderST :: Dsu s -> Int -> ST s Int+leaderST dsu@Dsu {..} a = do+  pa <- VGM.read parentOrSizeDsu a+  if pa < 0+    then pure a+    else do+      lpa <- leaderST dsu pa+      VGM.write parentOrSizeDsu a lpa+      pure lpa++{-# INLINEABLE sizeST #-}+sizeST :: (HasCallStack) => Dsu s -> Int -> ST s Int+sizeST dsu@Dsu {..} a = do+  let !_ = ACIA.checkVertex "AtCoder.Dsu.sizeST" a nDsu+  la <- leaderST dsu a+  sizeLa <- VGM.read parentOrSizeDsu la+  pure (-sizeLa)++{-# INLINEABLE groupsST #-}+groupsST :: Dsu s -> ST s (V.Vector (VU.Vector Int))+groupsST dsu@Dsu {..} = do   groupSize <- VUM.replicate nDsu (0 :: Int)   leaders <- VU.generateM nDsu $ \i -> do     li <- leaderST dsu i
src/AtCoder/Extra/DynSparseSegTree/Persistent.hs view
@@ -113,7 +113,7 @@ {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSparseSegTree (PrimState m) a -> P.Index -> Int -> a -> m P.Index write dst root i x = stToPrim $ do-     Raw.modifyMST dst root (pure . const x) i+  Raw.modifyMST dst root (pure . const x) i  -- | \(O(\log L)\) Modifies the monoid value of the node at \(i\). --@@ -124,7 +124,7 @@ {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSparseSegTree (PrimState m) a -> P.Index -> (a -> a) -> Int -> m P.Index modify dst root f i = stToPrim $ do-     Raw.modifyMST dst root (pure . f) i+  Raw.modifyMST dst root (pure . f) i  -- | \(O(\log L)\) Modifies the monoid value of the node at \(i\). --@@ -135,7 +135,7 @@ {-# INLINE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Raw.DynSparseSegTree (PrimState m) a -> P.Index -> (a -> m a) -> Int -> m P.Index modifyM dst root f i = do-    Raw.modifyMST dst root f i+  Raw.modifyMST dst root f i  -- | \(O(\log L)\) Returns the monoid product in \([l, r)\). --
src/AtCoder/Extra/Graph.hs view
@@ -123,7 +123,7 @@ -- [1,2,4,0,3] -- -- @since 1.1.0.0-{-# INLINABLE topSort #-}+{-# INLINEABLE topSort #-} topSort :: Int -> (Int -> VU.Vector Int) -> VU.Vector Int topSort n gr = runST $ do   inDeg <- VUM.replicate n (0 :: Int)@@ -172,7 +172,7 @@ -- [[3,2],[0,3,1]] -- -- @since 1.1.1.0-{-# INLINABLE blockCut #-}+{-# INLINEABLE blockCut #-} blockCut :: Int -> (Int -> VU.Vector Int) -> Csr () blockCut n gr = runST $ do   low <- VUM.replicate n (0 :: Int)
src/AtCoder/Extra/HashMap.hs view
@@ -77,8 +77,8 @@  import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (void, when)-import Control.Monad.ST (ST) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bit (Bit (..)) import Data.Bits (Bits (xor, (.&.)), (.>>.)) import Data.Vector.Generic qualified as VG@@ -107,9 +107,9 @@     usedHM :: !(VUM.MVector s Bit)   } -{-# INLINE decrementRestCapacity #-}-decrementRestCapacity :: (HasCallStack, PrimMonad m) => VUM.MVector (PrimState m) Int -> String -> m ()-decrementRestCapacity restCap funcName = do+{-# INLINE decrementRestCapacityST #-}+decrementRestCapacityST :: (HasCallStack) => VUM.MVector s Int -> String -> ST s ()+decrementRestCapacityST restCap funcName = do   rest <- VGM.unsafeRead restCap 0   let !_ = ACIA.runtimeAssert (rest > 0) $ "AtCoder.Extra.HashMap." ++ funcName ++ ": out of capacity"   VGM.unsafeWrite restCap 0 (rest - 1)@@ -120,28 +120,14 @@ -- @since 1.1.0.0 {-# INLINE new #-} new :: (PrimMonad m, VU.Unbox a) => Int -> m (HashMap (PrimState m) a)-new n = do-  let !k0 = 1-  let !k = until (>= 2 * n) (* 2) k0-  -- we need extra space-  let !maxCapHM = k `div` 2-  restCapHM <- VUM.replicate 1 maxCapHM-  let !maskHM = k - 1-  keyHM <- VUM.unsafeNew k-  valHM <- VUM.unsafeNew k-  usedHM <- VUM.replicate k $ Bit False-  pure HashMap {..}+new n = stToPrim $ newST n  -- | \(O(n)\) Creates a `HashMap` of capacity \(n\) with initial entries. -- -- @since 1.1.0.0 {-# INLINE build #-} build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (HashMap (PrimState m) a)-build n xs = do-  hm <- new n-  VU.forM_ xs $ \(!i, !x) -> do-    insert hm i x-  pure hm+build n xs = stToPrim $ buildST n xs  -- | \(O(1)\) Returns the maximum number of elements the hash map can store. --@@ -159,36 +145,6 @@   !rest <- VUM.unsafeRead restCapHM 0   pure $ maxCapHM - rest --- | \(O(1)\) (Internal) Hash value calculation.------ @since 1.1.0.0-{-# INLINE hash #-}-hash :: HashMap a s -> Int -> Int-hash hm x = fromIntegral $ (x3 `xor` (x3 .>>. 31)) .&. fromIntegral (maskHM hm)-  where-    fixedRandom, x1, x2, x3 :: Word64-    fixedRandom = 321896566547-    x1 = fromIntegral x + fixedRandom-    x2 = (x1 `xor` (x1 .>>. 30)) * 0xbf58476d1ce4e5b9-    x3 = (x2 `xor` (x2 .>>. 27)) * 0x94d049bb133111eb---- | \(O(1)\) (Internal) Hashed slot search.------ ==== Constraint--- - The rest capacity must be non-zero. Otherwise it loops forever.-{-# INLINE indexST #-}-indexST :: (HasCallStack) => HashMap s a -> Int -> ST s Int-indexST hm@HashMap {..} k = do-  inner (hash hm k)-  where-    inner !h = do-      Bit b <- VGM.read usedHM h-      -- already there?-      k' <- VGM.read keyHM h-      if b && k' /= k-        then inner $ (h + 1) .&. maskHM-        else pure h- -- | \(O(1)\) Return the value to which the specified key is mapped, or `Nothing` if this map -- contains no mapping for the key. --@@ -198,12 +154,7 @@ -- @since 1.1.0.0 {-# INLINE lookup #-} lookup :: (HasCallStack, VU.Unbox a, PrimMonad m) => HashMap (PrimState m) a -> Int -> m (Maybe a)-lookup hm@HashMap {..} k = do-  i <- stToPrim $ indexST hm k-  Bit b <- VGM.read usedHM i-  if b-    then Just <$> VGM.read valHM i-    else pure Nothing+lookup hm k = stToPrim $ lookupST hm k  -- | \(O(1)\) Checks whether the hash map contains the element. --@@ -213,12 +164,7 @@ -- @since 1.1.0.0 {-# INLINE member #-} member :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool-member hm@HashMap {..} k = do-  i <- stToPrim $ indexST hm k-  Bit b <- VGM.read usedHM i-  -- TODO: is this key check necessary-  k' <- VGM.read keyHM i-  pure $ b && k' == k+member hm k = stToPrim $ memberST hm k  -- | \(O(1)\) Checks whether the hash map does not contain the element. --@@ -228,7 +174,7 @@ -- @since 1.1.0.0 {-# INLINE notMember #-} notMember :: (HasCallStack, PrimMonad m) => HashMap (PrimState m) a -> Int -> m Bool-notMember hm k = not <$> member hm k+notMember hm k = stToPrim $ not <$> memberST hm k  -- | \(O(1)\) Inserts a \((k, v)\) pair. --@@ -238,7 +184,7 @@ -- @since 1.1.0.0 {-# INLINE insert #-} insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m ()-insert hm k v = void $ exchange hm k v+insert hm k v = void . stToPrim $ exchangeST hm k v  -- | \(O(1)\) Inserts a \((k, v)\) pair. If the key exists, the function will insert the pair -- \((k, f(v_{\mathrm{new}}, v_{\mathrm{old}}))\).@@ -249,18 +195,7 @@ -- @since 1.1.0.0 {-# INLINE insertWith #-} insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()-insertWith hm@HashMap {..} f k v = do-  i <- stToPrim $ indexST hm k-  Bit b <- VGM.exchange usedHM i $ Bit True-  if b-    then do-      -- modify the existing entry-      VGM.modify valHM (f v) i-    else do-      -- insert the new \((k, v)\) pair-      decrementRestCapacity restCapHM "insertWith"-      VGM.write keyHM i k-      VGM.write valHM i v+insertWith hm f k v = stToPrim $ insertWithST hm f k v  -- | \(O(1)\) Inserts a \((k, v)\) pair and returns the old value, or `Nothing` if no such entry -- exists.@@ -271,30 +206,14 @@ -- @since 1.1.0.0 {-# INLINE exchange #-} exchange :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> Int -> a -> m (Maybe a)-exchange hm@HashMap {..} k v = do-  i <- stToPrim $ indexST hm k-  Bit b <- VGM.exchange usedHM i $ Bit True-  if b-    then do-      -- overwrite the existing entry-      Just <$> VGM.exchange valHM i v-    else do-      -- insert the new (key, value) pair-      decrementRestCapacity restCapHM "exchange"-      VGM.write keyHM i k-      VGM.write valHM i v-      pure Nothing+exchange hm k v = stToPrim $ exchangeST hm k v  -- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists. -- -- @since 1.1.0.0 {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> a) -> Int -> m ()-modify hm@HashMap {..} f k = do-  i <- stToPrim $ indexST hm k-  Bit b <- VGM.read usedHM i-  when b $ do-    VGM.modify valHM f i+modify hm f k = stToPrim $ modifyST hm f k  -- | \(O(1)\) Modifies the element at the given key. Does nothing if no such entry exists. --@@ -303,7 +222,7 @@ modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> (a -> m a) -> Int -> m () modifyM hm@HashMap {..} f k = do   i <- stToPrim $ indexST hm k-  Bit b <- VGM.read usedHM i+  Bit b <- stToPrim $ VGM.read usedHM i   when b $ do     VGM.modifyM valHM f i @@ -312,9 +231,7 @@ -- @since 1.1.0.0 {-# INLINE clear #-} clear :: (PrimMonad m) => HashMap (PrimState m) a -> m ()-clear HashMap {..} = do-  VGM.set usedHM $ Bit False-  VUM.unsafeWrite restCapHM 0 maxCapHM+clear hm = stToPrim $ clearST hm  -- | \(O(n)\) Enumerates the keys in the hash map. --@@ -342,27 +259,160 @@ -- @since 1.1.0.0 {-# INLINE unsafeKeys #-} unsafeKeys :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector Int)-unsafeKeys HashMap {..} = do-  used <- VU.unsafeFreeze usedHM-  keys_ <- VU.unsafeFreeze keyHM-  pure $ VU.ifilter (const . unBit . (used VG.!)) keys_+unsafeKeys hm = stToPrim $ unsafeKeysST hm  -- | \(O(n)\) Enumerates the elements (values) in the hash map. -- -- @since 1.1.0.0 {-# INLINE unsafeElems #-} unsafeElems :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector a)-unsafeElems HashMap {..} = do-  used <- VU.unsafeFreeze usedHM-  vals <- VU.unsafeFreeze valHM-  pure $ VU.ifilter (const . unBit . (used VG.!)) vals+unsafeElems hm = stToPrim $ unsafeElemsST hm  -- | \(O(n)\) Enumerates the key-value pairs in the hash map. -- -- @since 1.1.0.0 {-# INLINE unsafeAssocs #-} unsafeAssocs :: (PrimMonad m, VU.Unbox a) => HashMap (PrimState m) a -> m (VU.Vector (Int, a))-unsafeAssocs HashMap {..} = do+unsafeAssocs hm = stToPrim $ unsafeAssocsST hm++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (VU.Unbox a) => Int -> ST s (HashMap s a)+newST n = do+  let !k0 = 1+  let !k = until (>= 2 * n) (* 2) k0+  -- we need extra space+  let !maxCapHM = k `div` 2+  restCapHM <- VUM.replicate 1 maxCapHM+  let !maskHM = k - 1+  keyHM <- VUM.unsafeNew k+  valHM <- VUM.unsafeNew k+  usedHM <- VUM.replicate k $ Bit False+  pure HashMap {..}++{-# INLINEABLE buildST #-}+buildST :: (VU.Unbox a) => Int -> VU.Vector (Int, a) -> ST s (HashMap s a)+buildST n xs = do+  hm <- newST n+  VU.forM_ xs $ \(!i, !x) -> do+    void $ exchangeST hm i x+  pure hm++-- TODO: no need of INLINE?++-- | \(O(1)\) (Internal) Hash value calculation.+{-# INLINEABLE hash #-}+hash :: HashMap a s -> Int -> Int+hash hm x = fromIntegral $ (x3 `xor` (x3 .>>. 31)) .&. fromIntegral (maskHM hm)+  where+    fixedRandom, x1, x2, x3 :: Word64+    fixedRandom = 321896566547+    x1 = fromIntegral x + fixedRandom+    x2 = (x1 `xor` (x1 .>>. 30)) * 0xbf58476d1ce4e5b9+    x3 = (x2 `xor` (x2 .>>. 27)) * 0x94d049bb133111eb++-- TODO: INLINE indexST?++-- | \(O(1)\) (Internal) Hashed slot search.+--+-- ==== Constraint+-- - The rest capacity must be non-zero. Otherwise it loops forever.+{-# INLINEABLE indexST #-}+indexST :: (HasCallStack) => HashMap s a -> Int -> ST s Int+indexST hm@HashMap {..} k = do+  inner (hash hm k)+  where+    inner !h = do+      Bit b <- VGM.read usedHM h+      -- already there?+      k' <- VGM.read keyHM h+      if b && k' /= k+        then inner $ (h + 1) .&. maskHM+        else pure h++{-# INLINEABLE lookupST #-}+lookupST :: (HasCallStack, VU.Unbox a) => HashMap s a -> Int -> ST s (Maybe a)+lookupST hm@HashMap {..} k = do+  i <- indexST hm k+  Bit b <- VGM.read usedHM i+  if b+    then Just <$> VGM.read valHM i+    else pure Nothing++{-# INLINEABLE memberST #-}+memberST :: (HasCallStack) => HashMap s a -> Int -> ST s Bool+memberST hm@HashMap {..} k = do+  i <- indexST hm k+  Bit b <- VGM.read usedHM i+  -- TODO: is this key check necessary+  k' <- VGM.read keyHM i+  pure $ b && k' == k++{-# INLINEABLE insertWithST #-}+insertWithST :: (HasCallStack, VU.Unbox a) => HashMap s a -> (a -> a -> a) -> Int -> a -> ST s ()+insertWithST hm@HashMap {..} f k v = do+  i <- indexST hm k+  Bit b <- VGM.exchange usedHM i $ Bit True+  if b+    then do+      -- modify the existing entry+      VGM.modify valHM (f v) i+    else do+      -- insert the new \((k, v)\) pair+      decrementRestCapacityST restCapHM "insertWith"+      VGM.write keyHM i k+      VGM.write valHM i v++{-# INLINEABLE exchangeST #-}+exchangeST :: (HasCallStack, VU.Unbox a) => HashMap s a -> Int -> a -> ST s (Maybe a)+exchangeST hm@HashMap {..} k v = do+  i <- indexST hm k+  Bit b <- VGM.exchange usedHM i $ Bit True+  if b+    then do+      -- overwrite the existing entry+      Just <$> VGM.exchange valHM i v+    else do+      -- insert the new (key, value) pair+      decrementRestCapacityST restCapHM "exchange"+      VGM.write keyHM i k+      VGM.write valHM i v+      pure Nothing++{-# INLINEABLE modifyST #-}+modifyST :: (HasCallStack, VU.Unbox a) => HashMap s a -> (a -> a) -> Int -> ST s ()+modifyST hm@HashMap {..} f k = do+  i <- indexST hm k+  Bit b <- VGM.read usedHM i+  when b $ do+    VGM.modify valHM f i++{-# INLINEABLE clearST #-}+clearST :: HashMap s a -> ST s ()+clearST HashMap {..} = do+  VGM.set usedHM $ Bit False+  VUM.unsafeWrite restCapHM 0 maxCapHM++{-# INLINEABLE unsafeKeysST #-}+unsafeKeysST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector Int)+unsafeKeysST HashMap {..} = do+  used <- VU.unsafeFreeze usedHM+  keys_ <- VU.unsafeFreeze keyHM+  pure $ VU.ifilter (const . unBit . (used VG.!)) keys_++{-# INLINEABLE unsafeElemsST #-}+unsafeElemsST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector a)+unsafeElemsST HashMap {..} = do+  used <- VU.unsafeFreeze usedHM+  vals <- VU.unsafeFreeze valHM+  pure $ VU.ifilter (const . unBit . (used VG.!)) vals++{-# INLINEABLE unsafeAssocsST #-}+unsafeAssocsST :: (VU.Unbox a) => HashMap s a -> ST s (VU.Vector (Int, a))+unsafeAssocsST HashMap {..} = do   used <- VU.unsafeFreeze usedHM   keys_ <- VU.unsafeFreeze keyHM   vals <- VU.unsafeFreeze valHM
src/AtCoder/Extra/IntMap.hs view
@@ -82,7 +82,8 @@  import AtCoder.Extra.IntSet qualified as IS import Control.Monad (when)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Maybe (fromJust) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU@@ -101,158 +102,130 @@ -- | \(O(n)\) Creates an `IntMap` for an interval \([0, n)\). -- -- @since 1.1.0.0-{-# INLINE new #-}+{-# INLINEABLE new #-} new :: (PrimMonad m, VU.Unbox a) => Int -> m (IntMap (PrimState m) a)-new cap = do-  setIM <- IS.new cap-  valIM <- VUM.unsafeNew cap-  pure IntMap {..}+new cap = stToPrim $ newST cap  -- | \(O(n + m \log n)\) Creates an `IntMap` for an interval \([0, n)\) with initial values. -- -- @since 1.1.0.0-{-# INLINE build #-}+{-# INLINEABLE build #-} build :: (PrimMonad m, VU.Unbox a) => Int -> VU.Vector (Int, a) -> m (IntMap (PrimState m) a)-build cap xs = do-  im <- new cap-  VU.forM_ xs $ \(!k, !v) -> do-    insert im k v-  pure im+build cap xs = stToPrim $ buildST cap xs  -- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the map. -- -- @since 1.1.0.0-{-# INLINE capacity #-}+{-# INLINEABLE capacity #-} capacity :: IntMap s a -> Int capacity = IS.capacity . setIM  -- | \(O(1)\) Returns the number of elements in the map. -- -- @since 1.1.0.0-{-# INLINE size #-}+{-# INLINEABLE size #-} size :: (PrimMonad m) => IntMap (PrimState m) a -> m Int size = IS.size . setIM  -- | \(O(1)\) Returns whether the map is empty. -- -- @since 1.1.0.0-{-# INLINE null #-}+{-# INLINEABLE null #-} null :: (PrimMonad m) => IntMap (PrimState m) a -> m Bool null = IS.null . setIM  -- | \(O(\log n)\) Looks up the value for a key. -- -- @since 1.1.0.0-{-# INLINE lookup #-}+{-# INLINEABLE lookup #-} lookup :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe a)-lookup im@IntMap {..} k = do-  member im k >>= \case-    True -> Just <$> VGM.read valIM k-    False -> pure Nothing+lookup im k = stToPrim $ lookupST im k  -- | \(O(\log n)\) Tests whether a key \(k\) is in the map. -- -- @since 1.1.0.0-{-# INLINE member #-}+{-# INLINEABLE member #-} member :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool member = IS.member . setIM  -- | \(O(\log n)\) Tests whether a key \(k\) is not in the map. -- -- @since 1.1.0.0-{-# INLINE notMember #-}+{-# INLINEABLE notMember #-} notMember :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool notMember = IS.notMember . setIM  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest key \(k\) such that \(k \ge k_0\). -- -- @since 1.1.0.0-{-# INLINE lookupGE #-}+{-# INLINEABLE lookupGE #-} lookupGE :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))-lookupGE IntMap {..} k = do-  IS.lookupGE setIM k >>= \case-    Just i -> Just . (i,) <$> VGM.read valIM i-    Nothing -> pure Nothing+lookupGE im k = stToPrim $ lookupGEST im k  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the smallest \(k\) such that \(k \gt k_0\). -- -- @since 1.1.0.0-{-# INLINE lookupGT #-}+{-# INLINEABLE lookupGT #-} lookupGT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))-lookupGT is k = lookupGE is (k + 1)+lookupGT is k = stToPrim $ lookupGEST is (k + 1)  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \le k_0\). -- -- @since 1.1.0.0-{-# INLINE lookupLE #-}+{-# INLINEABLE lookupLE #-} lookupLE :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))-lookupLE IntMap {..} k = do-  IS.lookupLE setIM k >>= \case-    Just i -> Just . (i,) <$> VGM.read valIM i-    Nothing -> pure Nothing+lookupLE im k = stToPrim $ lookupLEST im k  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the largest key \(k\) such that \(k \lt k_0\). -- -- @since 1.1.0.0-{-# INLINE lookupLT #-}+{-# INLINEABLE lookupLT #-} lookupLT :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> m (Maybe (Int, a))-lookupLT is k = lookupLE is (k - 1)+lookupLT im k = stToPrim $ lookupLEST im (k - 1)  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the minimum key \(k\). -- -- @since 1.1.0.0-{-# INLINE lookupMin #-}+{-# INLINEABLE lookupMin #-} lookupMin :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))-lookupMin is = lookupGE is 0+lookupMin im = stToPrim $ lookupMinST im  -- | \(O(\log n)\) Looks up the \((k, v)\) pair with the maximum key \(k\). -- -- @since 1.1.0.0-{-# INLINE lookupMax #-}+{-# INLINEABLE lookupMax #-} lookupMax :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))-lookupMax im = lookupLE im (IS.capacity (setIM im) - 1)+lookupMax im = stToPrim $ lookupMaxST im  -- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already -- exists, it is overwritten. -- -- @since 1.1.0.0-{-# INLINE insert #-}+{-# INLINEABLE insert #-} insert :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> Int -> a -> m ()-insert IntMap {..} k v = do-  IS.insert setIM k-  VGM.write valIM k v+insert im k v = stToPrim $ insertST im k v  -- | \(O(\log n)\) Inserts a \((k, v)\) pair into the map. If an entry with the same key already -- exists, it overwritten with \(f(v_{\mathrm{new}}, v_{\mathrm{old}})\). -- -- @since 1.1.0.0-{-# INLINE insertWith #-}+{-# INLINEABLE insertWith #-} insertWith :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a -> a) -> Int -> a -> m ()-insertWith IntMap {..} f k v = do-  b <- IS.member setIM k-  if b-    then do-      VGM.modify valIM (f v) k-    else do-      IS.insert setIM k-      VGM.write valIM k v+insertWith im f k v = stToPrim $ insertWithST im f k v  -- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already -- does not exist, nothing is performed. -- -- @since 1.1.0.0-{-# INLINE modify #-}+{-# INLINEABLE modify #-} modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> a) -> Int -> m ()-modify IntMap {..} f k = do-  b <- IS.member setIM k-  when b $ do-    VGM.modify valIM f k+modify im f k = stToPrim $ modifyST im f k  -- | \(O(\log n)\) Modifies the value associated with a key. If an entry with the same key already -- does not exist, nothing is performed. -- -- @since 1.1.0.0-{-# INLINE modifyM #-}+{-# INLINEABLE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> (a -> m a) -> Int -> m () modifyM IntMap {..} f k = do   b <- IS.member setIM k@@ -263,77 +236,180 @@ -- such key exists. Returns whether the key existed. -- -- @since 1.1.0.0-{-# INLINE delete #-}+{-# INLINEABLE delete #-} delete :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m Bool-delete im = IS.delete (setIM im)+delete im = stToPrim . deleteST im  -- | \(O(\log n)\) Deletes the \((k, v)\) pair with the key \(k\) from the map. Does nothing if no -- such key exists. -- -- @since 1.1.0.0-{-# INLINE delete_ #-}+{-# INLINEABLE delete_ #-} delete_ :: (PrimMonad m) => IntMap (PrimState m) a -> Int -> m ()-delete_ im = IS.delete_ (setIM im)+delete_ im = stToPrim . deleteST_ im  -- | \(O(\log n)\) Deletes the \((k, v)\) pair with the minimum key \(k\) in the map. -- -- @since 1.1.0.0-{-# INLINE deleteMin #-}+{-# INLINEABLE deleteMin #-} deleteMin :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))-deleteMin is = do-  lookupMin is-    >>= mapM-      ( \(!key, !val) -> do-          delete_ is key-          pure (key, val)-      )+deleteMin is = stToPrim $ deleteMinST is  -- | \(O(\log n)\) Deletes the \((k, v)\) pair with maximum key \(k\) in the map. -- -- @since 1.1.0.0-{-# INLINE deleteMax #-}+{-# INLINEABLE deleteMax #-} deleteMax :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (Maybe (Int, a))-deleteMax is = do-  lookupMax is-    >>= mapM-      ( \(!k, !v) -> do-          delete_ is k-          pure (k, v)-      )+deleteMax is = stToPrim $ deleteMaxST is  -- | \(O(n \log n)\) Enumerates the keys in the map. -- -- @since 1.1.0.0-{-# INLINE keys #-}+{-# INLINEABLE keys #-} keys :: (PrimMonad m) => IntMap (PrimState m) a -> m (VU.Vector Int)-keys = IS.keys . setIM+keys = stToPrim . keysST  -- | \(O(n \log n)\) Enumerates the elements (values) in the map. -- -- @since 1.1.0.0-{-# INLINE elems #-}+{-# INLINEABLE elems #-} elems :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector a)-elems im@IntMap {..} = do+elems = stToPrim . elemsST++-- | \(O(n \log n)\) Enumerates the key-value pairs in the map.+--+-- @since 1.1.0.0+{-# INLINEABLE assocs #-}+assocs :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector (Int, a))+assocs = stToPrim . assocsST++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (VU.Unbox a) => Int -> ST s (IntMap s a)+newST cap = do+  setIM <- IS.new cap+  valIM <- VUM.unsafeNew cap+  pure IntMap {..}++{-# INLINEABLE buildST #-}+buildST :: (VU.Unbox a) => Int -> VU.Vector (Int, a) -> ST s (IntMap s a)+buildST cap xs = do+  im <- new cap+  VU.forM_ xs $ \(!k, !v) -> do+    insertST im k v+  pure im++-- | \(O(\log n)\) Looks up the value for a key.+--+-- @since 1.1.0.0+{-# INLINEABLE lookupST #-}+lookupST :: (VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe a)+lookupST im@IntMap {..} k = do+  member im k >>= \case+    True -> Just <$> VGM.read valIM k+    False -> pure Nothing++{-# INLINEABLE lookupGEST #-}+lookupGEST :: (VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe (Int, a))+lookupGEST IntMap {..} k = do+  IS.lookupGE setIM k >>= \case+    Just i -> Just . (i,) <$> VGM.read valIM i+    Nothing -> pure Nothing++{-# INLINEABLE lookupLEST #-}+lookupLEST :: (HasCallStack, VU.Unbox a) => IntMap s a -> Int -> ST s (Maybe (Int, a))+lookupLEST IntMap {..} k = do+  IS.lookupLE setIM k >>= \case+    Just i -> Just . (i,) <$> VGM.read valIM i+    Nothing -> pure Nothing++{-# INLINEABLE lookupMinST #-}+lookupMinST :: (VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))+lookupMinST is = lookupGEST is 0++{-# INLINEABLE lookupMaxST #-}+lookupMaxST :: (VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))+lookupMaxST im = lookupLEST im (IS.capacity (setIM im) - 1)++{-# INLINEABLE insertST #-}+insertST :: (HasCallStack, VU.Unbox a) => IntMap s a -> Int -> a -> ST s ()+insertST IntMap {..} k v = do+  IS.insert setIM k+  VGM.write valIM k v++{-# INLINEABLE insertWithST #-}+insertWithST :: (HasCallStack, VU.Unbox a) => IntMap s a -> (a -> a -> a) -> Int -> a -> ST s ()+insertWithST IntMap {..} f k v = do+  b <- IS.member setIM k+  if b+    then do+      VGM.modify valIM (f v) k+    else do+      IS.insert setIM k+      VGM.write valIM k v++{-# INLINEABLE modifyST #-}+modifyST :: (HasCallStack, VU.Unbox a) => IntMap s a -> (a -> a) -> Int -> ST s ()+modifyST IntMap {..} f k = do+  b <- IS.member setIM k+  when b $ do+    VGM.modify valIM f k++{-# INLINEABLE deleteST #-}+deleteST :: IntMap s a -> Int -> ST s Bool+deleteST im = IS.delete (setIM im)++{-# INLINEABLE deleteST_ #-}+deleteST_ :: IntMap s a -> Int -> ST s ()+deleteST_ im = IS.delete_ (setIM im)++{-# INLINEABLE deleteMinST #-}+deleteMinST :: (HasCallStack, VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))+deleteMinST is = do+  lookupMinST is+    >>= mapM+      ( \(!key, !val) -> do+          deleteST_ is key+          pure (key, val)+      )++{-# INLINEABLE deleteMaxST #-}+deleteMaxST :: (HasCallStack, VU.Unbox a) => IntMap s a -> ST s (Maybe (Int, a))+deleteMaxST is = do+  lookupMaxST is+    >>= mapM+      ( \(!k, !v) -> do+          deleteST_ is k+          pure (k, v)+      )++{-# INLINEABLE keysST #-}+keysST :: IntMap s a -> ST s (VU.Vector Int)+keysST = IS.keys . setIM++{-# INLINEABLE elemsST #-}+elemsST :: (VU.Unbox a) => IntMap s a -> ST s (VU.Vector a)+elemsST im@IntMap {..} = do   n <- IS.size setIM   VU.unfoldrExactNM     n     ( \i -> do-        (!i', !x') <- fromJust <$> lookupGT im i+        (!i', !x') <- fromJust <$> lookupGEST im (i + 1)         pure (x', i')     )     (-1) --- | \(O(n \log n)\) Enumerates the key-value pairs in the map.------ @since 1.1.0.0-{-# INLINE assocs #-}-assocs :: (PrimMonad m, VU.Unbox a) => IntMap (PrimState m) a -> m (VU.Vector (Int, a))-assocs im@IntMap {..} = do+{-# INLINEABLE assocsST #-}+assocsST :: (VU.Unbox a) => IntMap s a -> ST s (VU.Vector (Int, a))+assocsST im@IntMap {..} = do   n <- IS.size setIM   VU.unfoldrExactNM     n     ( \i -> do-        (!i', !x') <- fromJust <$> lookupGT im i+        (!i', !x') <- fromJust <$> lookupGEST im (i + 1)         pure ((i', x'), i')     )     (-1)
src/AtCoder/Extra/IntSet.hs view
@@ -76,7 +76,8 @@  import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (unless, void)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bifunctor (bimap) import Data.Bits   ( Bits (clearBit, setBit, testBit),@@ -143,33 +144,14 @@ -- @since 1.1.0.0 {-# INLINE new #-} new :: (PrimMonad m) => Int -> m (IntSet (PrimState m))-new 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)+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 = do-  set <- new n-  VU.forM_ vs (insert set)-  pure set+build n vs = stToPrim $ buildST n vs  -- | \(O(1)\) Returns the capacity \(n\), where the interval \([0, n)\) is covered by the set. --@@ -197,25 +179,143 @@ -- @since 1.1.0.0 {-# INLINE member #-} member :: (PrimMonad m) => IntSet (PrimState m) -> Int -> m Bool-member IntSet {..} k-  | ACIA.testIndex k capacityIS = do-      let (!q, !r) = k `divMod` wordSize-      (`testBit` r) <$> VGM.unsafeRead (VG.unsafeHead vecIS) q-  | otherwise = pure False+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 = not <$> member dis k+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 IntSet {..} i0+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@@ -239,19 +339,13 @@       where         (!q, !r) = i `divMod` wordSize --- | \(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 = lookupGE is (k + 1)+{-# INLINEABLE lookupGTST #-}+lookupGTST :: IntSet s -> Int -> ST s (Maybe Int)+lookupGTST is k = lookupGEST is (k + 1) --- | \(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 IntSet {..} i0+{-# 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@@ -274,35 +368,22 @@       where         (!q, !r) = i `divMod` wordSize --- | \(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 = lookupLE is (k - 1)+{-# INLINEABLE lookupLTST #-}+lookupLTST :: IntSet s -> Int -> ST s (Maybe Int)+lookupLTST is k = lookupLEST is (k - 1) --- | \(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 = lookupGE is 0+{-# INLINEABLE lookupMinST #-}+lookupMinST :: IntSet s -> ST s (Maybe Int)+lookupMinST is = lookupGEST is 0 --- | \(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 = lookupLE is (capacityIS is - 1)+{-# INLINEABLE lookupMaxST #-}+lookupMaxST :: IntSet s -> ST s (Maybe Int)+lookupMaxST is = lookupLEST is (capacityIS is - 1) --- | \(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@IntSet {..} k = do-  b <- member is k+{-# 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'_@@ -316,14 +397,10 @@   where     !_ = ACIA.checkIndex "AtCoder.Extra.IntSet.insert" k capacityIS --- | \(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@IntSet {..} k = do-  b_ <- member is k+{-# 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@@ -342,50 +419,38 @@       pure True     else pure False --- | \(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 = void $ delete is k+{-# INLINEABLE deleteST_ #-}+deleteST_ :: IntSet s -> Int -> ST s ()+deleteST_ is k = void $ 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 = do-  lookupMin is+{-# INLINEABLE deleteMinST #-}+deleteMinST :: IntSet s -> ST s (Maybe Int)+deleteMinST is = do+  lookupMinST is     >>= mapM       ( \key -> do-          delete_ is key+          deleteST_ is key           pure key       ) --- | \(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 = do-  lookupMax is+{-# INLINEABLE deleteMaxST #-}+deleteMaxST :: IntSet s -> ST s (Maybe Int)+deleteMaxST is = do+  lookupMaxST is     >>= mapM       ( \key -> do-          delete_ is key+          deleteST_ is key           pure key       ) --- | \(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@IntSet {sizeIS} = do+{-# 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 <$> lookupGT is i+        i' <- fromJust <$> lookupGTST is i         pure (i', i')     )     (-1)
src/AtCoder/Extra/IntervalMap.hs view
@@ -92,6 +92,7 @@ import AtCoder.Extra.IntMap qualified as IM import Control.Monad (foldM_) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Vector.Unboxed qualified as VU import GHC.Stack (HasCallStack) import Prelude hiding (lookup, read)@@ -148,7 +149,7 @@   where     step dim !l !xs' = do       let !l' = l + VU.length xs'-      IM.insert dim l (l', VU.head xs')+      stToPrim $ IM.insert dim l (l', VU.head xs')       onAdd l l' (VU.head xs')       pure l' @@ -171,7 +172,7 @@ -- @since 1.1.0.0 {-# INLINE contains #-} contains :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> m Bool-contains itm i = containsInterval itm i (i + 1)+contains itm i = stToPrim $ containsIntervalST itm i (i + 1)  -- | \(O(\log n)\) Returns whether an interval \([l, r)\) is fully contained within any of the -- intervals.@@ -179,27 +180,14 @@ -- @since 1.1.0.0 {-# INLINE containsInterval #-} containsInterval :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m Bool-containsInterval (IntervalMap dim) l r-  | l >= r = pure False-  | otherwise = do-      res <- IM.lookupLE dim l-      pure $ case res of-        Just (!_, (!r', !_)) -> r <= r'-        _ -> False+containsInterval itm l r = stToPrim $ containsIntervalST itm l r  -- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\). -- -- @since 1.1.0.0 {-# INLINE lookup #-} lookup :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe (Int, Int, a))-lookup (IntervalMap im) l r-  | l >= r = pure Nothing-  | otherwise = do-      res <- IM.lookupLE im l-      pure $ case res of-        Just (!l', (!r', !a))-          | r <= r' -> Just (l', r', a)-        _ -> Nothing+lookup itm l r = stToPrim $ lookupST itm l r  -- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value. -- Throws an error if no such interval exists.@@ -207,11 +195,7 @@ -- @since 1.1.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m a-read itm l r = do-  res <- readMaybe itm l r-  pure $ case res of-    Just !a -> a-    Nothing -> error $ "[read] not a member: " ++ show (l, r)+read itm l r = stToPrim $ readST itm l r  -- | \(O(\log n)\) Looks up an interval that fully contains \([l, r)\) and reads out the value. -- Returns `Nothing` if no such interval exists.@@ -219,14 +203,7 @@ -- @since 1.1.0.0 {-# INLINE readMaybe #-} readMaybe :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> m (Maybe a)-readMaybe (IntervalMap dim) l r-  | l >= r = pure Nothing-  | otherwise = do-      res <- IM.lookupLE dim l-      pure $ case res of-        Just (!_, (!r', !a))-          | r <= r' -> Just a-        _ -> Nothing+readMaybe itm l r = stToPrim $ readMaybeST itm l r  -- | Amortized \(O(\log n)\) Inserts an interval \([l, r)\) with associated value \(v\) into the -- map. Overwrites any overlapping intervals.@@ -244,7 +221,7 @@ -- hooks. -- -- @since 1.1.0.0-{-# INLINABLE insertM #-}+{-# INLINEABLE insertM #-} insertM ::   (PrimMonad m, Eq a, VU.Unbox a) =>   -- | The map@@ -376,7 +353,7 @@ -- changes via @onAdd@ and @onDel@ hooks. -- -- @since 1.1.0.0-{-# INLINABLE deleteM #-}+{-# INLINEABLE deleteM #-} deleteM ::   (PrimMonad m, VU.Unbox a) =>   -- | The map@@ -453,11 +430,7 @@ -- @since 1.1.0.0 {-# INLINE overwrite #-} overwrite :: (PrimMonad m, Eq a, VU.Unbox a) => IntervalMap (PrimState m) a -> Int -> Int -> a -> m ()-overwrite itm l r x = do-  res <- lookup itm l r-  case res of-    Just (!l', !r', !_) -> insert itm l' r' x-    Nothing -> pure ()+overwrite itm l r x = stToPrim $ overwriteST itm l r x  -- | \(O(\log n)\). Shorthand for overwriting the value of an interval that contains \([l, r)\). -- Tracks interval state changes via @onAdd@ and @onDel@ hooks.@@ -480,7 +453,7 @@   (Int -> Int -> a -> m ()) ->   m () overwriteM itm l r x onAdd onDel = do-  res <- lookup itm l r+  res <- stToPrim $ lookupST itm l r   case res of     Just (!l', !r', !_) -> insertM itm l' r' x onAdd onDel     Nothing -> pure ()@@ -492,3 +465,55 @@ {-# INLINE freeze #-} freeze :: (PrimMonad m, VU.Unbox a) => IntervalMap (PrimState m) a -> m (VU.Vector (Int, (Int, a))) freeze = IM.assocs . unITM++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE containsIntervalST #-}+containsIntervalST :: (VU.Unbox a) => IntervalMap s a -> Int -> Int -> ST s Bool+containsIntervalST (IntervalMap dim) l r+  | l >= r = pure False+  | otherwise = do+      res <- IM.lookupLE dim l+      pure $ case res of+        Just (!_, (!r', !_)) -> r <= r'+        _ -> False++{-# INLINEABLE lookupST #-}+lookupST :: (VU.Unbox a) => IntervalMap s a -> Int -> Int -> ST s (Maybe (Int, Int, a))+lookupST (IntervalMap im) l r+  | l >= r = pure Nothing+  | otherwise = do+      res <- IM.lookupLE im l+      pure $ case res of+        Just (!l', (!r', !a))+          | r <= r' -> Just (l', r', a)+        _ -> Nothing++{-# INLINEABLE readST #-}+readST :: (HasCallStack, VU.Unbox a) => IntervalMap s a -> Int -> Int -> ST s a+readST itm l r = do+  res <- readMaybeST itm l r+  pure $ case res of+    Just !a -> a+    Nothing -> error $ "AtCoder.Extra.IntervalMap.readST: not a member: " ++ show (l, r)++{-# INLINEABLE readMaybeST #-}+readMaybeST :: (VU.Unbox a) => IntervalMap s a -> Int -> Int -> ST s (Maybe a)+readMaybeST (IntervalMap dim) l r+  | l >= r = pure Nothing+  | otherwise = do+      res <- IM.lookupLE dim l+      pure $ case res of+        Just (!_, (!r', !a))+          | r <= r' -> Just a+        _ -> Nothing++{-# INLINEABLE overwriteST #-}+overwriteST :: (Eq a, VU.Unbox a) => IntervalMap s a -> Int -> Int -> a -> ST s ()+overwriteST itm l r x = do+  res <- lookupST itm l r+  case res of+    Just (!l', !r', !_) -> insert itm l' r' x+    Nothing -> pure ()
src/AtCoder/Extra/KdTree.hs view
@@ -155,7 +155,7 @@ -- | \(O(n \log n)\) Collects points in \([x_l, x_r) \times [y_l, y_r)\). -- -- @since 1.2.2.0-{-# INLINE findPointsIn #-}+{-# INLINEABLE findPointsIn #-} findPointsIn ::   (HasCallStack) =>   -- | `KdTree`@@ -200,7 +200,7 @@ -- point, or `Nothing` if the `KdTree` has no point. -- -- @since 1.2.2.0-{-# INLINE findNearestPoint #-}+{-# INLINEABLE findNearestPoint #-} findNearestPoint ::   (HasCallStack) =>   -- | `KdTree`
src/AtCoder/Extra/MultiSet.hs view
@@ -46,7 +46,7 @@ -- -- >>> MS.inc ms 11 -- >>> MS.sub ms 11 2--- *** Exception: AtCoder.Extra.Multiset.sub: the count of `11` is becoming a negative value: `-1`+-- *** Exception: AtCoder.Extra.Multiset.subST: the count of `11` is becoming a negative value: `-1` -- ... -- -- Decrementing a non-existing key does nothing and does not throw an exception:@@ -100,7 +100,8 @@  import AtCoder.Extra.HashMap qualified as HM import Control.Monad (when)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Functor ((<&>)) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU@@ -121,10 +122,7 @@ -- @since 1.1.0.0 {-# INLINE new #-} new :: (PrimMonad m) => Int -> m (MultiSet (PrimState m))-new n = do-  mapMS <- HM.new n-  cntMS <- VUM.replicate 1 0-  pure $ MultiSet {..}+new n = stToPrim $ newST n  -- | \(O(1)\) Returns the maximum number of distinct keys that can be inserted into the internal -- hash map.@@ -147,27 +145,21 @@ -- @since 1.1.0.0 {-# INLINE lookup #-} lookup :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m (Maybe Int)-lookup MultiSet {..} k = do-  HM.lookup mapMS k <&> \case-    Just i | i > 0 -> Just i-    _ -> Nothing+lookup ms k = stToPrim $ lookupST ms k  -- | \(O(1)\) Tests whether \(k\) is in the set. -- -- @since 1.1.0.0 {-# INLINE member #-} member :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool-member MultiSet {..} k = do-  HM.lookup mapMS k <&> \case-    Just i -> i > 0-    _ -> False+member ms k = stToPrim $ memberST ms k  -- | \(O(1)\) Tests whether \(k\) is not in the set. -- -- @since 1.1.0.0 {-# INLINE notMember #-} notMember :: (PrimMonad m) => MultiSet (PrimState m) -> Int -> m Bool-notMember ms k = not <$> member ms k+notMember ms k = stToPrim $ not <$> memberST ms k  -- | \(O(1)\) Increments the count of a key. --@@ -189,18 +181,7 @@ -- @since 1.1.0.0 {-# INLINE add #-} add :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()-add ms@MultiSet {..} k v = case compare v 0 of-  LT -> sub ms k (-v)-  EQ -> pure ()-  GT -> do-    HM.lookup mapMS k >>= \case-      Just n ->  do-        HM.insert mapMS k $ n + v-        when (n <= 0) $ do-          VGM.unsafeModify cntMS (+ 1) 0-      Nothing -> do-        HM.insert mapMS k v-        VGM.unsafeModify cntMS (+ 1) 0+add ms k v = stToPrim $ addST ms k v  -- | \(O(1)\) Decrements the count of a key \(k\) by \(c\). If \(c\) is negative, it falls back to -- `add`.@@ -208,35 +189,14 @@ -- @since 1.1.0.0 {-# INLINE sub #-} sub :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()-sub ms@MultiSet {..} k v = case compare v 0 of-  LT -> add ms k (-v)-  EQ -> pure ()-  GT -> do-    HM.lookup mapMS k >>= \case-      Just 0 -> pure () -- ignored-      Just n -> case compare n v of-        GT -> do-          HM.insert mapMS k (n - v)-        EQ -> do-          HM.insert mapMS k 0-          VGM.unsafeModify cntMS (subtract 1) 0-        LT -> error $ "AtCoder.Extra.Multiset.sub: the count of `" ++ show k ++ "` is becoming a negative value: `" ++ show (n - v) ++ "`"-      _ -> pure ()+sub ms k v = stToPrim $ subST ms k v  -- | \(O(1)\) Inserts a key-count pair into the set. `MultiSet` is actually a count map. -- -- @since 1.1.0.0 {-# INLINE insert #-} insert :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> Int -> m ()-insert MultiSet {..} k v-  | v <= 0 = error $ "AtCoder.Extra.Multiset.insert: new count must be positive`" ++ show k ++ "`: `" ++ show v ++ "`"-  | otherwise = do-      HM.lookup mapMS k >>= \case-        Just n | n > 0 -> do-          HM.insert mapMS k v-        _ -> do-          HM.insert mapMS k v-          VGM.unsafeModify cntMS (+ 1) 0+insert ms k v = stToPrim $ insertST ms k v  -- | \(O(1)\) Deletes a key. Note that it does not undo its insertion and does not increase the -- number of distinct keys that can be inserted into the internal hash map.@@ -244,12 +204,7 @@ -- @since 1.1.0.0 {-# INLINE delete #-} delete :: (HasCallStack, PrimMonad m) => MultiSet (PrimState m) -> Int -> m ()-delete MultiSet {..} k = do-  HM.lookup mapMS k >>= \case-    Just i | i > 0 -> do-      HM.insert mapMS k 0-      VGM.unsafeModify cntMS (subtract 1) 0-    _ -> pure ()+delete ms k = stToPrim $ deleteST ms k  -- | \(O(n)\) Enumerates the keys in the set. --@@ -292,3 +247,81 @@ {-# INLINE unsafeAssocs #-} unsafeAssocs :: (PrimMonad m) => MultiSet (PrimState m) -> m (VU.Vector (Int, Int)) unsafeAssocs = (VU.filter (\(!_, !n) -> n > 0) <$>) . HM.unsafeAssocs . mapMS++-- -------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: Int -> ST s (MultiSet s)+newST n = do+  mapMS <- HM.new n+  cntMS <- VUM.replicate 1 0+  pure $ MultiSet {..}++{-# INLINEABLE lookupST #-}+lookupST :: MultiSet s -> Int -> ST s (Maybe Int)+lookupST MultiSet {..} k = do+  HM.lookup mapMS k <&> \case+    Just i | i > 0 -> Just i+    _ -> Nothing++{-# INLINEABLE memberST #-}+memberST :: MultiSet s -> Int -> ST s Bool+memberST MultiSet {..} k = do+  HM.lookup mapMS k <&> \case+    Just i -> i > 0+    _ -> False++{-# INLINEABLE addST #-}+addST :: (HasCallStack) => MultiSet s -> Int -> Int -> ST s ()+addST ms@MultiSet {..} k v = case compare v 0 of+  LT -> subST ms k (-v)+  EQ -> pure ()+  GT -> do+    HM.lookup mapMS k >>= \case+      Just n -> do+        HM.insert mapMS k $ n + v+        when (n <= 0) $ do+          VGM.unsafeModify cntMS (+ 1) 0+      Nothing -> do+        HM.insert mapMS k v+        VGM.unsafeModify cntMS (+ 1) 0++{-# INLINEABLE subST #-}+subST :: (HasCallStack) => MultiSet s -> Int -> Int -> ST s ()+subST ms@MultiSet {..} k v = case compare v 0 of+  LT -> addST ms k (-v)+  EQ -> pure ()+  GT -> do+    HM.lookup mapMS k >>= \case+      Just 0 -> pure () -- ignored+      Just n -> case compare n v of+        GT -> do+          HM.insert mapMS k (n - v)+        EQ -> do+          HM.insert mapMS k 0+          VGM.unsafeModify cntMS (subtract 1) 0+        LT -> error $ "AtCoder.Extra.Multiset.subST: the count of `" ++ show k ++ "` is becoming a negative value: `" ++ show (n - v) ++ "`"+      _ -> pure ()++{-# INLINEABLE insertST #-}+insertST :: (HasCallStack) => MultiSet s -> Int -> Int -> ST s ()+insertST MultiSet {..} k v+  | v <= 0 = error $ "AtCoder.Extra.Multiset.insertST: new count must be positive`" ++ show k ++ "`: `" ++ show v ++ "`"+  | otherwise = do+      HM.lookup mapMS k >>= \case+        Just n | n > 0 -> do+          HM.insert mapMS k v+        _ -> do+          HM.insert mapMS k v+          VGM.unsafeModify cntMS (+ 1) 0++{-# INLINEABLE deleteST #-}+deleteST :: (HasCallStack) => MultiSet s -> Int -> ST s ()+deleteST MultiSet {..} k = do+  HM.lookup mapMS k >>= \case+    Just i | i > 0 -> do+      HM.insert mapMS k 0+      VGM.unsafeModify cntMS (subtract 1) 0+    _ -> pure ()
src/AtCoder/Extra/Pdsu.hs view
@@ -135,39 +135,13 @@   where     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.leader" v0 $ nPdsu pdsu -{-# INLINE leaderST #-}-leaderST :: (Semigroup a, VU.Unbox a) => Pdsu s a -> Int -> ST s Int-leaderST Pdsu {..} v0 = inner v0-  where-    inner v = do-      p <- VGM.read parentOrSizePdsu v-      if {- size? -} p < 0-        then pure v-        else do-          -- NOTE(perf): Path compression.-          -- Handle the nodes closer to the root first and move them onto just under the root-          !r <- inner p-          when (p /= r) $ do-            !pp <- VGM.read potentialPdsu p-            -- Move `v` to just under the root:-            VGM.write parentOrSizePdsu v {- root -} r-            -- INVARIANT: new coming monoids always come from the left. And we're performing-            -- reverse folding.-            VGM.modify potentialPdsu (<> pp) v-          pure r- -- | \(O(\alpha(n))\) Returns \(p(v)\), the potential value of vertex \(v\) relative to the -- reprensetative of its group. -- -- @since 1.1.0.0 {-# INLINE pot #-} pot :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m a-pot dsu@Pdsu {..} v1 = stToPrim $ do-  -- Perform path compression-  _ <- leaderST dsu v1-  VGM.read potentialPdsu v1-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.pot" v1 nPdsu+pot dsu v1 = stToPrim $ potST dsu v1  -- | \(O(\alpha(n))\) Returns whether the vertices \(a\) and \(b\) are in the same connected -- component.@@ -175,15 +149,7 @@ -- @since 1.1.0.0 {-# INLINE same #-} same :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m Bool-same !dsu !v1 !v2 = stToPrim $ do-  l1 <- leaderST dsu v1-  l2 <- leaderST dsu v2-  pure $ l1 == l2-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.same" v1 $ nPdsu dsu-    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.same" v2 $ nPdsu dsu---- TODO: call it unsafeDiff+same dsu v1 v2 = stToPrim $ sameST dsu v1 v2  -- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\) -- if the two vertices belong to the same group. Returns `Nothing` when the two vertices are not@@ -192,11 +158,7 @@ -- @since 1.1.0.0 {-# INLINE diff #-} diff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m (Maybe a)-diff !dsu !v1 !v2 = do-  b <- same dsu v1 v2-  if b-    then Just <$> unsafeDiff dsu v1 v2-    else pure Nothing+diff dsu v1 v2 = stToPrim $ diffST dsu v1 v2  -- | \(O(\alpha(n))\) Returns the potential of \(v_1\) relative to \(v_2\): \(p(v_1) \cdot p^{-1}(v_2)\) -- if the two vertices belong to the same group. Returns meaningless value if the two vertices are@@ -205,10 +167,7 @@ -- @since 1.1.0.0 {-# INLINE unsafeDiff #-} unsafeDiff :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> m a-unsafeDiff !dsu !v1 !v2 = do-  p1 <- pot dsu v1-  p2 <- pot dsu v2-  pure $ p1 <> invertPdsu dsu p2+unsafeDiff dsu v1 v2 = stToPrim $ unsafeDiffST dsu v1 v2  -- | \(O(\alpha(n))\) Merges \(v_1\) to \(v_2\) with differential (relative) potential -- \(\mathrm{dp}\): \(p(v1) := \mathrm{dp} \cdot p(v2)\). Returns `True` if they're newly merged.@@ -221,7 +180,87 @@     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v10 $ nPdsu dsu     !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.merge" v20 $ nPdsu dsu -{-# INLINE mergeST #-}+-- | \(O(\alpha(n))\) `merge` with the return value discarded.+--+-- @since 1.1.0.0+{-# INLINE merge_ #-}+merge_ :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m ()+merge_ !dsu !v1 !v2 !dp = stToPrim $ do+  _ <- mergeST dsu v1 v2 dp+  pure ()++-- | \(O(\alpha(n))\) Returns `True` if the two vertices belong to different groups or they belong+-- to the same group under the condition \(p(v_1) = dp \cdot p(v_2)\). It's just a convenient+-- helper function.+--+-- @since 1.1.0.0+{-# INLINE canMerge #-}+canMerge :: (HasCallStack, PrimMonad m, Semigroup a, Eq a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool+canMerge dsu v1 v2 dp = stToPrim $ canMergeST dsu v1 v2 dp++-- | \(O(\alpha(n))\) Returns the number of vertices belonging to the same group.+--+-- @since 1.1.0.0+{-# INLINE size #-}+size :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int+size !dsu !v = stToPrim $ do+  l <- leaderST dsu v+  negate <$> VGM.read (parentOrSizePdsu dsu) l+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.size" v $ nPdsu dsu++-- | \(O(n)\) Divides the graph into connected components and returns the list of them.+--+-- @since 1.1.0.0+{-# INLINE groups #-}+groups :: (PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> m (V.Vector (VU.Vector Int))+groups dsu = stToPrim $ groupsST dsu++-- | \(O(n)\) Clears the `Pdsu` to the initial state.+--+-- @since 1.1.0.0+{-# INLINE clear #-}+clear :: forall m a. (PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> m ()+clear !dsu = do+  VGM.set (potentialPdsu dsu) (mempty @a)+  VGM.set (parentOrSizePdsu dsu) (-1 {- size -})++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINE leaderST #-}+-- NOTE(perf): INLINE makes it a bit faster+leaderST :: (Semigroup a, VU.Unbox a) => Pdsu s a -> Int -> ST s Int+leaderST Pdsu {..} v0 = inner v0+  where+    inner v = do+      p <- VGM.read parentOrSizePdsu v+      if {- size? -} p < 0+        then pure v+        else do+          -- NOTE(perf): Path compression.+          -- Handle the nodes closer to the root first and move them onto just under the root+          !r <- inner p+          when (p /= r) $ do+            !pp <- VGM.read potentialPdsu p+            -- Move `v` to just under the root:+            VGM.write parentOrSizePdsu v {- root -} r+            -- INVARIANT: new coming monoids always come from the left. And we're performing+            -- reverse folding.+            VGM.modify potentialPdsu (<> pp) v+          pure r++{-# INLINEABLE potST #-}+potST :: (HasCallStack, Semigroup a, VU.Unbox a) => Pdsu s a -> Int -> ST s a+potST dsu@Pdsu {..} v1 = do+  -- Perform path compression+  _ <- leaderST dsu v1+  VGM.read potentialPdsu v1+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.potST" v1 nPdsu++{-# INLINEABLE mergeST #-} mergeST :: (HasCallStack, Monoid a, Ord a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Bool mergeST dsu@Pdsu {..} v10 v20 !dp0 = inner v10 v20 dp0   where@@ -261,24 +300,10 @@             else do               inner v2 v1 $ invertPdsu dp --- | \(O(\alpha(n))\) `merge` with the return value discarded.------ @since 1.1.0.0-{-# INLINE merge_ #-}-merge_ :: (HasCallStack, PrimMonad m, Monoid a, Ord a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m ()-merge_ !dsu !v1 !v2 !dp = do-  _ <- merge dsu v1 v2 dp-  pure ()---- | \(O(\alpha(n))\) Returns `True` if the two vertices belong to different groups or they belong--- to the same group under the condition \(p(v_1) = dp \cdot p(v_2)\). It's just a convenient--- helper function.------ @since 1.1.0.0-{-# INLINE canMerge #-}-canMerge :: (HasCallStack, PrimMonad m, Semigroup a, Eq a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> Int -> a -> m Bool-canMerge !dsu !v1 !v2 !dp = do-  b <- same dsu v1 v2+{-# INLINEABLE canMergeST #-}+canMergeST :: (HasCallStack, Semigroup a, Eq a, VU.Unbox a) => Pdsu s a -> Int -> Int -> a -> ST s Bool+canMergeST dsu v1 v2 dp = do+  b <- sameST dsu v1 v2   if not b     then pure True     else do@@ -286,23 +311,34 @@       !p2 <- VGM.read (potentialPdsu dsu) v2       pure $ p1 == dp <> p2 --- | \(O(\alpha(n))\) Returns the number of vertices belonging to the same group.------ @since 1.1.0.0-{-# INLINE size #-}-size :: (HasCallStack, PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> Int -> m Int-size !dsu !v = stToPrim $ do-  l <- leaderST dsu v-  negate <$> VGM.read (parentOrSizePdsu dsu) l+{-# INLINEABLE sameST #-}+sameST :: (HasCallStack, Semigroup a, VU.Unbox a) => Pdsu s a -> Int -> Int -> ST s Bool+sameST !dsu !v1 !v2 = stToPrim $ do+  l1 <- leaderST dsu v1+  l2 <- leaderST dsu v2+  pure $ l1 == l2   where-    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.size" v $ nPdsu dsu+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.sameST" v1 $ nPdsu dsu+    !_ = ACIA.checkIndex "AtCoder.Extra.Pdsu.sameST" v2 $ nPdsu dsu --- | \(O(n)\) Divides the graph into connected components and returns the list of them.------ @since 1.1.0.0-{-# INLINE groups #-}-groups :: (PrimMonad m, Semigroup a, VU.Unbox a) => Pdsu (PrimState m) a -> m (V.Vector (VU.Vector Int))-groups dsu@Pdsu {..} = stToPrim $ do+{-# INLINEABLE diffST #-}+diffST :: (HasCallStack, Monoid a, VU.Unbox a) => Pdsu s a -> Int -> Int -> ST s (Maybe a)+diffST dsu v1 v2 = do+  b <- sameST dsu v1 v2+  if b+    then Just <$> unsafeDiffST dsu v1 v2+    else pure Nothing++{-# INLINEABLE unsafeDiffST #-}+unsafeDiffST :: (HasCallStack, Monoid a, VU.Unbox a) => Pdsu s a -> Int -> Int -> ST s a+unsafeDiffST !dsu !v1 !v2 = do+  p1 <- potST dsu v1+  p2 <- potST dsu v2+  pure $ p1 <> invertPdsu dsu p2++{-# INLINEABLE groupsST #-}+groupsST :: (Semigroup a, VU.Unbox a) => Pdsu s a -> ST s (V.Vector (VU.Vector Int))+groupsST dsu@Pdsu {..} = do   groupSize <- VUM.replicate nPdsu (0 :: Int)   leaders <- VU.generateM nPdsu $ \i -> do     li <- leaderST dsu i@@ -316,12 +352,3 @@     VGM.write (result VG.! li) i' i     VGM.write groupSize li i'   V.filter (not . VU.null) <$> V.mapM VU.unsafeFreeze result---- | \(O(n)\) Clears the `Pdsu` to the initial state.------ @since 1.1.0.0-{-# INLINE clear #-}-clear :: forall m a. (PrimMonad m, Monoid a, VU.Unbox a) => Pdsu (PrimState m) a -> m ()-clear !dsu = do-  VGM.set (potentialPdsu dsu) (mempty @a)-  VGM.set (parentOrSizePdsu dsu) (-1 {- size -})
src/AtCoder/Extra/Pool.hs view
@@ -39,7 +39,8 @@ where  import AtCoder.Internal.Buffer qualified as B-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Coerce import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM@@ -87,18 +88,12 @@ -- | \(O(n)\) Creates a pool with the specified @capacity@. {-# INLINE new #-} new :: (VU.Unbox a, PrimMonad m) => Int -> m (Pool (PrimState m) a)-new cap = do-  dataPool <- VUM.unsafeNew cap-  freePool <- B.new cap-  nextPool <- VUM.replicate 1 (Index 0)-  pure Pool {..}+new cap = stToPrim $ newST cap  -- | \(O(1)\) Resets the pool to the initial state. {-# INLINE clear #-} clear :: (PrimMonad m) => Pool (PrimState m) a -> m ()-clear Pool {..} = do-  B.clear freePool-  VGM.unsafeWrite nextPool 0 $ Index 0+clear pool = stToPrim $ clearST pool  -- | \(O(1)\) Returns the maximum number of elements the pool can store. {-# INLINE capacity #-}@@ -108,10 +103,7 @@ -- | \(O(1)\) Returns the number of elements in the pool. {-# INLINE size #-} size :: (PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> m Int-size Pool {..} = do-  !nFree <- B.length freePool-  Index !next <- VGM.unsafeRead nextPool 0-  pure $ next - nFree+size pool = stToPrim $ sizeST pool  -- | \(O(1)\) Allocates a new element. --@@ -119,18 +111,7 @@ -- - The number of elements must not exceed the `capacity`. {-# INLINE alloc #-} alloc :: (HasCallStack, PrimMonad m, VU.Unbox a) => Pool (PrimState m) a -> a -> m Index-alloc Pool {..} !x = do-  B.popBack freePool >>= \case-    Just i -> pure i-    Nothing -> do-      Index i <- VGM.unsafeRead nextPool 0-      if i < VGM.length dataPool-        then do-          VGM.unsafeWrite nextPool 0 $ coerce (i + 1)-          VGM.write dataPool i x-          pure $ coerce i-        else do-          error "AtCoder.Extra.Pool.alloc: capacity out of bounds"+alloc pool x = stToPrim $ allocST pool x  -- | \(O(1)\) Frees an element. Be sure to not free a deleted element. --@@ -205,3 +186,43 @@ {-# INLINE invalidateHandle #-} invalidateHandle :: (PrimMonad m) => Handle (PrimState m) -> m () invalidateHandle (Handle h) = VGM.unsafeWrite h 0 undefIndex++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (VU.Unbox a) => Int -> ST s (Pool s a)+newST cap = do+  dataPool <- VUM.unsafeNew cap+  freePool <- B.new cap+  nextPool <- VUM.replicate 1 (Index 0)+  pure Pool {..}++{-# INLINEABLE clearST #-}+clearST :: Pool s a -> ST s ()+clearST Pool {..} = do+  B.clear freePool+  VGM.unsafeWrite nextPool 0 $ Index 0++{-# INLINEABLE sizeST #-}+sizeST :: (VU.Unbox a) => Pool s a -> ST s Int+sizeST Pool {..} = do+  !nFree <- B.length freePool+  Index !next <- VGM.unsafeRead nextPool 0+  pure $ next - nFree++{-# INLINEABLE allocST #-}+allocST :: (HasCallStack, VU.Unbox a) => Pool s a -> a -> ST s Index+allocST Pool {..} !x = do+  B.popBack freePool >>= \case+    Just i -> pure i+    Nothing -> do+      Index i <- VGM.unsafeRead nextPool 0+      if i < VGM.length dataPool+        then do+          VGM.unsafeWrite nextPool 0 $ coerce (i + 1)+          VGM.write dataPool i x+          pure $ coerce i+        else do+          error "AtCoder.Extra.Pool.allocST: capacity out of bounds"
src/AtCoder/Extra/Semigroup/Matrix.hs view
@@ -434,32 +434,33 @@                               pure $! m - det_                             else pure det_             det' <- swapLoop i det-            det'' <- VU.foldM'-              ( \ !acc j -> do-                  let visitDiag !det_ = do-                        aii <- read2d view i i-                        if aii == 0-                          then pure det_-                          else do-                            aji <- read2d view j i-                            let !c = m - aji `div` aii-                            rowI <- VGM.unsafeRead view i-                            rowJ <- VGM.unsafeRead view j-                            -- NOTE: it's a reverse loop!-                            VGM.ifoldrM'-                              ( \k_ aik () -> do-                                  VGM.unsafeModify rowJ ((`mod` m) . (+ aik * c)) (k_ + i)-                              )-                              ()-                              (VGM.unsafeDrop i rowI)-                            VGM.unsafeSwap view i j-                            visitDiag (m - det_)-                  acc' <- visitDiag acc-                  VGM.unsafeSwap view i j-                  pure $! m - acc'-              )-              det'-              (VU.generate (n - (i + 1)) (+ (i + 1)))+            det'' <-+              VU.foldM'+                ( \ !acc j -> do+                    let visitDiag !det_ = do+                          aii <- read2d view i i+                          if aii == 0+                            then pure det_+                            else do+                              aji <- read2d view j i+                              let !c = m - aji `div` aii+                              rowI <- VGM.unsafeRead view i+                              rowJ <- VGM.unsafeRead view j+                              -- NOTE: it's a reverse loop!+                              VGM.ifoldrM'+                                ( \k_ aik () -> do+                                    VGM.unsafeModify rowJ ((`mod` m) . (+ aik * c)) (k_ + i)+                                )+                                ()+                                (VGM.unsafeDrop i rowI)+                              VGM.unsafeSwap view i j+                              visitDiag (m - det_)+                    acc' <- visitDiag acc+                    VGM.unsafeSwap view i j+                    pure $! m - acc'+                )+                det'+                (VU.generate (n - (i + 1)) (+ (i + 1)))              inner (i + 1) det''
src/AtCoder/Extra/Seq.hs view
@@ -173,8 +173,8 @@   ) where +import AtCoder.Extra.Pool (Handle (..), invalidateHandle, newHandle, nullHandle) import AtCoder.Extra.Pool qualified as P-import AtCoder.Extra.Pool (Handle (..), newHandle, nullHandle, invalidateHandle) import AtCoder.Extra.Seq.Raw (Seq (..)) import AtCoder.Extra.Seq.Raw qualified as Seq import AtCoder.LazySegTree (SegAct (..))
src/AtCoder/Extra/Seq/Map.hs view
@@ -51,6 +51,7 @@     delete_,      -- ** Products+     -- sliceST,     prod,     prodMaybe,
src/AtCoder/Extra/Tree/Lct.hs view
@@ -212,7 +212,7 @@   VU.Vector (Vertex, Vertex) ->   -- | Link/cut tree   m (Lct (PrimState m) a)-build = buildInv id+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`).@@ -229,7 +229,205 @@   VU.Vector (Vertex, Vertex) ->   -- | Link/cut tree   m (Lct (PrimState m) a)-buildInv !invOpLct xs es = do+buildInv invOpLct xs es = stToPrim $ buildST invOpLct xs es++-- -------------------------------------------------------------------------------------------------+-- Write+-- -------------------------------------------------------------------------------------------------++-- | 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 lct v x = stToPrim $ do+  -- make @v@ the new root of the underlying tree:+  evertST lct v+  VGM.unsafeWrite (vLct lct) v x+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.write" v (nLct 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 lct f v = stToPrim $ do+  -- make @v@ the new root of the underlying tree:+  evertST lct v+  VGM.unsafeModify (vLct lct) f v+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modify" v (nLct 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 lct f v = do+  -- make @v@ the new root of the underlying tree:+  stToPrim $ evertST lct v+  VGM.unsafeModifyM (vLct lct) f v+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modifyM" v (nLct lct)++-- -------------------------------------------------------------------------------------------------+-- Link/cut operations+-- -------------------------------------------------------------------------------------------------++-- | 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 lct c p = stToPrim $ linkST lct c p+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" c (nLct lct)+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" p (nLct lct)++-- | 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 lct u v = stToPrim $ cutST lct u v+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" u (nLct lct)+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" v (nLct lct)++-- | 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 lct v = stToPrim $ evertST lct v+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.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 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 lct v = stToPrim $ exposeST lct v+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v (nLct 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_ lct v0 = stToPrim $ do+  _ <- exposeST lct v0+  pure ()+  where+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v0 (nLct lct)++-- -------------------------------------------------------------------------------------------------+-- Jump, 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 lct c0 = stToPrim $ rootST lct c0++-- | \(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 lct x = stToPrim $ parentST lct x++-- | \(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 lct u v k = stToPrim $ jumpST lct u v 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 lct u v = stToPrim $ do+  ru <- rootST lct u+  rv <- rootST lct v+  let !_ = ACIA.runtimeAssert (ru == rv) $ "AtCoder.Extra.Lct.lca: given two vertices in different connected components " ++ show (u, v)+  _ <- exposeST lct u+  exposeST lct v++-- -------------------------------------------------------------------------------------------------+-- Monoid product+-- -------------------------------------------------------------------------------------------------++-- | 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 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.Lct.prodPath" u (nLct lct)+    !_ = ACIA.checkIndex "AtCoder.Extra.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 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 lct v rootOrParent = stToPrim $ prodSubtreeST lct v rootOrParent++-- -------------------------------------------------------------------------------------------------+-- 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@@ -247,9 +445,7 @@     link lct u v   pure lct --- ---------------------------------------------------------------------------------------------------- Balancing--- -------------------------------------------------------------------------------------------------+-- * Balancing  -- | \(O(1)\) Rotates up a non-root node. {-# INLINEABLE rotateST #-}@@ -360,7 +556,7 @@ -- * Node helpers  -- | \(O(1)\)-{-# INLINEABLE isRootNodeST #-}+{-# INLINE isRootNodeST #-} isRootNodeST :: Lct s a -> Vertex -> ST s Bool isRootNodeST lct v = do   (== RootNodeLct) <$> nodePlaceST lct v@@ -474,51 +670,7 @@   let !sub' = invOpLct sub   VGM.unsafeModify midLct (<> sub') 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 lct v x = stToPrim $ do-  -- make @v@ the new root of the underlying tree:-  evertST lct v-  VGM.unsafeWrite (vLct lct) v x-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.write" v (nLct 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 lct f v = stToPrim $ do-  -- make @v@ the new root of the underlying tree:-  evertST lct v-  VGM.unsafeModify (vLct lct) f v-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modify" v (nLct 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 lct f v = do-  -- make @v@ the new root of the underlying tree:-  stToPrim $ evertST lct v-  VGM.unsafeModifyM (vLct lct) f v-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.modifyM" v (nLct lct)---- ---------------------------------------------------------------------------------------------------- Link/cut operations--- -------------------------------------------------------------------------------------------------+-- * Link/cut  -- | Amortized \(O(\log n)\). {-# INLINEABLE linkST #-}@@ -542,17 +694,6 @@   VGM.unsafeWrite rLct p c   updateNodeST lct 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 lct c p = stToPrim $ linkST lct c p-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" c (nLct lct)-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.link" p (nLct lct)- {-# INLINEABLE cutST #-} cutST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> ST s () cutST lct@Lct {pLct, lLct} u v = do@@ -584,16 +725,6 @@   VGM.unsafeWrite lLct v undefLct   updateNodeST lct 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 lct u v = stToPrim $ cutST lct u v-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" u (nLct lct)-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.cut" v (nLct 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 ()@@ -604,15 +735,6 @@   reverseNodeST lct v   pushNodeST lct 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 lct v = stToPrim $ evertST lct v-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.evert" v (nLct lct)- {-# INLINEABLE exposeST #-} exposeST :: (Monoid a, VU.Unbox a) => Lct s a -> Vertex -> ST s Vertex exposeST lct@Lct {pLct, rLct} v0 = do@@ -658,39 +780,11 @@    pure 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 lct v = stToPrim $ exposeST lct v-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v (nLct 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_ lct v0 = stToPrim $ do-  _ <- exposeST lct v0-  pure ()-  where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.expose_" v0 (nLct lct)---- ---------------------------------------------------------------------------------------------------- Jumo, LCA--- -------------------------------------------------------------------------------------------------+-- * Jump, 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 lct@Lct {lLct} c0 = stToPrim $ do+{-# 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@@ -704,14 +798,11 @@   splayST lct c'   pure c'   where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.root" c0 (nLct lct)+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.rootST" c0 (nLct 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 lct@Lct {lLct, rLct} x = stToPrim $ do+{-# 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@@ -728,7 +819,7 @@                 inner yr       Just <$> inner xl   where-    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.parent" x (nLct lct)+    !_ = ACIA.checkIndex "AtCoder.Extra.Lct.parentST" x (nLct lct)  {-# INLINEABLE jumpST #-} jumpST :: (HasCallStack, Monoid a, VU.Unbox a) => Lct s a -> Vertex -> Vertex -> Int -> ST s Vertex@@ -759,71 +850,18 @@   splayST lct res   pure 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 lct u v k = stToPrim $ jumpST lct u v 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 lct u v = stToPrim $ do-  ru <- root lct u-  rv <- root lct v-  let !_ = ACIA.runtimeAssert (ru == rv) $ "AtCoder.Extra.Lct.lca: given two vertices in different connected components " ++ show (u, v)-  _ <- exposeST lct u-  exposeST lct 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 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.Lct.prodPath" u (nLct lct)-    !_ = ACIA.checkIndex "AtCoder.Extra.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 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) =>+{-# INLINEABLE prodSubtreeST #-}+prodSubtreeST ::+  (HasCallStack, Monoid a, VU.Unbox a) =>   -- | Link/cut tree-  Lct (PrimState m) a ->+  Lct s a ->   -- | Vertex   Vertex ->   -- | Root or parent   Vertex ->   -- | Subtree's monoid product-  m a-prodSubtree lct@Lct {nLct, subtreeProdLct} v rootOrParent = stToPrim $ do+  ST s a+prodSubtreeST lct@Lct {nLct, subtreeProdLct} v rootOrParent = do   if v == rootOrParent     then do       -- `v` will be the root
src/AtCoder/Extra/Tree/TreeMonoid.hs view
@@ -111,7 +111,8 @@ import AtCoder.Internal.Assert qualified as ACIA import AtCoder.SegTree qualified as ST import Control.Monad-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Monoid (Dual (..)) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM@@ -166,23 +167,6 @@       Show     ) --- | \(O(n)\)-{-# INLINE buildImpl #-}-buildImpl ::-  (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) =>-  Hld.Hld ->-  Commutativity ->-  Hld.WeightPolicy ->-  VU.Vector a ->-  m (TreeMonoid a (PrimState m))-buildImpl hldTM commuteTM weightPolicyTM weights = do-  segFTM <- ST.build weights-  segBTM <--    case commuteTM of-      Commute -> ST.build VU.empty-      NonCommute -> ST.build $ VU.map Dual weights-  pure TreeMonoid {..}- -- | \(O(n)\) Creates a `TreeMonoid` with weights on vertices. -- -- @since 1.1.0.0@@ -197,14 +181,7 @@   VU.Vector a ->   -- | A `TreeMonoid` with weights on vertices.   m (TreeMonoid a (PrimState m))-fromVerts hld@Hld.Hld {indexHld} commuteTM xs_ = do-  let !_ = ACIA.runtimeAssert (VU.length indexHld == VU.length xs_) $ "AtCoder.Extra.Tree.TreeMonoid.fromVerts: vertex number mismatch (`" ++ show (VU.length indexHld) ++ "` and `" ++ show (VU.length xs_) ++ "`)"-  let !xs = VU.create $ do-        vec <- VUM.unsafeNew $ VU.length xs_-        VU.iforM_ xs_ $ \i x -> do-          VGM.write vec (indexHld VG.! i) x-        pure vec-  buildImpl hld commuteTM Hld.WeightsAreOnVertices xs+fromVerts hld commuteTM xs_ = stToPrim $ fromVertsST hld commuteTM xs_  -- | \(O(n)\) Creates a `TreeMonoid` with weignts on edges. The edges are not required to be -- duplicated: only one of \((u, v, w)\) or \((v, u, w)\) is needed.@@ -221,83 +198,49 @@   VU.Vector (Vertex, Vertex, a) ->   -- | A `TreeMonoid` with weights on edges.   m (TreeMonoid a (PrimState m))-fromEdges hld@Hld.Hld {indexHld} commuteTM edges = do-  let !xs = VU.create $ do-        vec <- VUM.unsafeNew $ VU.length indexHld-        VU.forM_ edges $ \(!u, !v, !w) -> do-          let u' = indexHld VG.! u-          let v' = indexHld VG.! v-          VGM.write vec (max u' v') w-        pure vec-  buildImpl hld commuteTM Hld.WeightsAreOnEdges xs+fromEdges hld commuteTM edges = stToPrim $ fromEdgesST hld commuteTM edges  -- | \(O(\log^2 n)\) Returns the product of the path between two vertices \(u\), \(v\) (invlusive). -- -- @since 1.1.0.0 {-# INLINE prod #-} prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> Vertex -> m a-prod TreeMonoid {..} u v = do-  case commuteTM of-    Commute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (ST.prod segFTM) u v-    NonCommute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (\l r -> getDual <$> ST.prod segBTM l r) u v+prod tm u v = stToPrim $ prodST tm u v  -- | \(O(\log n)\) Returns the product of the subtree rooted at the given `Vertex`. -- -- @since 1.1.0.0 {-# INLINE prodSubtree #-} prodSubtree :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a-prodSubtree TreeMonoid {..} subtreeRoot = do-  let (!l, !r) = Hld.subtreeSegmentInclusive hldTM subtreeRoot-  case weightPolicyTM of-    Hld.WeightsAreOnVertices -> ST.prod segFTM l (r + 1)-    Hld.WeightsAreOnEdges -> do-      -- ignore the root of the subtree, which has the minimum index among the subtree vertices-      if l == r-        then pure mempty-        else ST.prod segFTM (l + 1) (r + 1)+prodSubtree tm subtreeRoot = stToPrim $ prodSubtreeST tm subtreeRoot  -- | \(O(1)\) Reads a monoid value of a `Vertex`. -- -- @since 1.1.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> m a-read TreeMonoid {..} i_ = do-  let !i = Hld.indexHld hldTM VG.! i_-  ST.read segFTM i+read tm i_ = stToPrim $ readST tm i_  -- | \(O(\log n)\) Writes to the monoid value of a vertex. -- -- @since 1.1.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m ()-write TreeMonoid {..} i_ x = do-  let !i = Hld.indexHld hldTM VG.! i_-  ST.write segFTM i x-  when (commuteTM == NonCommute) $ do-    ST.write segBTM i $ Dual x+write tm i_ x = stToPrim $ writeST tm i_ x  -- | \(O(\log n)\) Writes to the monoid value of a vertex and returns the old value. -- -- @since 1.1.0.0 {-# INLINE exchange #-} exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> Vertex -> a -> m a-exchange TreeMonoid {..} i_ x = do-  let !i = Hld.indexHld hldTM VG.! i_-  !res <- ST.exchange segFTM i x-  when (commuteTM == NonCommute) $ do-    ST.write segBTM i $ Dual x-  pure res+exchange tm i_ x = stToPrim $ exchangeST tm i_ x  -- | \(O(\log n)\) Modifies the monoid value of a vertex with a pure function. -- -- @since 1.1.0.0 {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => TreeMonoid a (PrimState m) -> (a -> a) -> Int -> m ()-modify TreeMonoid {..} f i_ = do-  let !i = Hld.indexHld hldTM VG.! i_-  ST.modify segFTM f i-  when (commuteTM == NonCommute) $ do-    ST.modify segBTM (Dual . f . getDual) i+modify tm f i_ = stToPrim $ modifyST tm f i_  -- | \(O(\log n)\) Modifies the monoid value of a vertex with a monadic function. --@@ -309,3 +252,106 @@   ST.modifyM segFTM f i   when (commuteTM == NonCommute) $ do     ST.modifyM segBTM ((Dual <$>) . f . getDual) i++-- -------------------------------------------------------------------------------------------------+-- INLINE+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE buildST #-}+buildST ::+  (HasCallStack, Monoid a, VU.Unbox a) =>+  Hld.Hld ->+  Commutativity ->+  Hld.WeightPolicy ->+  VU.Vector a ->+  ST s (TreeMonoid a s)+buildST hldTM commuteTM weightPolicyTM weights = do+  segFTM <- ST.build weights+  segBTM <-+    case commuteTM of+      Commute -> ST.build VU.empty+      NonCommute -> ST.build $ VU.map Dual weights+  pure TreeMonoid {..}++{-# INLINEABLE fromVertsST #-}+fromVertsST ::+  (HasCallStack, Monoid a, VU.Unbox a) =>+  Hld.Hld ->+  Commutativity ->+  VU.Vector a ->+  ST s (TreeMonoid a s)+fromVertsST hld@Hld.Hld {indexHld} commuteTM xs_ = do+  let !_ = ACIA.runtimeAssert (VU.length indexHld == VU.length xs_) $ "AtCoder.Extra.Tree.TreeMonoid.fromVertsST: vertex number mismatch (`" ++ show (VU.length indexHld) ++ "` and `" ++ show (VU.length xs_) ++ "`)"+  let !xs = VU.create $ do+        vec <- VUM.unsafeNew $ VU.length xs_+        VU.iforM_ xs_ $ \i x -> do+          VGM.write vec (indexHld VG.! i) x+        pure vec+  buildST hld commuteTM Hld.WeightsAreOnVertices xs++{-# INLINEABLE fromEdgesST #-}+fromEdgesST ::+  (HasCallStack, Monoid a, VU.Unbox a) =>+  Hld.Hld ->+  Commutativity ->+  VU.Vector (Vertex, Vertex, a) ->+  ST s (TreeMonoid a s)+fromEdgesST hld@Hld.Hld {indexHld} commuteTM edges = do+  let !xs = VU.create $ do+        vec <- VUM.unsafeNew $ VU.length indexHld+        VU.forM_ edges $ \(!u, !v, !w) -> do+          let u' = indexHld VG.! u+          let v' = indexHld VG.! v+          VGM.write vec (max u' v') w+        pure vec+  buildST hld commuteTM Hld.WeightsAreOnEdges xs++{-# INLINEABLE prodST #-}+prodST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> Vertex -> ST s a+prodST TreeMonoid {..} u v = do+  case commuteTM of+    Commute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (ST.prod segFTM) u v+    NonCommute -> Hld.prod weightPolicyTM hldTM (ST.prod segFTM) (\l r -> getDual <$> ST.prod segBTM l r) u v++{-# INLINEABLE prodSubtreeST #-}+prodSubtreeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> ST s a+prodSubtreeST TreeMonoid {..} subtreeRoot = do+  let (!l, !r) = Hld.subtreeSegmentInclusive hldTM subtreeRoot+  case weightPolicyTM of+    Hld.WeightsAreOnVertices -> ST.prod segFTM l (r + 1)+    Hld.WeightsAreOnEdges -> do+      -- ignore the root of the subtree, which has the minimum index among the subtree vertices+      if l == r+        then pure mempty+        else ST.prod segFTM (l + 1) (r + 1)++{-# INLINEABLE readST #-}+readST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> ST s a+readST TreeMonoid {..} i_ = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.read segFTM i++{-# INLINEABLE writeST #-}+writeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> a -> ST s ()+writeST TreeMonoid {..} i_ x = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.write segFTM i x+  when (commuteTM == NonCommute) $ do+    ST.write segBTM i $ Dual x++{-# INLINEABLE exchangeST #-}+exchangeST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> Vertex -> a -> ST s a+exchangeST TreeMonoid {..} i_ x = do+  let !i = Hld.indexHld hldTM VG.! i_+  !res <- ST.exchange segFTM i x+  when (commuteTM == NonCommute) $ do+    ST.write segBTM i $ Dual x+  pure res++{-# INLINEABLE modifyST #-}+modifyST :: (HasCallStack, Monoid a, VU.Unbox a) => TreeMonoid a s -> (a -> a) -> Int -> ST s ()+modifyST TreeMonoid {..} f i_ = do+  let !i = Hld.indexHld hldTM VG.! i_+  ST.modify segFTM f i+  when (commuteTM == NonCommute) $ do+    ST.modify segBTM (Dual . f . getDual) i
src/AtCoder/Extra/WaveletMatrix.hs view
@@ -99,14 +99,14 @@ -- original array if you can. -- -- @since 1.1.0.0-{-# INLINABLE access #-}+{-# INLINEABLE access #-} access :: WaveletMatrix -> Int -> Maybe Int access WaveletMatrix {..} i = (xDictWM VG.!) <$> Rwm.access rawWM i  -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r)\). -- -- @since 1.1.0.0-{-# INLINABLE rank #-}+{-# INLINEABLE rank #-} rank ::   -- | A wavelet matrix   WaveletMatrix ->@@ -123,7 +123,7 @@ -- | \(O(\log |S|)\) Returns the number of \(y\) in \([l, r) \times [y_1, y_2)\). -- -- @since 1.1.0.0-{-# INLINABLE rankBetween #-}+{-# INLINEABLE rankBetween #-} rankBetween ::   -- | A wavelet matrix   WaveletMatrix ->@@ -151,7 +151,7 @@ -- not found. -- -- @since 1.1.0.0-{-# INLINABLE select #-}+{-# INLINEABLE select #-} select :: WaveletMatrix -> Int -> Maybe Int select wm = selectKth wm 0 @@ -159,7 +159,7 @@ -- if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINABLE selectKth #-}+{-# INLINEABLE selectKth #-} selectKth ::   -- | A wavelet matrix   WaveletMatrix ->@@ -180,7 +180,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINABLE selectIn #-}+{-# INLINEABLE selectIn #-} selectIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -198,7 +198,7 @@ -- (0-based) of \(y\) in the sequence, or `Nothing` if no such occurrence exists. -- -- @since 1.1.0.0-{-# INLINABLE selectKthIn #-}+{-# INLINEABLE selectKthIn #-} selectKthIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -223,7 +223,7 @@ -- largest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINABLE kthLargestIn #-}+{-# INLINEABLE kthLargestIn #-} kthLargestIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -243,7 +243,7 @@ -- \(k\)-th (0-based) largest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINABLE ikthLargestIn #-}+{-# INLINEABLE ikthLargestIn #-} ikthLargestIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -263,7 +263,7 @@ -- smallest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINABLE kthSmallestIn #-}+{-# INLINEABLE kthSmallestIn #-} kthSmallestIn ::   -- | A wavelet matrix   WaveletMatrix ->@@ -283,7 +283,7 @@ -- \(k\)-th (0-based) smallest value. Note that duplicated values are treated as distinct occurrences. -- -- @since 1.1.0.0-{-# INLINABLE ikthSmallestIn #-}+{-# INLINEABLE ikthSmallestIn #-} ikthSmallestIn ::   WaveletMatrix ->   -- | \(l\)@@ -301,7 +301,7 @@ -- | \(O(\log |S|)\) -- -- @since 1.1.0.0-{-# INLINABLE unsafeKthSmallestIn #-}+{-# INLINEABLE unsafeKthSmallestIn #-} unsafeKthSmallestIn :: WaveletMatrix -> Int -> Int -> Int -> Int unsafeKthSmallestIn WaveletMatrix {..} l r k =   xDictWM VG.! Rwm.unsafeKthSmallestIn rawWM l r k@@ -309,7 +309,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0]\). -- -- @since 1.1.0.0-{-# INLINABLE lookupLE #-}+{-# INLINEABLE lookupLE #-} lookupLE ::   -- | A wavelet matrix   WaveletMatrix ->@@ -334,7 +334,7 @@ -- | \(O(\log |S|)\) Looks up the maximum \(y\) in \([l, r) \times (-\infty, y_0)\). -- -- @since 1.1.0.0-{-# INLINABLE lookupLT #-}+{-# INLINEABLE lookupLT #-} lookupLT ::   -- | A wavelet matrix   WaveletMatrix ->@@ -351,7 +351,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times [y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINABLE lookupGE #-}+{-# INLINEABLE lookupGE #-} lookupGE ::   -- | A wavelet matrix   WaveletMatrix ->@@ -376,7 +376,7 @@ -- | \(O(\log |S|)\) Looks up the minimum \(y\) in \([l, r) \times (y_0, \infty)\). -- -- @since 1.1.0.0-{-# INLINABLE lookupGT #-}+{-# INLINEABLE lookupGT #-} lookupGT ::   -- | A wavelet matrix   WaveletMatrix ->@@ -394,7 +394,7 @@ -- ascending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINABLE assocsIn #-}+{-# INLINEABLE assocsIn #-} assocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] assocsIn WaveletMatrix {..} l r = Rwm.assocsWith rawWM l r (xDictWM VG.!) @@ -402,6 +402,6 @@ -- descending order of \(y\). Note that it's only fast when the \(|S|\) is very small. -- -- @since 1.1.0.0-{-# INLINABLE descAssocsIn #-}+{-# INLINEABLE descAssocsIn #-} descAssocsIn :: WaveletMatrix -> Int -> Int -> [(Int, Int)] descAssocsIn WaveletMatrix {..} l r = Rwm.descAssocsInWith rawWM l r (xDictWM VG.!)
src/AtCoder/FenwickTree.hs view
@@ -59,6 +59,7 @@ import Control.Monad (when) import Control.Monad.Fix (fix) import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bits import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU@@ -90,11 +91,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => Int -> m (FenwickTree (PrimState m) a)-new nFt-  | nFt >= 0 = do-      dataFt <- VUM.replicate nFt 0-      pure FenwickTree {..}-  | otherwise = error $ "AtCoder.FenwickTree.new: given negative size `" ++ show nFt ++ "`"+new = stToPrim . newST  -- | Creates `FenwickTree` with initial values. --@@ -102,12 +99,9 @@ -- - \(O(n)\) -- -- @since 1.0.0.0-build :: (PrimMonad m, Num a, VU.Unbox a) => VU.Vector a -> m (FenwickTree (PrimState m) a) {-# INLINE build #-}-build xs = do-  ft <- new $ VU.length xs-  VU.iforM_ xs $ add ft-  pure ft+build :: (PrimMonad m, Num a, VU.Unbox a) => VU.Vector a -> m (FenwickTree (PrimState m) a)+build = stToPrim . buildST  -- | Adds \(x\) to \(p\)-th value of the array. --@@ -120,26 +114,14 @@ -- @since 1.0.0.0 {-# INLINE add #-} add :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> a -> m ()-add FenwickTree {..} p0 x = do-  let !_ = ACIA.checkIndex "AtCoder.FenwickTree.add" p0 nFt-  let p1 = p0 + 1-  flip fix p1 $ \loop p -> do-    when (p <= nFt) $ do-      VGM.modify dataFt (+ x) (p - 1)-      loop $! p + (p .&. (-p))+add ft p0 x = stToPrim $ addST ft p0 x  -- | \(O(\log n)\) Calculates the sum in a half-open interval @[0, r)@. -- -- @since 1.0.0.0 {-# INLINE prefixSum #-} prefixSum :: (PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> m a-prefixSum FenwickTree {..} = inner 0-  where-    inner !acc !r-      | r <= 0 = pure acc-      | otherwise = do-          dx <- VGM.read dataFt (r - 1)-          inner (acc + dx) (r - r .&. (-r))+prefixSum ft r = stToPrim $ prefixSumST ft r  -- | Calculates the sum in a half-open interval \([l, r)\). --@@ -152,9 +134,7 @@ -- @since 1.0.0.0 {-# INLINE sum #-} sum :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m a-sum ft@FenwickTree {nFt} l r-  | not (ACIA.testInterval l r nFt) = ACIA.errorInterval "AtCoder.FenwickTree.sum" l r nFt-  | otherwise = unsafeSum ft l r+sum ft l r = stToPrim $ sumST ft l r  -- | Total variant of `sum`. Calculates the sum in a half-open interval \([l, r)\). It returns -- `Nothing` if the interval is invalid.@@ -165,17 +145,7 @@ -- @since 1.0.0.0 {-# INLINE sumMaybe #-} sumMaybe :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m (Maybe a)-sumMaybe ft@FenwickTree {nFt} l r-  | not (ACIA.testInterval l r nFt) = pure Nothing-  | otherwise = Just <$> unsafeSum ft l r---- | Internal implementation of `sum`.-{-# INLINE unsafeSum #-}-unsafeSum :: (HasCallStack, PrimMonad m, Num a, VU.Unbox a) => FenwickTree (PrimState m) a -> Int -> Int -> m a-unsafeSum ft l r = do-  xr <- prefixSum ft r-  xl <- prefixSum ft l-  pure $! xr - xl+sumMaybe ft l r = stToPrim $ sumMaybeST ft l r  -- | (Extra API) Applies a binary search on the Fenwick tree. It returns an index \(r\) that -- satisfies both of the following.@@ -359,3 +329,61 @@                   else inner2 (i + bit k) k t        (+ 1) <$> inner2 i0 k0 s0++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (HasCallStack, Num a, VU.Unbox a) => Int -> ST s (FenwickTree s a)+newST nFt+  | nFt >= 0 = do+      dataFt <- VUM.replicate nFt 0+      pure FenwickTree {..}+  | otherwise = error $ "AtCoder.FenwickTree.newST: given negative size `" ++ show nFt ++ "`"++{-# INLINEABLE buildST #-}+buildST :: (Num a, VU.Unbox a) => VU.Vector a -> ST s (FenwickTree s a)+buildST xs = do+  ft <- new $ VU.length xs+  VU.iforM_ xs $ add ft+  pure ft++{-# INLINEABLE addST #-}+addST :: (HasCallStack, Num a, VU.Unbox a) => FenwickTree s a -> Int -> a -> ST s ()+addST FenwickTree {..} p0 x = do+  let !_ = ACIA.checkIndex "AtCoder.FenwickTree.addST" p0 nFt+  let p1 = p0 + 1+  flip fix p1 $ \loop p -> do+    when (p <= nFt) $ do+      VGM.modify dataFt (+ x) (p - 1)+      loop $! p + (p .&. (-p))++{-# INLINEABLE prefixSumST #-}+prefixSumST :: (Num a, VU.Unbox a) => FenwickTree s a -> Int -> ST s a+prefixSumST FenwickTree {..} = inner 0+  where+    inner !acc !r+      | r <= 0 = pure acc+      | otherwise = do+          dx <- VGM.read dataFt (r - 1)+          inner (acc + dx) (r - r .&. (-r))++{-# INLINEABLE sumST #-}+sumST :: (HasCallStack, Num a, VU.Unbox a) => FenwickTree s a -> Int -> Int -> ST s a+sumST ft@FenwickTree {nFt} l r+  | not (ACIA.testInterval l r nFt) = ACIA.errorInterval "AtCoder.FenwickTree.sumST" l r nFt+  | otherwise = unsafeSumST ft l r++{-# INLINEABLE sumMaybeST #-}+sumMaybeST :: (HasCallStack, Num a, VU.Unbox a) => FenwickTree s a -> Int -> Int -> ST s (Maybe a)+sumMaybeST ft@FenwickTree {nFt} l r+  | not (ACIA.testInterval l r nFt) = pure Nothing+  | otherwise = Just <$> unsafeSumST ft l r++{-# INLINEABLE unsafeSumST #-}+unsafeSumST :: (HasCallStack, Num a, VU.Unbox a) => FenwickTree s a -> Int -> Int -> ST s a+unsafeSumST ft l r = do+  xr <- prefixSumST ft r+  xl <- prefixSumST ft l+  pure $! xr - xl
src/AtCoder/Internal/Barrett.hs view
@@ -16,11 +16,14 @@ module AtCoder.Internal.Barrett   ( -- * Barrett     Barrett,+     -- * Constructor     new32,     new64,+     -- * Accessor     umod,+     -- * Barrett reduction     mulMod,   )
src/AtCoder/Internal/Buffer.hs view
@@ -88,7 +88,8 @@ where  import AtCoder.Internal.Assert qualified as ACIA-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM@@ -108,20 +109,14 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (PrimMonad m, VU.Unbox a) => Int -> m (Buffer (PrimState m) a)-new n = do-  lenB <- VUM.replicate 1 (0 :: Int)-  vecB <- VUM.unsafeNew n-  pure Buffer {..}+new n = stToPrim $ newST n  -- | \(O(n)\) Creates a buffer with capacity \(n\) with initial values. -- -- @since 1.0.0.0 {-# INLINE build #-} build :: (PrimMonad m, VU.Unbox a) => VU.Vector a -> m (Buffer (PrimState m) a)-build xs = do-  lenB <- VUM.replicate 1 $ VU.length xs-  vecB <- VU.thaw xs-  pure Buffer {..}+build xs = stToPrim $ buildST xs  -- | \(O(1)\) Returns the array size. --@@ -135,8 +130,7 @@ -- @since 1.0.0.0 {-# INLINE length #-} length :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m Int-length Buffer {..} = do-  VGM.read lenB 0+length Buffer {..} = VGM.read lenB 0  -- | \(O(1)\) Returns `True` if the buffer is empty. --@@ -150,13 +144,7 @@ -- @since 1.0.0.0 {-# INLINE back #-} back :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a)-back Buffer {..} = do-  len <- VGM.read lenB 0-  if len == 0-    then pure Nothing-    else do-      x <- VGM.read vecB (len - 1)-      pure $ Just x+back = stToPrim . backST  -- | \(O(1)\) Yields the element at the given position. Will throw an exception if the index is out -- of range.@@ -164,54 +152,35 @@ -- @since 1.0.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> Int -> m a-read Buffer {..} i = do-  len <- VGM.read lenB 0-  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.read" i len-  VGM.read vecB i+read buf i = stToPrim $ readST buf i  -- | \(O(1)\) Yields the element at the given position, or `Nothing` if the index is out of range. -- -- @since 1.2.1.0 {-# INLINE readMaybe #-} readMaybe :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> Int -> m (Maybe a)-readMaybe Buffer {..} i = do-  len <- VGM.read lenB 0-  if ACIA.testIndex i len-    then Just <$> VGM.unsafeRead vecB i-    else pure Nothing+readMaybe buf i = stToPrim $ readMaybeST buf i  -- | \(O(1)\) Appends an element to the back. -- -- @since 1.0.0.0 {-# INLINE pushBack #-} pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> a -> m ()-pushBack Buffer {..} e = do-  len <- VGM.read lenB 0-  VGM.write vecB len e-  VGM.write lenB 0 (len + 1)+pushBack buf e = stToPrim $ pushBackST buf e  -- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty. -- -- @since 1.0.0.0 {-# INLINE popBack #-} popBack :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (Maybe a)-popBack Buffer {..} = do-  len <- VGM.read lenB 0-  if len == 0-    then pure Nothing-    else do-      x <- VGM.read vecB (len - 1)-      VGM.write lenB 0 (len - 1)-      pure $ Just x+popBack buf = stToPrim $ popBackST buf  -- | \(O(1)\) Removes the last element from the buffer and discards it. -- -- @since 1.1.1.0 {-# INLINE popBack_ #-} popBack_ :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m ()-popBack_ buf = do-  _ <- popBack buf-  pure ()+popBack_ buf = stToPrim $ popBackST_ buf  -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of bounds.@@ -219,10 +188,7 @@ -- @since 1.0.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> Int -> a -> m ()-write Buffer {..} i e = do-  len <- VGM.read lenB 0-  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.write" i len-  VGM.write vecB i e+write buf i e = stToPrim $ writeST buf i e  -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of bounds.@@ -230,19 +196,16 @@ -- @since 1.0.0.0 {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> (a -> a) -> Int -> m ()-modify Buffer {..} f i = do-  len <- VGM.read lenB 0-  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.modify" i len-  VGM.modify vecB f i+modify buf f i = stToPrim $ modifyST buf f i  -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of bounds. -- -- @since 1.0.0.0-{-# INLINE modifyM #-}+{-# INLINEABLE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> (a -> m a) -> Int -> m () modifyM Buffer {..} f i = do-  len <- VGM.read lenB 0+  len <- stToPrim $ VGM.read lenB 0   let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.modifyM" i len   VGM.modifyM vecB f i @@ -251,7 +214,7 @@ -- @since 1.0.0.0 {-# INLINE clear #-} clear :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m ()-clear Buffer {..} = do+clear Buffer {..} = stToPrim $ do   VGM.write lenB 0 0  -- | \(O(n)\) Yields an immutable copy of the mutable vector.@@ -259,9 +222,7 @@ -- @since 1.0.0.0 {-# INLINE freeze #-} freeze :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (VU.Vector a)-freeze Buffer {..} = do-  len <- VGM.read lenB 0-  VU.freeze $ VUM.take len vecB+freeze = stToPrim . freezeST  -- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable -- vector may not be used after this operation.@@ -269,6 +230,97 @@ -- @since 1.0.0.0 {-# INLINE unsafeFreeze #-} unsafeFreeze :: (PrimMonad m, VU.Unbox a) => Buffer (PrimState m) a -> m (VU.Vector a)-unsafeFreeze Buffer {..} = do+unsafeFreeze = stToPrim . unsafeFreezeST++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (VU.Unbox a) => Int -> ST s (Buffer s a)+newST n = do+  lenB <- VUM.replicate 1 (0 :: Int)+  vecB <- VUM.unsafeNew n+  pure Buffer {..}++{-# INLINEABLE buildST #-}+buildST :: (VU.Unbox a) => VU.Vector a -> ST s (Buffer s a)+buildST xs = do+  lenB <- VUM.replicate 1 $ VU.length xs+  vecB <- VU.thaw xs+  pure Buffer {..}++{-# INLINEABLE backST #-}+backST :: (VU.Unbox a) => Buffer s a -> ST s (Maybe a)+backST Buffer {..} = do+  len <- VGM.read lenB 0+  if len == 0+    then pure Nothing+    else do+      x <- VGM.read vecB (len - 1)+      pure $ Just x++{-# INLINEABLE readST #-}+readST :: (HasCallStack, VU.Unbox a) => Buffer s a -> Int -> ST s a+readST Buffer {..} i = do+  len <- VGM.read lenB 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.read" i len+  VGM.read vecB i++{-# INLINEABLE readMaybeST #-}+readMaybeST :: (VU.Unbox a) => Buffer s a -> Int -> ST s (Maybe a)+readMaybeST Buffer {..} i = do+  len <- VGM.read lenB 0+  if ACIA.testIndex i len+    then Just <$> VGM.unsafeRead vecB i+    else pure Nothing++{-# INLINEABLE pushBackST #-}+pushBackST :: (HasCallStack, VU.Unbox a) => Buffer s a -> a -> ST s ()+pushBackST Buffer {..} e = do+  len <- VGM.read lenB 0+  VGM.write vecB len e+  VGM.write lenB 0 (len + 1)++{-# INLINEABLE popBackST #-}+popBackST :: (VU.Unbox a) => Buffer s a -> ST s (Maybe a)+popBackST Buffer {..} = do+  len <- VGM.read lenB 0+  if len == 0+    then pure Nothing+    else do+      x <- VGM.read vecB (len - 1)+      VGM.write lenB 0 (len - 1)+      pure $ Just x++{-# INLINEABLE popBackST_ #-}+popBackST_ :: (VU.Unbox a) => Buffer s a -> ST s ()+popBackST_ buf = do+  _ <- popBack buf+  pure ()++{-# INLINEABLE writeST #-}+writeST :: (HasCallStack, VU.Unbox a) => Buffer s a -> Int -> a -> ST s ()+writeST Buffer {..} i e = do+  len <- VGM.read lenB 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.write" i len+  VGM.write vecB i e++{-# INLINEABLE modifyST #-}+modifyST :: (HasCallStack, VU.Unbox a) => Buffer s a -> (a -> a) -> Int -> ST s ()+modifyST Buffer {..} f i = do+  len <- VGM.read lenB 0+  let !_ = ACIA.checkIndex "AtCoder.Internal.Buffer.modify" i len+  VGM.modify vecB f i++{-# INLINEABLE freezeST #-}+freezeST :: (VU.Unbox a) => Buffer s a -> ST s (VU.Vector a)+freezeST Buffer {..} = do+  len <- VGM.read lenB 0+  VU.freeze $ VUM.take len vecB++{-# INLINEABLE unsafeFreezeST #-}+unsafeFreezeST :: (VU.Unbox a) => Buffer s a -> ST s (VU.Vector a)+unsafeFreezeST Buffer {..} = do   len <- VGM.read lenB 0   VU.unsafeFreeze $ VUM.take len vecB
src/AtCoder/Internal/Convolution.hs view
@@ -58,7 +58,7 @@ -- | \(O(\log m)\) Creates an `FftInfo`. -- -- @since 1.0.0.0-{-# INLINABLE newInfo #-}+{-# INLINEABLE newInfo #-} newInfo :: forall s p. (AM.Modulus p) => ST s (FftInfo p) newInfo = do   let !g = AM.primitiveRootModulus (proxy# @p)@@ -111,7 +111,7 @@   pure FftInfo {..}  -- | @since 1.0.0.0-{-# INLINABLE butterfly #-}+{-# INLINEABLE butterfly #-} butterfly ::   forall s p.   (AM.Modulus p) =>@@ -175,7 +175,7 @@           loop $ len + 2  -- | @since 1.0.0.0-{-# INLINABLE butterflyInv #-}+{-# INLINEABLE butterflyInv #-} butterflyInv ::   forall s p.   (AM.Modulus p) =>@@ -240,7 +240,7 @@           loop $ len - 2  -- | @since 1.0.0.0-{-# INLINABLE convolutionNaive #-}+{-# INLINEABLE convolutionNaive #-} convolutionNaive ::   forall p.   (AM.Modulus p) =>@@ -263,7 +263,7 @@   pure ans  -- | @since 1.0.0.0-{-# INLINE convolutionFft #-}+{-# INLINEABLE convolutionFft #-} convolutionFft ::   forall p.   (AM.Modulus p) =>
src/AtCoder/Internal/Csr.hs view
@@ -134,7 +134,7 @@ {-# INLINE build1 #-} build1 :: (HasCallStack) => Int -> VU.Vector (Int, Int) -> Csr Int build1 n edges = build n $ VU.zip3 us vs (VU.replicate (VU.length us) (1 :: Int))- where+  where     (!us, !vs) = VU.unzip edges  -- | \(O(1)\) Returns the adjacent vertices.
src/AtCoder/Internal/GrowVec.hs view
@@ -83,7 +83,8 @@  import AtCoder.Internal.Assert qualified as ACIA import Control.Monad (when)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU@@ -162,23 +163,14 @@ -- @since 1.0.0.0 {-# INLINE read #-} read :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m a-read GrowVec {..} i = do-  vec <- readMutVar vecGV-  let len = VUM.length vec-  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.read" i len-  VGM.read vec i+read gv i = stToPrim $ readST gv i  -- | \(O(1)\) Yields the element at the given position, or `Nothing` if the index is out of range. -- -- @since 1.2.1.0 {-# INLINE readMaybe #-} readMaybe :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> m (Maybe a)-readMaybe GrowVec {..} i = do-  vec <- readMutVar vecGV-  len <- VGM.unsafeRead posGV 0-  if ACIA.testIndex i len-    then Just <$> VGM.unsafeRead vec i-    else pure Nothing+readMaybe gv i = stToPrim $ readMaybeST gv i  -- | \(O(1)\) Writes to the element at the given position. Will throw an exception if the index is -- out of range.@@ -186,19 +178,77 @@ -- @since 1.0.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> Int -> a -> m ()-write GrowVec {..} i x = do-  vec <- readMutVar vecGV-  let len = VUM.length vec-  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.write" i len-  VGM.write vec i x+write gv i x = stToPrim $ writeST gv i x  -- | Amortized \(O(1)\). Grow the capacity twice -- -- @since 1.0.0.0 {-# INLINE pushBack #-} pushBack :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> a -> m ()-pushBack GrowVec {..} e = do+pushBack gv e = stToPrim $ pushBackST gv e++-- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty.+--+-- @since 1.0.0.0+{-# INLINE popBack #-}+popBack :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (Maybe a)+popBack = stToPrim . popBackST++-- | \(O(1)\) `popBack` with the return value discarded.+--+-- @since 1.0.0.0+{-# INLINE popBack_ #-}+popBack_ :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m ()+popBack_ = stToPrim . popBackST_++-- | \(O(n)\) Yields an immutable copy of the mutable vector.+--+-- @since 1.0.0.0+{-# INLINE freeze #-}+freeze :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (VU.Vector a)+freeze = stToPrim . freezeST++-- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable+-- vector may not be used after this operation.+--+-- @since 1.0.0.0+{-# INLINE unsafeFreeze #-}+unsafeFreeze :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (VU.Vector a)+unsafeFreeze = stToPrim . unsafeFreezeST++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE readST #-}+readST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s a+readST GrowVec {..} i = do+  vec <- readMutVar vecGV+  let len = VUM.length vec+  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.read" i len+  VGM.read vec i++{-# INLINEABLE readMaybeST #-}+readMaybeST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> ST s (Maybe a)+readMaybeST GrowVec {..} i = do+  vec <- readMutVar vecGV   len <- VGM.unsafeRead posGV 0+  if ACIA.testIndex i len+    then Just <$> VGM.unsafeRead vec i+    else pure Nothing++{-# INLINEABLE writeST #-}+writeST :: (HasCallStack, VU.Unbox a) => GrowVec s a -> Int -> a -> ST s ()+writeST GrowVec {..} i x = do+  vec <- readMutVar vecGV+  let len = VUM.length vec+  let !_ = ACIA.checkIndex "AtCoder.Internal.GrowVec.write" i len+  VGM.write vec i x++{-# INLINEABLE pushBackST #-}+pushBackST :: (VU.Unbox a) => GrowVec s a -> a -> ST s ()+pushBackST GrowVec {..} e = do+  len <- VGM.unsafeRead posGV 0   vec <- do     vec <- readMutVar vecGV     if VUM.length vec > len@@ -217,12 +267,9 @@     )     0 --- | \(O(1)\) Removes the last element from the buffer and returns it, or `Nothing` if it is empty.------ @since 1.0.0.0-{-# INLINE popBack #-}-popBack :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (Maybe a)-popBack GrowVec {..} = do+{-# INLINEABLE popBackST #-}+popBackST :: (VU.Unbox a) => GrowVec s a -> ST s (Maybe a)+popBackST GrowVec {..} = do   pos <- VGM.unsafeRead posGV 0   if pos <= 0     then pure Nothing@@ -231,32 +278,22 @@       vec <- readMutVar vecGV       Just <$> VGM.read vec (pos - 1) --- | \(O(1)\) `popBack` with the return value discarded.------ @since 1.0.0.0-{-# INLINE popBack_ #-}-popBack_ :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m ()-popBack_ GrowVec {..} = do+{-# INLINEABLE popBackST_ #-}+popBackST_ :: (VU.Unbox a) => GrowVec s a -> ST s ()+popBackST_ GrowVec {..} = do   pos <- VGM.unsafeRead posGV 0   VGM.unsafeWrite posGV 0 $ max 0 $ pos - 1 --- | \(O(n)\) Yields an immutable copy of the mutable vector.------ @since 1.0.0.0-{-# INLINE freeze #-}-freeze :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (VU.Vector a)-freeze GrowVec {..} = do+{-# INLINEABLE freezeST #-}+freezeST :: (VU.Unbox a) => GrowVec s a -> ST s (VU.Vector a)+freezeST GrowVec {..} = do   len <- VGM.unsafeRead posGV 0   vec <- readMutVar vecGV   VU.freeze $ VUM.take len vec --- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable--- vector may not be used after this operation.------ @since 1.0.0.0-{-# INLINE unsafeFreeze #-}-unsafeFreeze :: (PrimMonad m, VU.Unbox a) => GrowVec (PrimState m) a -> m (VU.Vector a)-unsafeFreeze GrowVec {..} = do+{-# INLINEABLE unsafeFreezeST #-}+unsafeFreezeST :: (VU.Unbox a) => GrowVec s a -> ST s (VU.Vector a)+unsafeFreezeST GrowVec {..} = do   len <- VGM.unsafeRead posGV 0   vec <- readMutVar vecGV   VU.unsafeFreeze $ VUM.take len vec
src/AtCoder/Internal/McfCsr.hs view
@@ -43,6 +43,24 @@     costCsr :: !(VU.Vector cost)   } +-- | \(O(n + m)\) Creates `Csr`.+--+-- @since 1.0.0.0+{-# INLINE build #-}+build :: (HasCallStack, PrimMonad m, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> m (VU.Vector Int, Csr (PrimState m) cap cost)+build n edges = stToPrim $ buildST n edges++-- | \(O(1)\) Returns a vector of @(to, rev, cost)@.+--+-- @since 1.0.0.0+{-# INLINE adj #-}+adj :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost) => Csr s cap cost -> Int -> VU.Vector (Int, Int, cost)+adj Csr {..} v = VU.slice offset len vec+  where+    offset = startCsr VG.! v+    len = startCsr VG.! (v + 1) - offset+    vec = VU.zip3 toCsr revCsr costCsr+ {-# INLINEABLE buildST #-} buildST :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> ST s (VU.Vector Int, Csr s cap cost) buildST n edges = do@@ -89,21 +107,3 @@   revCsr <- VU.unsafeFreeze revVec   costCsr <- VU.unsafeFreeze costVec   pure (edgeIdx, Csr {..})---- | \(O(n + m)\) Creates `Csr`.------ @since 1.0.0.0-{-# INLINE build #-}-build :: (HasCallStack, PrimMonad m, Num cap, VU.Unbox cap, VU.Unbox cost, Num cost) => Int -> VU.Vector (Int, Int, cap, cap, cost) -> m (VU.Vector Int, Csr (PrimState m) cap cost)-build n edges = stToPrim $ buildST n edges---- | \(O(1)\) Returns a vector of @(to, rev, cost)@.------ @since 1.0.0.0-{-# INLINE adj #-}-adj :: (HasCallStack, Num cap, VU.Unbox cap, VU.Unbox cost) => Csr s cap cost -> Int -> VU.Vector (Int, Int, cost)-adj Csr {..} v = VU.slice offset len vec-  where-    offset = startCsr VG.! v-    len = startCsr VG.! (v + 1) - offset-    vec = VU.zip3 toCsr revCsr costCsr
src/AtCoder/Internal/MinHeap.hs view
@@ -120,6 +120,45 @@ clear :: (PrimMonad m, VU.Unbox a) => Heap (PrimState m) a -> m () clear Heap {sizeBH_} = VGM.unsafeWrite sizeBH_ 0 0 +-- | \(O(\log n)\) Inserts an element to the heap.+--+-- @since 1.0.0.0+{-# INLINE push #-}+push :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> a -> m ()+push heap x = stToPrim $ pushST heap x++-- | \(O(\log n)\) Removes the last element from the heap and returns it, or `Nothing` if it is+-- empty.+--+-- @since 1.0.0.0+{-# INLINE pop #-}+pop :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> m (Maybe a)+pop heap = stToPrim $ popST heap++-- | \(O(\log n)\) `pop` with the return value discarded.+--+-- @since 1.0.0.0+{-# INLINE pop_ #-}+pop_ :: (HasCallStack, Ord a, VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m ()+pop_ heap = do+  _ <- stToPrim $ popST heap+  pure ()++-- | \(O(1)\) Returns the smallest value in the heap, or `Nothing` if it is empty.+--+-- @since 1.0.0.0+{-# INLINE peek #-}+peek :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m (Maybe a)+peek heap = do+  isNull <- null heap+  if isNull+    then pure Nothing+    else Just <$> VGM.read (dataBH heap) 0++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------+ {-# INLINEABLE pushST #-} pushST :: (HasCallStack, Ord a, VU.Unbox a) => Heap s a -> a -> ST s () pushST Heap {..} x = do@@ -134,13 +173,6 @@           siftUp iParent   siftUp i0 --- | \(O(\log n)\) Inserts an element to the heap.------ @since 1.0.0.0-{-# INLINE push #-}-push :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> a -> m ()-push heap x = stToPrim $ pushST heap x- {-# INLINEABLE popST #-} popST :: (HasCallStack, Ord a, VU.Unbox a) => Heap s a -> ST s (Maybe a) popST heap@Heap {..} = do@@ -178,31 +210,3 @@        siftDown 0       pure $ Just root---- | \(O(\log n)\) Removes the last element from the heap and returns it, or `Nothing` if it is--- empty.------ @since 1.0.0.0-{-# INLINE pop #-}-pop :: (HasCallStack, PrimMonad m, Ord a, VU.Unbox a) => Heap (PrimState m) a -> m (Maybe a)-pop heap = stToPrim $ popST heap---- | \(O(\log n)\) `pop` with the return value discarded.------ @since 1.0.0.0-{-# INLINE pop_ #-}-pop_ :: (HasCallStack, Ord a, VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m ()-pop_ heap = do-  _ <- stToPrim $ popST heap-  pure ()---- | \(O(1)\) Returns the smallest value in the heap, or `Nothing` if it is empty.------ @since 1.0.0.0-{-# INLINE peek #-}-peek :: (VU.Unbox a, PrimMonad m) => Heap (PrimState m) a -> m (Maybe a)-peek heap = do-  isNull <- null heap-  if isNull-    then pure Nothing-    else Just <$> VGM.read (dataBH heap) 0
src/AtCoder/Internal/Queue.hs view
@@ -65,7 +65,8 @@   ) where -import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Vector.Generic.Mutable qualified as VGM import Data.Vector.Unboxed qualified as VU import Data.Vector.Unboxed.Mutable qualified as VUM@@ -86,10 +87,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (PrimMonad m, VU.Unbox a) => Int -> m (Queue (PrimState m) a)-new n = do-  posQ <- VUM.replicate 2 (0 :: Int)-  vecQ <- VUM.unsafeNew n-  pure Queue {..}+new n = stToPrim $ newST n  -- | \(O(1)\) Returns the array size. --@@ -103,10 +101,7 @@ -- @since 1.0.0.0 {-# INLINE length #-} length :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m Int-length Queue {..} = do-  l <- VGM.unsafeRead posQ 0-  r <- VGM.unsafeRead posQ 1-  pure $ r - l+length que = stToPrim $ lengthST que  -- | \(O(1)\) Returns `True` if the buffer is empty. --@@ -120,56 +115,28 @@ -- @since 1.0.0.0 {-# INLINE pushBack #-} pushBack :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> a -> m ()-pushBack Queue {..} e = do-  VGM.unsafeModifyM-    posQ-    ( \r -> do-        VGM.write vecQ r e-        pure $ r + 1-    )-    1+pushBack que e = stToPrim $ pushBackST que e  -- | \(O(1)\) Appends an element to the back. Will throw an exception if the index is out of range. -- -- @since 1.0.0.0 {-# INLINE pushFront #-} pushFront :: (HasCallStack, PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> a -> m ()-pushFront Queue {..} e = do-  l0 <- VGM.unsafeRead posQ 0-  if l0 == 0-    then error "AtCoder.Internal.Queue.pushFront: no empty front space"-    else do-      VGM.unsafeModifyM-        posQ-        ( \l -> do-            VGM.write vecQ (l - 1) e-            pure $ l - 1-        )-        0+pushFront que e = stToPrim $ pushFrontST que e  -- | \(O(1)\) Removes the first element from the queue and returns it, or `Nothing` if it is empty. -- -- @since 1.0.0.0 {-# INLINE popFront #-} popFront :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (Maybe a)-popFront Queue {..} = do-  l <- VGM.unsafeRead posQ 0-  r <- VGM.unsafeRead posQ 1-  if l >= r-    then pure Nothing-    else do-      x <- VGM.read vecQ l-      VGM.unsafeWrite posQ 0 (l + 1)-      pure $ Just x+popFront que = stToPrim $ popFrontST que  -- | \(O(1)\) `popFront` with the return value discarded. -- -- @since 1.0.0.0 {-# INLINE popFront_ #-} popFront_ :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m ()-popFront_ que = do-  _ <- popFront que-  pure ()+popFront_ que = stToPrim $ popFrontST_ que  -- | \(O(1)\) Sets the `length` to zero. --@@ -184,10 +151,7 @@ -- @since 1.0.0.0 {-# INLINE freeze #-} freeze :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (VU.Vector a)-freeze Queue {..} = do-  l <- VGM.unsafeRead posQ 0-  r <- VGM.unsafeRead posQ 1-  VU.freeze $ VUM.take (r - l) $ VUM.drop l vecQ+freeze que = stToPrim $ freezeST que  -- | \(O(1)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable -- vector may not be used after this operation.@@ -195,7 +159,85 @@ -- @since 1.0.0.0 {-# INLINE unsafeFreeze #-} unsafeFreeze :: (PrimMonad m, VU.Unbox a) => Queue (PrimState m) a -> m (VU.Vector a)-unsafeFreeze Queue {..} = do+unsafeFreeze que = stToPrim $ unsafeFreezeST que++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (VU.Unbox a) => Int -> ST s (Queue s a)+newST n = do+  posQ <- VUM.replicate 2 (0 :: Int)+  vecQ <- VUM.unsafeNew n+  pure Queue {..}++{-# INLINEABLE lengthST #-}+lengthST :: (VU.Unbox a) => Queue s a -> ST s Int+lengthST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  pure $ r - l++{-# INLINEABLE pushBackST #-}+pushBackST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s ()+pushBackST Queue {..} e = do+  VGM.unsafeModifyM+    posQ+    ( \r -> do+        VGM.write vecQ r e+        pure $ r + 1+    )+    1++{-# INLINEABLE pushFrontST #-}+pushFrontST :: (HasCallStack, VU.Unbox a) => Queue s a -> a -> ST s ()+pushFrontST Queue {..} e = do+  l0 <- VGM.unsafeRead posQ 0+  if l0 == 0+    then error "AtCoder.Internal.Queue.pushFront: no empty front space"+    else do+      VGM.unsafeModifyM+        posQ+        ( \l -> do+            VGM.write vecQ (l - 1) e+            pure $ l - 1+        )+        0++{-# INLINEABLE popFrontST #-}+popFrontST :: (VU.Unbox a) => Queue s a -> ST s (Maybe a)+popFrontST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  if l >= r+    then pure Nothing+    else do+      x <- VGM.read vecQ l+      VGM.unsafeWrite posQ 0 (l + 1)+      pure $ Just x++{-# INLINEABLE popFrontST_ #-}+popFrontST_ :: (VU.Unbox a) => Queue s a -> ST s ()+popFrontST_ que = do+  _ <- popFront que+  pure ()++{-# INLINEABLE clearST #-}+clearST :: (VU.Unbox a) => Queue s a -> ST s ()+clearST Queue {..} = do+  VGM.set posQ 0++{-# INLINEABLE freezeST #-}+freezeST :: (VU.Unbox a) => Queue s a -> ST s (VU.Vector a)+freezeST Queue {..} = do+  l <- VGM.unsafeRead posQ 0+  r <- VGM.unsafeRead posQ 1+  VU.freeze $ VUM.take (r - l) $ VUM.drop l vecQ++{-# INLINEABLE unsafeFreezeST #-}+unsafeFreezeST :: (VU.Unbox a) => Queue s a -> ST s (VU.Vector a)+unsafeFreezeST Queue {..} = do   l <- VGM.unsafeRead posQ 0   r <- VGM.unsafeRead posQ 1   VU.unsafeFreeze $ VUM.take (r - l) $ VUM.drop l vecQ
src/AtCoder/Internal/Scc.hs view
@@ -75,7 +75,18 @@   csr <- ACICSR.build' nScc <$> ACIGV.unsafeFreeze edgesScc   pure $ sccIdsCsr csr --- NOTE(perf): faster without INLINEABLE (somehow)+-- | \(O(n + m)\) Returns the strongly connected components.+--+-- @since 1.0.0.0+{-# INLINE scc #-}+scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int))+scc g = stToPrim $ sccST g++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE sccST #-} sccST :: SccGraph s -> ST s (V.Vector (VU.Vector Int)) sccST g = do   (!groupNum, !ids) <- sccIds g@@ -91,13 +102,6 @@     VGM.write is sccId $ i + 1     VGM.write (groups VG.! sccId) i v   V.mapM VU.unsafeFreeze groups---- | \(O(n + m)\) Returns the strongly connected components.------ @since 1.0.0.0-{-# INLINE scc #-}-scc :: (PrimMonad m) => SccGraph (PrimState m) -> m (V.Vector (VU.Vector Int))-scc g = stToPrim $ sccST g  -- | \(O(n + m)\) API) Returns a pair of @(# of scc, scc id)@. --
src/AtCoder/Internal/String.hs view
@@ -25,7 +25,7 @@ -- | \(O(n^2)\) Internal implementation of suffix array creation (naive). -- -- @since 1.0.0.0-{-# INLINABLE saNaive #-}+{-# INLINEABLE saNaive #-} saNaive :: (HasCallStack) => VU.Vector Int -> VU.Vector Int saNaive s =   let n = VU.length s@@ -48,7 +48,7 @@ -- | \(O(n \log n)\) Internal implementation of suffix array creation (doubling). -- -- @since 1.0.0.0-{-# INLINABLE saDoubling #-}+{-# INLINEABLE saDoubling #-} saDoubling :: (HasCallStack) => VU.Vector Int -> VU.Vector Int saDoubling s = VU.create $ do   let n = VU.length s@@ -87,7 +87,7 @@ -- | \(O(n)\) Internal implementation of suffix array creation (suffix array induced sorting). -- -- @since 1.0.0.0-{-# INLINABLE saIsImpl #-}+{-# INLINEABLE saIsImpl #-} saIsImpl ::   (HasCallStack) =>   -- | naive threshould
src/AtCoder/LazySegTree.hs view
@@ -321,21 +321,6 @@     lzLst :: !(VUM.MVector s f)   } -{-# INLINE buildST #-}-buildST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => VU.Vector a -> ST s (LazySegTree s f a)-buildST vs = do-  let nLst = VU.length vs-  let sizeLst = ACIBIT.bitCeil nLst-  let logLst = countTrailingZeros sizeLst-  dLst <- VUM.replicate (2 * sizeLst) mempty-  lzLst <- VUM.replicate sizeLst mempty-  VU.iforM_ vs $ \i v -> do-    VGM.write dLst (sizeLst + i) v-  let segtree = LazySegTree {..}-  for_ [sizeLst - 1, sizeLst - 2 .. 1] $ \i -> do-    updateST segtree i-  pure segtree- -- | Creates an array of length \(n\). All the elements are initialized to `mempty`. -- -- ==== Constraints@@ -364,16 +349,6 @@ build :: (PrimMonad m, Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => VU.Vector a -> m (LazySegTree (PrimState m) f a) build vs = stToPrim $ buildST vs -{-# INLINE writeST #-}-writeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s ()-writeST self@LazySegTree {..} p x = do-  let p' = p + sizeLst-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    pushST self $ p' .>>. i-  VGM.unsafeWrite dLst p' x-  for_ [1 .. logLst] $ \i -> do-    updateST self $ p' .>>. i- -- | Sets \(p\)-th value of the array to \(x\). -- -- ==== Constraints@@ -389,16 +364,6 @@   where     !_ = ACIA.checkIndex "AtCoder.LazySegTree.write" p (nLst self) -{-# INLINE modifyST #-}-modifyST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> (a -> a) -> Int -> ST s ()-modifyST self@LazySegTree {..} f p = do-  let p' = p + sizeLst-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    pushST self $ p' .>>. i-  VGM.unsafeModify dLst f p'-  for_ [1 .. logLst] $ \i -> do-    updateST self $ p' .>>. i- -- | (Extra API) Modifies \(p\)-th value with a function \(f\). -- -- ==== Constraints@@ -423,7 +388,7 @@ -- - \(O(\log n)\) -- -- @since 1.0.0.0-{-# INLINE modifyM #-}+{-# INLINEABLE modifyM #-} modifyM :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> (a -> m a) -> Int -> m () modifyM self@LazySegTree {..} f p = do   let !_ = ACIA.checkIndex "AtCoder.LazySegTree.modifyM" p nLst@@ -434,17 +399,6 @@   stToPrim $ for_ [1 .. logLst] $ \i -> do     updateST self $ p' .>>. i -{-# INLINE exchangeST #-}-exchangeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s a-exchangeST self@LazySegTree {..} p x = do-  let p' = p + sizeLst-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    pushST self $ p' .>>. i-  res <- VGM.unsafeExchange dLst p' x-  for_ [1 .. logLst] $ \i -> do-    updateST self $ p' .>>. i-  pure res- -- | (Extra API) Sets \(p\)-th value of the array to \(x\) and returns the old value. -- -- ==== Constraints@@ -460,14 +414,6 @@   where     !_ = ACIA.checkIndex "AtCoder.LazySegTree.exchange" p (nLst self) -{-# INLINE readST #-}-readST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> ST s a-readST self@LazySegTree {..} p = do-  let p' = p + sizeLst-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    pushST self $ p' .>>. i-  VGM.unsafeRead dLst p'- -- | Returns \(p\)-th value of the array. -- -- ==== Constraints@@ -515,31 +461,6 @@   | l0 == r0 = pure (Just mempty)   | otherwise = stToPrim $ Just <$> unsafeProdST self l0 r0 --- | Internal implementation of `prod`.-{-# INLINE unsafeProdST #-}-unsafeProdST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> ST s a-unsafeProdST self@LazySegTree {..} l0 r0 = do-  let l = l0 + sizeLst-  let r = r0 + sizeLst-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    when (((l .>>. i) .<<. i) /= l) $ pushST self $ l .>>. i-    when (((r .>>. i) .<<. i) /= r) $ pushST self $ (r - 1) .>>. i-  inner l (r - 1) mempty mempty-  where-    -- NOTE: we're using inclusive range [l, r] for simplicity-    inner l r !smL !smR-      | l > r = pure $! smL <> smR-      | otherwise = do-          !smL' <--            if testBit l 0-              then (smL <>) <$> VGM.read dLst l-              else pure smL-          !smR' <--            if not $ testBit r 0-              then (<> smR) <$> VGM.read dLst r-              else pure smR-          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) smL' smR'- -- | Returns the product of \([op(a[0], ..., a[n - 1])]\), assuming the properties of the monoid. It -- returns `mempty` if \(n = 0\). --@@ -551,20 +472,6 @@ allProd :: (PrimMonad m, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m a allProd LazySegTree {..} = VGM.read dLst 1 -{-# INLINE applyAtST #-}-applyAtST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> f -> ST s ()-applyAtST self@LazySegTree {..} p f = do-  let p' = p + sizeLst-  -- propagate-  for_ [logLst, logLst - 1 .. 1] $ \i -> do-    pushST self $ p' .>>. i-  -- FIXME: length should be always `1`-  let !len = bit $! logLst - (63 - countLeadingZeros p')-  VGM.modify dLst (segActWithLength len f) p'-  -- evaluate-  for_ [1 .. logLst] $ \i -> do-    updateST self $ p' .>>. i- -- | Applies @segAct f@ to an index \(p\). -- -- ==== Constraints@@ -580,33 +487,6 @@   where     !_ = ACIA.checkIndex "AtCoder.LazySegTree.applyAt" p (nLst self) -{-# INLINE applyInST #-}-applyInST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> f -> ST s ()-applyInST self@LazySegTree {..} l0 r0 f-  | l0 == r0 = pure ()-  | otherwise = do-      let l = l0 + sizeLst-      let r = r0 + sizeLst-      -- propagate-      for_ [logLst, logLst - 1 .. 1] $ \i -> do-        when (((l .>>. i) .<<. i) /= l) $ pushST self (l .>>. i)-        when (((r .>>. i) .<<. i) /= r) $ pushST self ((r - 1) .>>. i)-      inner l (r - 1)-      -- evaluate-      for_ [1 .. logLst] $ \i -> do-        when (((l .>>. i) .<<. i) /= l) $ updateST self (l .>>. i)-        when (((r .>>. i) .<<. i) /= r) $ updateST self ((r - 1) .>>. i)-  where-    -- NOTE: we're using inclusive range [l, r] for simplicity-    inner l r-      | l > r = pure ()-      | otherwise = do-          when (testBit l 0) $ do-            allApplyST self l f-          unless (testBit r 0) $ do-            allApplyST self r f-          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1)- -- | Applies @segAct f@ to an interval \([l, r)\). -- -- ==== Constraints@@ -657,7 +537,7 @@ -- - \(O(\log n)\) -- -- @since 1.0.0.0-{-# INLINE minLeftM #-}+{-# INLINEABLE minLeftM #-} minLeftM :: (HasCallStack, PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> Int -> (a -> m Bool) -> m Int minLeftM self@LazySegTree {..} r0 g = do   b <- g mempty@@ -781,11 +661,7 @@ -- @since 1.0.0.0 {-# INLINE freeze #-} freeze :: (PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m (VU.Vector a)-freeze self@LazySegTree {..} = do-  -- push all (we _could_ skip some elements)-  stToPrim $ for_ [1 .. sizeLst - 1] $ \i -> do-    pushST self i-  VU.freeze . VUM.take nLst $ VUM.drop sizeLst dLst+freeze = stToPrim . freezeST  -- | \(O(n)\) Unsafely converts a mutable vector to an immutable one without copying. The mutable -- vector may not be used after this operation.@@ -793,13 +669,148 @@ -- @since 1.0.0.0 {-# INLINE unsafeFreeze #-} unsafeFreeze :: (PrimMonad m, SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree (PrimState m) f a -> m (VU.Vector a)-unsafeFreeze self@LazySegTree {..} = do+unsafeFreeze = stToPrim . unsafeFreezeST++-- NOTE (perf): these functions have to be inlined after all:++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE buildST #-}+buildST :: (Monoid f, VU.Unbox f, Monoid a, VU.Unbox a) => VU.Vector a -> ST s (LazySegTree s f a)+buildST vs = do+  let nLst = VU.length vs+  let sizeLst = ACIBIT.bitCeil nLst+  let logLst = countTrailingZeros sizeLst+  dLst <- VUM.replicate (2 * sizeLst) mempty+  lzLst <- VUM.replicate sizeLst mempty+  VU.iforM_ vs $ \i v -> do+    VGM.write dLst (sizeLst + i) v+  let segtree = LazySegTree {..}+  for_ [sizeLst - 1, sizeLst - 2 .. 1] $ \i -> do+    updateST segtree i+  pure segtree++{-# INLINEABLE writeST #-}+writeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s ()+writeST self@LazySegTree {..} p x = do+  let p' = p + sizeLst+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    pushST self $ p' .>>. i+  VGM.unsafeWrite dLst p' x+  for_ [1 .. logLst] $ \i -> do+    updateST self $ p' .>>. i++{-# INLINEABLE modifyST #-}+modifyST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> (a -> a) -> Int -> ST s ()+modifyST self@LazySegTree {..} f p = do+  let p' = p + sizeLst+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    pushST self $ p' .>>. i+  VGM.unsafeModify dLst f p'+  for_ [1 .. logLst] $ \i -> do+    updateST self $ p' .>>. i++{-# INLINEABLE exchangeST #-}+exchangeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> a -> ST s a+exchangeST self@LazySegTree {..} p x = do+  let p' = p + sizeLst+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    pushST self $ p' .>>. i+  res <- VGM.unsafeExchange dLst p' x+  for_ [1 .. logLst] $ \i -> do+    updateST self $ p' .>>. i+  pure res++{-# INLINEABLE readST #-}+readST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> ST s a+readST self@LazySegTree {..} p = do+  let p' = p + sizeLst+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    pushST self $ p' .>>. i+  VGM.unsafeRead dLst p'++{-# INLINEABLE applyAtST #-}+applyAtST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> f -> ST s ()+applyAtST self@LazySegTree {..} p f = do+  let p' = p + sizeLst+  -- propagate+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    pushST self $ p' .>>. i+  -- FIXME: length should be always `1`+  let !len = bit $! logLst - (63 - countLeadingZeros p')+  VGM.modify dLst (segActWithLength len f) p'+  -- evaluate+  for_ [1 .. logLst] $ \i -> do+    updateST self $ p' .>>. i++{-# INLINEABLE applyInST #-}+applyInST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> f -> ST s ()+applyInST self@LazySegTree {..} l0 r0 f+  | l0 == r0 = pure ()+  | otherwise = do+      let l = l0 + sizeLst+      let r = r0 + sizeLst+      -- propagate+      for_ [logLst, logLst - 1 .. 1] $ \i -> do+        when (((l .>>. i) .<<. i) /= l) $ pushST self (l .>>. i)+        when (((r .>>. i) .<<. i) /= r) $ pushST self ((r - 1) .>>. i)+      inner l (r - 1)+      -- evaluate+      for_ [1 .. logLst] $ \i -> do+        when (((l .>>. i) .<<. i) /= l) $ updateST self (l .>>. i)+        when (((r .>>. i) .<<. i) /= r) $ updateST self ((r - 1) .>>. i)+  where+    -- NOTE: we're using inclusive range [l, r] for simplicity+    inner l r+      | l > r = pure ()+      | otherwise = do+          when (testBit l 0) $ do+            allApplyST self l f+          unless (testBit r 0) $ do+            allApplyST self r f+          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1)++-- | Internal implementation of `prod`.+{-# INLINEABLE unsafeProdST #-}+unsafeProdST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> Int -> Int -> ST s a+unsafeProdST self@LazySegTree {..} l0 r0 = do+  let l = l0 + sizeLst+  let r = r0 + sizeLst+  for_ [logLst, logLst - 1 .. 1] $ \i -> do+    when (((l .>>. i) .<<. i) /= l) $ pushST self $ l .>>. i+    when (((r .>>. i) .<<. i) /= r) $ pushST self $ (r - 1) .>>. i+  inner l (r - 1) mempty mempty+  where+    -- NOTE: we're using inclusive range [l, r] for simplicity+    inner l r !smL !smR+      | l > r = pure $! smL <> smR+      | otherwise = do+          !smL' <-+            if testBit l 0+              then (smL <>) <$> VGM.read dLst l+              else pure smL+          !smR' <-+            if not $ testBit r 0+              then (<> smR) <$> VGM.read dLst r+              else pure smR+          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) smL' smR'++{-# INLINEABLE freezeST #-}+freezeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> ST s (VU.Vector a)+freezeST self@LazySegTree {..} = do+  for_ [1 .. sizeLst - 1] $ \i -> do+    pushST self i+  VU.freeze . VUM.take nLst $ VUM.drop sizeLst dLst++{-# INLINEABLE unsafeFreezeST #-}+unsafeFreezeST :: (SegAct f a, VU.Unbox f, Monoid a, VU.Unbox a) => LazySegTree s f a -> ST s (VU.Vector a)+unsafeFreezeST self@LazySegTree {..} = do   -- push all (we _could_ skip some elements)-  stToPrim $ for_ [1 .. sizeLst - 1] $ \i -> do+  for_ [1 .. sizeLst - 1] $ \i -> do     pushST self i   VU.unsafeFreeze . VUM.take nLst $ VUM.drop sizeLst dLst---- NOTE (perf): these functions have to be inlined after all:  -- | \(O(1)\) {-# INLINE updateST #-}
src/AtCoder/MaxFlow.hs view
@@ -100,10 +100,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (PrimMonad m, VU.Unbox cap) => Int -> m (MfGraph (PrimState m) cap)-new nG = do-  gG <- V.replicateM nG (ACIGV.new 0)-  posG <- ACIGV.new 0-  pure MfGraph {..}+new nG = stToPrim $ newST nG  -- | Adds an edge oriented from the vertex @from@ to the vertex @to@ with the capacity @cap@ and the -- flow amount \(0\). It returns an integer \(k\) such that this is the \(k\)-th edge that is added.@@ -129,19 +126,7 @@   cap ->   -- | Edge index   m Int-addEdge MfGraph {..} from to cap = do-  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdge" "`from` vertex" from "the number of vertices" nG-  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdge" "`to` vertex" to "the number of vertices" nG-  let !_ = ACIA.runtimeAssert (0 <= cap) "AtCoder.MaxFlow.addEdge: given invalid edge `cap` less than `0`" -- not `Show cap`-  m <- ACIGV.length posG-  iEdge <- ACIGV.length (gG VG.! from)-  ACIGV.pushBack posG (from, iEdge)-  iRevEdge <- do-    len <- ACIGV.length (gG VG.! to)-    pure $ if from == to then len + 1 else len-  ACIGV.pushBack (gG VG.! from) (to, iRevEdge, cap)-  ACIGV.pushBack (gG VG.! to) (from, iEdge, 0)-  pure m+addEdge g from to cap = stToPrim $ addEdgeST g from to cap  -- | `addEdge` with the return value discarded. --@@ -165,8 +150,8 @@   -- | cap   cap ->   m ()-addEdge_ graph from to cap = do-  _ <- addEdge graph from to cap+addEdge_ graph from to cap = stToPrim $ do+  _ <- addEdgeST graph from to cap   pure ()  -- | Augments the flow from \(s\) to \(t\) as much as possible, until reaching the amount of@@ -195,11 +180,169 @@   cap ->   -- | Max flow   m cap-flow MfGraph {..} s t flowLimit = stToPrim $ do-  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flow" "`source` vertex" s "the number of vertices" nG-  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flow" "`sink` vertex" t "the number of vertices" nG-  let !_ = ACIA.runtimeAssert (s /= t) $ "AtCoder.MaxFlow.flow: `source` and `sink` vertex must be distinct: `" ++ show s ++ "`"+flow g s t flowLimit = stToPrim $ flowST g s t flowLimit +-- | `flow` with no capacity limit.+--+-- ==== Constraints+-- - \(s \neq t\)+-- - \(0 \leq s, t \lt n\)+--+-- ==== Complexity+-- - \(O((n + m) \sqrt{m})\) (if all the capacities are \(1\)),+-- - \(O(n^2 m)\) (general), or+-- - \(O(F(n + m))\), where \(F\) is the returned value+--+-- @since 1.0.0.0+{-# INLINE maxFlow #-}+maxFlow ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, Bounded cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Source @s@+  Int ->+  -- | Sink @t@+  Int ->+  -- | Max flow+  m cap+maxFlow graph s t = stToPrim $ flowST graph s t maxBound++-- | Returns a vector of length \(n\), such that the \(i\)-th element is `True` if and only if there+-- is a directed path from \(s\) to \(i\) in the residual network. The returned vector corresponds+-- to a \(s-t\) minimum cut after calling @'maxFlow' s t@.+--+-- ==== Complexity+-- - \(O(n + m)\), where \(m\) is the number of added edges.+--+-- @since 1.0.0.0+{-# INLINE minCut #-}+minCut ::+  (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Source @s@+  Int ->+  -- | Minimum cut+  m (VU.Vector Bit)+minCut g s = stToPrim $ minCutST g s++-- | \(O(1)\) Returns the current internal state of \(i\)-th edge: @(from, to, cap, flow)@. The+-- edges are ordered in the same order as added by `addEdge`.+--+-- ==== Constraints+-- - \(0 \leq i \lt m\)+--+-- ==== Complexity+-- - \(O(1)\)+--+-- @since 1.0.0.0+{-# INLINE getEdge #-}+getEdge ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Vertex+  Int ->+  -- | Tuple of @(from, to, cap, flow)@+  m (Int, Int, cap, cap)+getEdge g i = stToPrim $ getEdgeST g i++-- | Returns the current internal state of the edges: @(from, to, cap, flow)@. The edges are ordered+-- in the same order as added by `addEdge`.+--+-- ==== Complexity+-- - \(O(m)\), where \(m\) is the number of added edges.+--+-- @since 1.0.0.0+{-# INLINE edges #-}+edges ::+  (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Vector of @(from, to, cap, flow)@+  m (VU.Vector (Int, Int, cap, cap))+edges g = stToPrim $ edgesST g++-- | \(O(1)\) Changes the capacity and the flow amount of the $i$-th edge to @newCap@ and+-- @newFlow@, respectively. It oes not change the capacity or the flow amount of other edges.+--+-- ==== Constraints+-- - \(0 \leq \mathrm{newflow} \leq \mathrm{newcap}\)+--+-- ==== Complexity+-- - \(O(1)\)+--+-- @since 1.0.0.0+{-# INLINE changeEdge #-}+changeEdge ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Edge index+  Int ->+  -- | New capacity+  cap ->+  -- | New flow+  cap ->+  m ()+changeEdge g i newCap newFlow = stToPrim $ changeEdgeST g i newCap newFlow++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (PrimMonad m, VU.Unbox cap) => Int -> m (MfGraph (PrimState m) cap)+newST nG = do+  gG <- V.replicateM nG (ACIGV.new 0)+  posG <- ACIGV.new 0+  pure MfGraph {..}++{-# INLINEABLE addEdgeST #-}+addEdgeST ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | from+  Int ->+  -- | to+  Int ->+  -- | cap+  cap ->+  -- | Edge index+  m Int+addEdgeST MfGraph {..} from to cap = do+  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdgeST" "`from` vertex" from "the number of vertices" nG+  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.addEdgeST" "`to` vertex" to "the number of vertices" nG+  let !_ = ACIA.runtimeAssert (0 <= cap) "AtCoder.MaxFlow.addEdgeST: given invalid edge `cap` less than `0`" -- not `Show cap`+  m <- ACIGV.length posG+  iEdge <- ACIGV.length (gG VG.! from)+  ACIGV.pushBack posG (from, iEdge)+  iRevEdge <- do+    len <- ACIGV.length (gG VG.! to)+    pure $ if from == to then len + 1 else len+  ACIGV.pushBack (gG VG.! from) (to, iRevEdge, cap)+  ACIGV.pushBack (gG VG.! to) (from, iEdge, 0)+  pure m++{-# INLINEABLE flowST #-}+flowST ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>+  -- | Graph+  MfGraph (PrimState m) cap ->+  -- | Source @s@+  Int ->+  -- | Sink @t@+  Int ->+  -- | Flow limit+  cap ->+  -- | Max flow+  m cap+flowST MfGraph {..} s t flowLimit = stToPrim $ do+  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flowST" "`source` vertex" s "the number of vertices" nG+  let !_ = ACIA.checkCustom "AtCoder.MaxFlow.flowST" "`sink` vertex" t "the number of vertices" nG+  let !_ = ACIA.runtimeAssert (s /= t) $ "AtCoder.MaxFlow.flowST: `source` and `sink` vertex must be distinct: `" ++ show s ++ "`"+   level <- VUM.unsafeNew nG   que <- ACIQ.new nG   let bfs = do@@ -273,41 +416,8 @@               then pure flow_               else loop $! flow_ + f --- | `flow` with no capacity limit.------ ==== Constraints--- - \(s \neq t\)--- - \(0 \leq s, t \lt n\)------ ==== Complexity--- - \(O((n + m) \sqrt{m})\) (if all the capacities are \(1\)),--- - \(O(n^2 m)\) (general), or--- - \(O(F(n + m))\), where \(F\) is the returned value------ @since 1.0.0.0-{-# INLINE maxFlow #-}-maxFlow ::-  (HasCallStack, PrimMonad m, Num cap, Ord cap, Bounded cap, VU.Unbox cap) =>-  -- | Graph-  MfGraph (PrimState m) cap ->-  -- | Source @s@-  Int ->-  -- | Sink @t@-  Int ->-  -- | Max flow-  m cap-maxFlow graph s t = flow graph s t maxBound---- | Returns a vector of length \(n\), such that the \(i\)-th element is `True` if and only if there--- is a directed path from \(s\) to \(i\) in the residual network. The returned vector corresponds--- to a \(s-t\) minimum cut after calling @'maxFlow' s t@.------ ==== Complexity--- - \(O(n + m)\), where \(m\) is the number of added edges.------ @since 1.0.0.0-{-# INLINE minCut #-}-minCut ::+{-# INLINEABLE minCutST #-}+minCutST ::   (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>   -- | Graph   MfGraph (PrimState m) cap ->@@ -315,7 +425,7 @@   Int ->   -- | Minimum cut   m (VU.Vector Bit)-minCut MfGraph {..} s = stToPrim $ do+minCutST MfGraph {..} s = stToPrim $ do   visited <- VUM.replicate nG $ Bit False   que <- ACIQ.new nG -- we could use a growable queue here   ACIQ.pushBack que s@@ -334,18 +444,8 @@         loop   VU.unsafeFreeze visited --- | \(O(1)\) Returns the current internal state of \(i\)-th edge: @(from, to, cap, flow)@. The--- edges are ordered in the same order as added by `addEdge`.------ ==== Constraints--- - \(0 \leq i \lt m\)------ ==== Complexity--- - \(O(1)\)------ @since 1.0.0.0-{-# INLINE getEdge #-}-getEdge ::+{-# INLINEABLE getEdgeST #-}+getEdgeST ::   (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>   -- | Graph   MfGraph (PrimState m) cap ->@@ -353,7 +453,7 @@   Int ->   -- | Tuple of @(from, to, cap, flow)@   m (Int, Int, cap, cap)-getEdge MfGraph {..} i = stToPrim $ do+getEdgeST MfGraph {..} i = stToPrim $ do   m <- ACIGV.length posG   let !_ = ACIA.checkEdge "AtCoder.MaxFlow.getEdge" i m   (!from, !iEdge) <- ACIGV.read posG i@@ -361,36 +461,19 @@   revCap <- readCapacityST gG to iRevEdge   pure (from, to, cap + revCap, revCap) --- | Returns the current internal state of the edges: @(from, to, cap, flow)@. The edges are ordered--- in the same order as added by `addEdge`.------ ==== Complexity--- - \(O(m)\), where \(m\) is the number of added edges.------ @since 1.0.0.0-{-# INLINE edges #-}-edges ::+{-# INLINEABLE edgesST #-}+edgesST ::   (PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>   -- | Graph   MfGraph (PrimState m) cap ->   -- | Vector of @(from, to, cap, flow)@   m (VU.Vector (Int, Int, cap, cap))-edges g@MfGraph {posG} = do+edgesST g@MfGraph {posG} = do   len <- ACIGV.length posG   VU.generateM len (getEdge g) --- | \(O(1)\) Changes the capacity and the flow amount of the $i$-th edge to @newCap@ and--- @newFlow@, respectively. It oes not change the capacity or the flow amount of other edges.------ ==== Constraints--- - \(0 \leq \mathrm{newflow} \leq \mathrm{newcap}\)------ ==== Complexity--- - \(O(1)\)------ @since 1.0.0.0-{-# INLINE changeEdge #-}-changeEdge ::+{-# INLINEABLE changeEdgeST #-}+changeEdgeST ::   (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap) =>   -- | Graph   MfGraph (PrimState m) cap ->@@ -401,30 +484,27 @@   -- | New flow   cap ->   m ()-changeEdge MfGraph {..} i newCap newFlow = stToPrim $ do+changeEdgeST MfGraph {..} i newCap newFlow = stToPrim $ do   m <- ACIGV.length posG-  let !_ = ACIA.checkEdge "AtCoder.MaxFlow.changeEdge" i m-  let !_ = ACIA.runtimeAssert (0 <= newFlow && newFlow <= newCap) "AtCoder.MaxFlow.changeEdge: invalid flow or capacity" -- not Show+  let !_ = ACIA.checkEdge "AtCoder.MaxFlow.changeEdgeST" i m+  let !_ = ACIA.runtimeAssert (0 <= newFlow && newFlow <= newCap) "AtCoder.MaxFlow.changeEdgeST: invalid flow or capacity" -- not Show   (!from, !iEdge) <- ACIGV.read posG i   (!to, !iRevEdge, !_) <- ACIGV.read (gG VG.! from) iEdge   writeCapacityST gG from iEdge $! newCap - newFlow   writeCapacityST gG to iRevEdge $! newFlow --- | \(O(1)\) Internal helper. {-# INLINE readCapacityST #-} readCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec s (Int, Int, cap)) -> Int -> Int -> ST s cap readCapacityST gvs v i = do   (VUM.MV_3 _ _ _ c) <- readMutVar $ ACIGV.vecGV $ gvs VG.! v   VGM.read c i --- | \(O(1)\) Internal helper. {-# INLINE writeCapacityST #-} writeCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => V.Vector (ACIGV.GrowVec s (Int, Int, cap)) -> Int -> Int -> cap -> ST s () writeCapacityST gvs v i cap = do   (VUM.MV_3 _ _ _ c) <- readMutVar $ ACIGV.vecGV $ gvs VG.! v   VGM.write c i cap --- | \(O(1)\) Internal helper. {-# INLINE modifyCapacityST #-} modifyCapacityST :: (Num cap, Ord cap, VU.Unbox cap) => ACIGV.GrowVec s (Int, Int, cap) -> (cap -> cap) -> Int -> ST s () modifyCapacityST gv f i = do
src/AtCoder/MinCostFlow.hs view
@@ -49,6 +49,7 @@  -- TODO: add `maxCostFlow`. -- TODO: add `build`.+-- TODO: is this fast enough with `INLINEABLE`?  import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Buffer qualified as ACIB@@ -57,7 +58,8 @@ import AtCoder.Internal.MinHeap qualified as ACIMH import Control.Monad (unless, when) import Control.Monad.Fix (fix)-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bit (Bit (..)) import Data.Maybe (fromJust) import Data.Primitive.MutVar (readMutVar)@@ -92,7 +94,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (PrimMonad m, VU.Unbox cap, VU.Unbox cost) => Int -> m (McfGraph (PrimState m) cap cost)-new nG = do+new nG = stToPrim $ do   edgesG <- ACIGV.new 0   pure McfGraph {..} @@ -122,14 +124,7 @@   cost ->   -- | Edge index   m Int-addEdge McfGraph {..} from to cap cost = do-  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdge" "`from` vertex" from "the number of vertices" nG-  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdge" "`to` vertex" to "the number of vertices" nG-  let !_ = ACIA.runtimeAssert (0 <= cap) "AtCoder.MinCostFlow.addEdge: given invalid edge `cap` less than `0`"-  let !_ = ACIA.runtimeAssert (0 <= cost) "AtCoder.MinCostFlow.addEdge: given invalid edge `cost` less than `0`"-  m <- ACIGV.length edgesG-  ACIGV.pushBack edgesG (from, to, cap, 0, cost)-  pure m+addEdge g from to cap cost = stToPrim $ addEdgeST g from to cap cost  -- | `addEdge` with the return value discarded. --@@ -155,8 +150,8 @@   -- | cost   cost ->   m ()-addEdge_ graph from to cap cost = do-  _ <- addEdge graph from to cap cost+addEdge_ graph from to cap cost = stToPrim $ do+  _ <- addEdgeST graph from to cap cost   pure ()  -- | Augments the flow from \(s\) to \(t\) as much as possible, until reaching the amount of@@ -182,8 +177,8 @@   cap ->   -- | Tuple of @(cap, cost@)   m (cap, cost)-flow graph s t flowLimit = do-  res <- slope graph s t flowLimit+flow graph s t flowLimit = stToPrim $ do+  res <- slopeST graph s t flowLimit   pure $ VG.last res  -- | `flow` with no capacity limit.@@ -206,8 +201,8 @@   Int ->   -- | Tuple of @(cap, cost@)   m (cap, cost)-maxFlow graph s t = do-  res <- slope graph s t maxBound+maxFlow graph s t = stToPrim $ do+  res <- slopeST graph s t maxBound   pure $ VG.last res  -- | Let \(g\) be a function such that \(g(x)\) is the cost of the minimum cost \(s-t\) flow when@@ -246,14 +241,112 @@   cap ->   -- | Vector of @(cap, cost)@   m (VU.Vector (cap, cost))-slope McfGraph {..} s t flowLimit = do-  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slope" "`source` vertex" s "the number of vertices" nG-  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slope" "`sink` vertex" t "the number of vertices" nG-  let !_ = ACIA.runtimeAssert (s /= t) "AtCoder.MinCostFlow.slope: `source` and `sink` vertex must be distict"+slope g s t flowLimit = stToPrim $ slopeST g s t flowLimit +-- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are+-- ordered in the same order as added by `addEdge`.+--+-- ==== Constraints+-- - \(0 \leq i \lt m\)+--+-- ==== Complexity+-- - \(O(1)\)+--+-- @since 1.0.0.0+{-# INLINE getEdge #-}+getEdge ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>+  -- | Graph+  McfGraph (PrimState m) cap cost ->+  -- | Edge index+  Int ->+  -- | Tuple of @(from, to, cap, flow, cost)@+  m (Int, Int, cap, cap, cost)+getEdge g i = stToPrim $ getEdgeST g i++-- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are+-- ordered in the same order as added by `addEdge`.+--+-- ==== Complexity+-- - \(O(m)\), where \(m\) is the number of added edges.+--+-- @since 1.0.0.0+{-# INLINE edges #-}+edges ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>+  -- | Graph+  McfGraph (PrimState m) cap cost ->+  -- | Vector of @(from, to, cap, flow, cost)@+  m (VU.Vector (Int, Int, cap, cap, cost))+edges McfGraph {..} = stToPrim $ do+  ACIGV.freeze edgesG++-- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@, but without+-- making copy. The edges are ordered in the same order as added by `addEdge`.+--+-- ==== Complexity+-- - \(O(1)\)+--+-- @since 1.0.0.0+{-# INLINE unsafeEdges #-}+unsafeEdges ::+  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>+  -- | Graph+  McfGraph (PrimState m) cap cost ->+  -- | Vector of @(from, to, cap, flow, cost)@+  m (VU.Vector (Int, Int, cap, cap, cost))+unsafeEdges McfGraph {..} = stToPrim $ do+  ACIGV.unsafeFreeze edgesG++-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE addEdgeST #-}+addEdgeST ::+  (HasCallStack, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>+  -- | Graph+  McfGraph s cap cost ->+  -- | from+  Int ->+  -- | to+  Int ->+  -- | capacity+  cap ->+  -- | cost+  cost ->+  -- | Edge index+  ST s Int+addEdgeST McfGraph {..} from to cap cost = do+  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdgeST" "`from` vertex" from "the number of vertices" nG+  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.addEdgeST" "`to` vertex" to "the number of vertices" nG+  let !_ = ACIA.runtimeAssert (0 <= cap) "AtCoder.MinCostFlow.addEdgeST: given invalid edge `cap` less than `0`"+  let !_ = ACIA.runtimeAssert (0 <= cost) "AtCoder.MinCostFlow.addEdgeST: given invalid edge `cost` less than `0`"+  m <- ACIGV.length edgesG+  ACIGV.pushBack edgesG (from, to, cap, 0, cost)+  pure m++{-# INLINEABLE slopeST #-}+slopeST ::+  (HasCallStack, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>+  -- | Graph+  McfGraph s cap cost ->+  -- | Source @s@+  Int ->+  -- | Sink @t@+  Int ->+  -- | Flow limit+  cap ->+  -- | Vector of @(cap, cost)@+  ST s (VU.Vector (cap, cost))+slopeST McfGraph {..} s t flowLimit = do+  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slopeST" "`source` vertex" s "the number of vertices" nG+  let !_ = ACIA.checkCustom "AtCoder.MinCostFlow.slopeST" "`sink` vertex" t "the number of vertices" nG+  let !_ = ACIA.runtimeAssert (s /= t) "AtCoder.MinCostFlow.slopeST: `source` and `sink` vertex must be distict"+   edges_@(VU.V_5 _ _ _ caps _ _) <- ACIGV.unsafeFreeze edgesG   (!edgeIdx, !g) <- ACIMCSR.build nG edges_-  result <- internalSlopeMCF g nG s t flowLimit+  result <- internalSlopeST g nG s t flowLimit    (VUM.MV_5 _ _ _ _ flows _) <- readMutVar $ ACIGV.vecGV edgesG   VU.iforM_ (VU.zip caps edgeIdx) $ \v (!cap1, !iEdge) -> do@@ -262,26 +355,26 @@    pure result -{-# INLINE internalSlopeMCF #-}-internalSlopeMCF ::-  forall cap cost m.-  (HasCallStack, PrimMonad m, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>-  ACIMCSR.Csr (PrimState m) cap cost ->+{-# INLINEABLE internalSlopeST #-}+internalSlopeST ::+  forall cap cost s.+  (HasCallStack, Integral cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, Bounded cost, VU.Unbox cost) =>+  ACIMCSR.Csr s cap cost ->   Int ->   Int ->   Int ->   cap ->-  m (VU.Vector (cap, cost))-internalSlopeMCF csr@ACIMCSR.Csr {..} n s t flowLimit = do+  ST s (VU.Vector (cap, cost))+internalSlopeST csr@ACIMCSR.Csr {..} n s t flowLimit = do   duals <- VUM.replicate n 0-  dists <- VUM.unsafeNew n :: m (VUM.MVector (PrimState m) cost)-  prevE <- VUM.unsafeNew n :: m (VUM.MVector (PrimState m) Int)-  vis <- VUM.unsafeNew n :: m (VUM.MVector (PrimState m) Bit)+  dists <- VUM.unsafeNew n :: ST s (VUM.MVector s cost)+  prevE <- VUM.unsafeNew n :: ST s (VUM.MVector s Int)+  vis <- VUM.unsafeNew n :: ST s (VUM.MVector s Bit)    -- FIXME: maximum capacity of heap?   let nEdges = VU.length toCsr-  queMin <- ACIB.new nEdges :: m (ACIB.Buffer (PrimState m) Int)-  heap <- ACIMH.new nEdges :: m (ACIMH.Heap (PrimState m) (cost, Int))+  queMin <- ACIB.new nEdges :: ST s (ACIB.Buffer s Int)+  heap <- ACIMH.new nEdges :: ST s (ACIMH.Heap s (cost, Int))    let dualRef = do         VGM.set dists $ maxBound @cost@@ -345,14 +438,14 @@   result <- ACIGV.new 16   ACIGV.pushBack result (0 :: cap, 0 :: cost) -  let inner :: cap -> cost -> cost -> m ()+  let inner :: cap -> cost -> cost -> ST s ()       inner flow_ cost prevCostPerFlow =         when (flow_ < flowLimit) $ do           b <- dualRef           when b $ do             prevE' <- VU.unsafeFreeze prevE -            let minC :: cap -> Int -> m cap+            let minC :: cap -> Int -> ST s cap                 minC !acc v                   | v == s = pure acc                   | otherwise = do@@ -361,7 +454,7 @@                       minC (min acc cap) $ toCsr VG.! iPrev             c <- minC (flowLimit - flow_) t -            let subC :: Int -> m ()+            let subC :: Int -> ST s ()                 subC v = when (v /= s) $ do                   let iPrev = prevE' VG.! v                   VGM.modify capCsr (+ c) iPrev@@ -380,60 +473,16 @@   inner 0 0 (-1)   ACIGV.unsafeFreeze result --- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are--- ordered in the same order as added by `addEdge`.------ ==== Constraints--- - \(0 \leq i \lt m\)------ ==== Complexity--- - \(O(1)\)------ @since 1.0.0.0-{-# INLINE getEdge #-}-getEdge ::-  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>+{-# INLINEABLE getEdgeST #-}+getEdgeST ::+  (HasCallStack, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>   -- | Graph-  McfGraph (PrimState m) cap cost ->+  McfGraph s cap cost ->   -- | Edge index   Int ->   -- | Tuple of @(from, to, cap, flow, cost)@-  m (Int, Int, cap, cap, cost)-getEdge McfGraph {..} i = do+  ST s (Int, Int, cap, cap, cost)+getEdgeST McfGraph {..} i = do   m <- ACIGV.length edgesG-  let !_ = ACIA.checkEdge "AtCoder.MinCostFlow.getEdge" i m+  let !_ = ACIA.checkEdge "AtCoder.MinCostFlow.getEdgeST" i m   ACIGV.read edgesG i---- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@. The edges are--- ordered in the same order as added by `addEdge`.------ ==== Complexity--- - \(O(m)\), where \(m\) is the number of added edges.------ @since 1.0.0.0-{-# INLINE edges #-}-edges ::-  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>-  -- | Graph-  McfGraph (PrimState m) cap cost ->-  -- | Vector of @(from, to, cap, flow, cost)@-  m (VU.Vector (Int, Int, cap, cap, cost))-edges McfGraph {..} = do-  ACIGV.freeze edgesG---- | Returns the current internal state of the edges: @(from, to, cap, flow, cost)@, but without--- making copy. The edges are ordered in the same order as added by `addEdge`.------ ==== Complexity--- - \(O(1)\)------ @since 1.0.0.0-{-# INLINE unsafeEdges #-}-unsafeEdges ::-  (HasCallStack, PrimMonad m, Num cap, Ord cap, VU.Unbox cap, Num cost, Ord cost, VU.Unbox cost) =>-  -- | Graph-  McfGraph (PrimState m) cap cost ->-  -- | Vector of @(from, to, cap, flow, cost)@-  m (VU.Vector (Int, Int, cap, cap, cost))-unsafeEdges McfGraph {..} = do-  ACIGV.unsafeFreeze edgesG
src/AtCoder/SegTree.hs view
@@ -106,6 +106,7 @@ import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Bit qualified as ACIBIT import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bits (countTrailingZeros, testBit, (.&.), (.>>.)) import Data.Foldable (for_) import Data.Vector.Generic.Mutable qualified as VGM@@ -145,9 +146,7 @@ -- @since 1.0.0.0 {-# INLINE new #-} new :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => Int -> m (SegTree (PrimState m) a)-new nSt-  | nSt >= 0 = build $ VU.replicate nSt mempty-  | otherwise = error $ "AtCoder.SegTree.new: given negative size (`" ++ show nSt ++ "`)"+new n = stToPrim $ newST n  -- | Creates an array with initial values. --@@ -157,17 +156,7 @@ -- @since 1.0.0.0 {-# INLINE build #-} build :: (PrimMonad m, Monoid a, VU.Unbox a) => VU.Vector a -> m (SegTree (PrimState m) a)-build vs = do-  let nSt = VU.length vs-  let sizeSt = ACIBIT.bitCeil nSt-  let logSt = countTrailingZeros sizeSt-  dSt <- VUM.replicate (2 * sizeSt) mempty-  VU.iforM_ vs $ \i v -> do-    VGM.write dSt (sizeSt + i) v-  let segtree = SegTree {..}-  for_ [sizeSt - 1, sizeSt - 2 .. 1] $ \i -> do-    update segtree i-  pure segtree+build vs = stToPrim $ buildST vs  -- | Writes \(p\)-th value of the array to \(x\). --@@ -180,11 +169,7 @@ -- @since 1.0.0.0 {-# INLINE write #-} write :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m ()-write self@SegTree {..} p x = do-  let !_ = ACIA.checkIndex "AtCoder.SegTree.write" p nSt-  VGM.write dSt (p + sizeSt) x-  for_ [1 .. logSt] $ \i -> do-    update self ((p + sizeSt) .>>. i)+write self p x = stToPrim $ writeST self p x  -- | (Extra API) Modifies \(p\)-th value with a function \(f\). --@@ -197,11 +182,7 @@ -- @since 1.0.0.0 {-# INLINE modify #-} modify :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> (a -> a) -> Int -> m ()-modify self@SegTree {..} f p = do-  let !_ = ACIA.checkIndex "AtCoder.SegTree.modify" p nSt-  VGM.modify dSt f (p + sizeSt)-  for_ [1 .. logSt] $ \i -> do-    update self ((p + sizeSt) .>>. i)+modify self f p = stToPrim $ modifyST self f p  -- | (Extra API) Modifies \(p\)-th value with a monadic function \(f\). --@@ -217,8 +198,8 @@ modifyM self@SegTree {..} f p = do   let !_ = ACIA.checkIndex "AtCoder.SegTree.modifyM" p nSt   VGM.modifyM dSt f (p + sizeSt)-  for_ [1 .. logSt] $ \i -> do-    update self ((p + sizeSt) .>>. i)+  stToPrim $ for_ [1 .. logSt] $ \i -> do+    updateST self ((p + sizeSt) .>>. i)  -- | (Extra API) Writes \(p\)-th value of the array to \(x\) and returns the old value. --@@ -231,13 +212,7 @@ -- @since 1.1.0.0 {-# INLINE exchange #-} exchange :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> a -> m a-exchange self@SegTree {..} p x = do-  let !_ = ACIA.checkIndex "AtCoder.SegTree.exchange" p nSt-  ret <- VGM.exchange dSt (p + sizeSt) x-  VGM.write dSt (p + sizeSt) x-  for_ [1 .. logSt] $ \i -> do-    update self ((p + sizeSt) .>>. i)-  pure ret+exchange self p x = stToPrim $ exchangeST self p x  -- | Returns \(p\)-th value of the array. --@@ -267,7 +242,7 @@ {-# INLINE prod #-} prod :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m a prod self@SegTree {nSt} l0 r0-  | ACIA.testInterval l0 r0 nSt = unsafeProd self l0 r0+  | ACIA.testInterval l0 r0 nSt = stToPrim $ unsafeProdST self l0 r0   | otherwise = ACIA.errorInterval "AtCoder.SegTree.prod" l0 r0 nSt  -- | Total variant of `prod`. Returns \(a[l] \cdot ... \cdot a[r - 1]\), assuming the properties of@@ -281,29 +256,10 @@ {-# INLINE prodMaybe #-} prodMaybe :: (HasCallStack, PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m (Maybe a) prodMaybe self@SegTree {nSt} l0 r0-  | ACIA.testInterval l0 r0 nSt = Just <$> unsafeProd self l0 r0+  | ACIA.testInterval l0 r0 nSt = stToPrim $ Just <$> unsafeProdST self l0 r0   -- l0 == r0 = pure (Just mempty)   | otherwise = pure Nothing --- | Internal implementation of `prod`.-{-# INLINE unsafeProd #-}-unsafeProd :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> Int -> m a-unsafeProd SegTree {..} l0 r0 = inner (l0 + sizeSt) (r0 + sizeSt - 1) mempty mempty-  where-    -- NOTE: we're using inclusive range [l, r] for simplicity-    inner l r !smL !smR-      | l > r = pure $! smL <> smR-      | otherwise = do-          !smL' <--            if testBit l 0-              then (smL <>) <$> VGM.read dSt l-              else pure smL-          !smR' <--            if not $ testBit r 0-              then (<> smR) <$> VGM.read dSt r-              else pure smR-          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) smL' smR'- -- | Returns @a[0] <> ... <> a[n - 1]@, assuming the properties of the monoid. It returns `mempty` -- if \(n = 0\). --@@ -510,10 +466,77 @@ unsafeFreeze SegTree {..} = do   VU.unsafeFreeze . VUM.take nSt $ VUM.drop sizeSt dSt --- | \(O(1)\)-{-# INLINE update #-}-update :: (PrimMonad m, Monoid a, VU.Unbox a) => SegTree (PrimState m) a -> Int -> m ()-update SegTree {..} k = do+-- -------------------------------------------------------------------------------------------------+-- Internal+-- -------------------------------------------------------------------------------------------------++{-# INLINEABLE newST #-}+newST :: (HasCallStack, Monoid a, VU.Unbox a) => Int -> ST s (SegTree s a)+newST nSt+  | nSt >= 0 = build $ VU.replicate nSt mempty+  | otherwise = error $ "AtCoder.SegTree.newST: given negative size (`" ++ show nSt ++ "`)"++{-# INLINEABLE buildST #-}+buildST :: (Monoid a, VU.Unbox a) => VU.Vector a -> ST s (SegTree s a)+buildST vs = do+  let nSt = VU.length vs+  let sizeSt = ACIBIT.bitCeil nSt+  let logSt = countTrailingZeros sizeSt+  dSt <- VUM.replicate (2 * sizeSt) mempty+  VU.iforM_ vs $ \i v -> do+    VGM.write dSt (sizeSt + i) v+  let segtree = SegTree {..}+  for_ [sizeSt - 1, sizeSt - 2 .. 1] $ \i -> do+    updateST segtree i+  pure segtree++{-# INLINEABLE writeST #-}+writeST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> Int -> a -> ST s ()+writeST self@SegTree {..} p x = do+  let !_ = ACIA.checkIndex "AtCoder.SegTree.writeST" p nSt+  VGM.write dSt (p + sizeSt) x+  for_ [1 .. logSt] $ \i -> do+    updateST self ((p + sizeSt) .>>. i)++{-# INLINEABLE modifyST #-}+modifyST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> (a -> a) -> Int -> ST s ()+modifyST self@SegTree {..} f p = do+  let !_ = ACIA.checkIndex "AtCoder.SegTree.modifyST" p nSt+  VGM.modify dSt f (p + sizeSt)+  for_ [1 .. logSt] $ \i -> do+    updateST self ((p + sizeSt) .>>. i)++{-# INLINEABLE exchangeST #-}+exchangeST :: (HasCallStack, Monoid a, VU.Unbox a) => SegTree s a -> Int -> a -> ST s a+exchangeST self@SegTree {..} p x = do+  let !_ = ACIA.checkIndex "AtCoder.SegTree.exchangeST" p nSt+  ret <- VGM.exchange dSt (p + sizeSt) x+  VGM.write dSt (p + sizeSt) x+  for_ [1 .. logSt] $ \i -> do+    updateST self ((p + sizeSt) .>>. i)+  pure ret++{-# INLINEABLE unsafeProdST #-}+unsafeProdST :: (Monoid a, VU.Unbox a) => SegTree s a -> Int -> Int -> ST s a+unsafeProdST SegTree {..} l0 r0 = inner (l0 + sizeSt) (r0 + sizeSt - 1) mempty mempty+  where+    -- NOTE: we're using inclusive range [l, r] for simplicity+    inner l r !smL !smR+      | l > r = pure $! smL <> smR+      | otherwise = do+          !smL' <-+            if testBit l 0+              then (smL <>) <$> VGM.read dSt l+              else pure smL+          !smR' <-+            if not $ testBit r 0+              then (<> smR) <$> VGM.read dSt r+              else pure smR+          inner ((l + 1) .>>. 1) ((r - 1) .>>. 1) smL' smR'++{-# INLINE updateST #-}+updateST :: (Monoid a, VU.Unbox a) => SegTree s a -> Int -> ST s ()+updateST SegTree {..} k = do   opL <- VGM.read dSt $ 2 * k   opR <- VGM.read dSt $ 2 * k + 1   VGM.write dSt k $! opL <> opR
src/AtCoder/String.hs view
@@ -101,7 +101,7 @@ -- - \(O(n)\)-space -- -- @since 1.0.0.0-{-# INLINABLE suffixArrayOrd #-}+{-# INLINEABLE suffixArrayOrd #-} suffixArrayOrd :: (HasCallStack, Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int suffixArrayOrd s =   let n = VU.length s@@ -137,7 +137,7 @@ -- - \(O(n)\) -- -- @since 1.0.0.0-{-# INLINABLE lcpArray #-}+{-# INLINEABLE lcpArray #-} lcpArray ::   (HasCallStack, Ord a, VU.Unbox a) =>   -- | A vector representing a string@@ -209,7 +209,7 @@ -- - \(O(n)\) -- -- @since 1.0.0.0-{-# INLINABLE zAlgorithm #-}+{-# INLINEABLE zAlgorithm #-} zAlgorithm :: (Ord a, VU.Unbox a) => VU.Vector a -> VU.Vector Int zAlgorithm s   | n == 0 = VU.empty
src/AtCoder/TwoSat.hs view
@@ -23,10 +23,13 @@ module AtCoder.TwoSat   ( -- * TwoSat     TwoSat (nTs),+     -- * Constructor     new,+     -- * Clause building     addClause,+     -- * Solvers     satisfiable,     answer,@@ -36,7 +39,8 @@  import AtCoder.Internal.Assert qualified as ACIA import AtCoder.Internal.Scc qualified as ACISCC-import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim)+import Control.Monad.ST (ST) import Data.Bit (Bit (..)) import Data.Vector.Generic qualified as VG import Data.Vector.Generic.Mutable qualified as VGM@@ -102,7 +106,11 @@ -- @since 1.0.0.0 {-# INLINE satisfiable #-} satisfiable :: (PrimMonad m) => TwoSat (PrimState m) -> m Bool-satisfiable TwoSat {..} = do+satisfiable = stToPrim . satisfiableST++{-# INLINEABLE satisfiableST #-}+satisfiableST :: TwoSat s -> ST s Bool+satisfiableST TwoSat {..} = do   (!_, !ids) <- ACISCC.sccIds sccTs   let inner i         | i >= nTs = pure True