content: -- | Monadic vector functions with @PrimMonad m@ constraints. They can fuse well.
  --
  -- Related issue: https://www.github.com/haskell/vector/issues/416.
  --
  -- Note that these functions are not 100% guaranteed to be sound.
  --
  -- @since 1.5.2.0
  module AtCoder.Extra.Vector.Prim
    ( -- * Construction
  
      -- ** Monadic initialization
      replicateM,
      generateM,
      iterateNM,
  
      -- ** Unfolding
  
      -- TODO: implement unfoldrM
      constructNM,
      constructrNM,
  
      -- * Elementwise operations
  
      -- ** Monadic mapping
      mapM,
      mapM_,
      imapM,
      imapM_,
      forM,
      forM_,
      iforM,
      iforM_,
  
      -- ** Monadic zipping
      zipWithM,
      zipWithM_,
      izipWithM,
      izipWithM_,
  
      -- * Concat map
      concatMapM,
      iconcatMapM,
  
      -- * Working with predicates
  
      -- ** Filtering
      filterM,
      mapMaybeM,
      imapMaybeM,
  
      -- ** Monadic scanl
      prescanlM,
      prescanlM',
      postscanlM,
      postscanlM',
      scanlM,
      scanlM',
      scanl1M,
      scanl1M',
    )
  where
  
  import Control.Monad.Primitive (PrimMonad)
  import Data.Vector.Fusion.Bundle qualified as Bundle
  import Data.Vector.Fusion.Bundle.Monadic qualified as BundleM
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Generic.Mutable qualified as VGM
  import GHC.Stack (HasCallStack)
  import Prelude hiding (mapM, mapM_)
  
  -- | ==== Constraints
  -- - \(n \ge 0\)
  --
  -- @since 1.5.2.0
  {-# INLINE replicateM #-}
  replicateM :: (PrimMonad m, VG.Vector v a) => Int -> m a -> m (v a)
  replicateM n f = unstreamPrimM (BundleM.replicateM n f)
  
  -- | ==== Constraints
  -- - \(n \ge 0\)
  --
  -- @since 1.5.2.0
  {-# INLINE generateM #-}
  generateM :: (PrimMonad m, VG.Vector v a) => Int -> (Int -> m a) -> m (v a)
  generateM n f = unstreamPrimM (BundleM.generateM n f)
  
  -- | ==== Constraints
  -- - \(n \ge 0\)
  --
  -- @since 1.5.2.0
  {-# INLINE iterateNM #-}
  iterateNM :: (PrimMonad m, VG.Vector v a) => Int -> (a -> m a) -> a -> m (v a)
  iterateNM n f x = unstreamPrimM (BundleM.iterateNM n f x)
  
  -- | ==== Constraints
  -- - \(n \ge 0\)
  --
  -- @since 1.5.2.0
  {-# INLINE constructNM #-}
  constructNM :: forall m v a. (PrimMonad m, VG.Vector v a) => Int -> (v a -> m a) -> m (v a)
  constructNM n f = do
    v <- VGM.new n
    v' <- VG.unsafeFreeze v
    fill v' 0
    where
      fill :: v a -> Int -> m (v a)
      fill !v i
        | i < n = do
            x <- f (VG.unsafeTake i v)
            VG.elemseq v x $ do
              v' <- VG.unsafeThaw v
              VGM.unsafeWrite v' i x
              v'' <- VG.unsafeFreeze v'
              fill v'' (i + 1)
      fill v _ = pure v
  
  -- | ==== Constraints
  -- - \(n \ge 0\)
  --
  -- @since 1.5.2.0
  {-# INLINE constructrNM #-}
  constructrNM :: forall m v a. (PrimMonad m, VG.Vector v a) => Int -> (v a -> m a) -> m (v a)
  constructrNM n f = do
    v <- n `seq` VGM.new n
    v' <- VG.unsafeFreeze v
    fill v' 0
    where
      fill :: v a -> Int -> m (v a)
      fill !v i
        | i < n = do
            x <- f (VG.unsafeSlice (n - i) i v)
            VG.elemseq v x $ do
              v' <- VG.unsafeThaw v
              VGM.unsafeWrite v' (n - i - 1) x
              v'' <- VG.unsafeFreeze v'
              fill v'' (i + 1)
      fill v _ = pure v
  
  -- | @since 1.5.2.0
  {-# INLINE mapM #-}
  mapM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> m b) -> v a -> m (v b)
  mapM f = unstreamPrimM . Bundle.mapM f . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE imapM #-}
  imapM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (Int -> a -> m b) -> v a -> m (v b)
  imapM f = unstreamPrimM . Bundle.mapM (uncurry f) . Bundle.indexed . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE mapM_ #-}
  mapM_ :: (PrimMonad m, VG.Vector v a) => (a -> m b) -> v a -> m ()
  mapM_ f = Bundle.mapM_ f . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE imapM_ #-}
  imapM_ :: (PrimMonad m, VG.Vector v a) => (Int -> a -> m b) -> v a -> m ()
  imapM_ f = Bundle.mapM_ (uncurry f) . Bundle.indexed . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE forM #-}
  forM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => v a -> (a -> m b) -> m (v b)
  forM as f = mapM f as
  
  -- | @since 1.5.2.0
  {-# INLINE forM_ #-}
  forM_ :: (PrimMonad m, VG.Vector v a) => v a -> (a -> m b) -> m ()
  forM_ as f = mapM_ f as
  
  -- | @since 1.5.2.0
  iforM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => v a -> (Int -> a -> m b) -> m (v b)
  {-# INLINE iforM #-}
  iforM as f = imapM f as
  
  -- | @since 1.5.2.0
  {-# INLINE iforM_ #-}
  iforM_ :: (PrimMonad m, VG.Vector v a) => v a -> (Int -> a -> m b) -> m ()
  iforM_ as f = imapM_ f as
  
  -- | @since 1.5.2.0
  {-# INLINE zipWithM #-}
  zipWithM :: (PrimMonad m, VG.Vector v a, VG.Vector v b, VG.Vector v c) => (a -> b -> m c) -> v a -> v b -> m (v c)
  zipWithM f = \as bs -> unstreamPrimM $ Bundle.zipWithM f (VG.stream as) (VG.stream bs)
  
  -- | @since 1.5.2.0
  {-# INLINE izipWithM #-}
  izipWithM :: (PrimMonad m, VG.Vector v a, VG.Vector v b, VG.Vector v c) => (Int -> a -> b -> m c) -> v a -> v b -> m (v c)
  izipWithM m as bs = unstreamPrimM . Bundle.zipWithM (uncurry m) (Bundle.indexed (VG.stream as)) $ VG.stream bs
  
  -- | @since 1.5.2.0
  {-# INLINE zipWithM_ #-}
  zipWithM_ :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m c) -> v a -> v b -> m ()
  zipWithM_ f = \as bs -> Bundle.zipWithM_ f (VG.stream as) (VG.stream bs)
  
  -- | @since 1.5.2.0
  {-# INLINE izipWithM_ #-}
  izipWithM_ :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (Int -> a -> b -> m c) -> v a -> v b -> m ()
  izipWithM_ m as bs = Bundle.zipWithM_ (uncurry m) (Bundle.indexed (VG.stream as)) $ VG.stream bs
  
  -- | Maps each element to a vector and concatenate the results.
  --
  -- @since 1.5.2.0
  {-# INLINE concatMapM #-}
  concatMapM ::
    (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) =>
    (a -> m (v b)) ->
    v a ->
    m (v b)
  concatMapM f =
    unstreamPrimM
      . BundleM.concatVectors
      . BundleM.mapM f
      . BundleM.fromVector
  
  -- | Maps each element to a vector and concatenate the results.
  --
  -- @since 1.5.2.0
  {-# INLINE iconcatMapM #-}
  iconcatMapM ::
    (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) =>
    (Int -> a -> m (v b)) ->
    v a ->
    m (v b)
  iconcatMapM f =
    unstreamPrimM
      . BundleM.concatVectors
      . BundleM.mapM (uncurry f)
      . BundleM.indexed
      . BundleM.fromVector
  
  -- | @since 1.5.2.0
  {-# INLINE filterM #-}
  filterM :: (PrimMonad m, VG.Vector v a) => (a -> m Bool) -> v a -> m (v a)
  filterM f = unstreamPrimM . Bundle.filterM f . VG.stream
  
  -- | @since 1.5.2.0
  mapMaybeM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> m (Maybe b)) -> v a -> m (v b)
  {-# INLINE mapMaybeM #-}
  mapMaybeM f = unstreamPrimM . Bundle.mapMaybeM f . VG.stream
  
  -- | @since 0.12.2.0
  {-# INLINE imapMaybeM #-}
  imapMaybeM :: (PrimMonad m, VG.Vector v a, VG.Vector v b) => (Int -> a -> m (Maybe b)) -> v a -> m (v b)
  imapMaybeM f = unstreamPrimM . Bundle.mapMaybeM (\(i, a) -> f i a) . Bundle.indexed . VG.stream
  
  -- | https://github.com/haskell/vector/issues/416
  {-# INLINE [1] unstreamPrimM #-}
  unstreamPrimM :: (PrimMonad m, VG.Vector v a) => BundleM.Bundle m u a -> m (v a)
  unstreamPrimM s = VGM.munstream s >>= VG.unsafeFreeze
  
  -- | @since 1.5.2.0
  {-# INLINE prescanlM #-}
  prescanlM :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  prescanlM f x0 =
    unstreamPrimM
      . prescanlMB f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE prescanlM' #-}
  prescanlM' :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  prescanlM' f x0 =
    unstreamPrimM
      . prescanlMB' f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE postscanlM #-}
  postscanlM :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  postscanlM f x0 =
    unstreamPrimM
      . postscanlMB f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE postscanlM' #-}
  postscanlM' :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  postscanlM' f x0 =
    unstreamPrimM
      . postscanlMB' f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE scanlM #-}
  scanlM :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  scanlM f x0 =
    unstreamPrimM
      . scanlMB f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE scanlM' #-}
  scanlM' :: (HasCallStack, PrimMonad m, VG.Vector v a, VG.Vector v b) => (a -> b -> m a) -> a -> v b -> m (v a)
  scanlM' f x0 =
    unstreamPrimM
      . scanlMB' f x0
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE scanl1M #-}
  scanl1M :: (HasCallStack, PrimMonad m, VG.Vector v a) => (a -> a -> m a) -> v a -> m (v a)
  scanl1M f =
    unstreamPrimM
      . scanl1MB f
      . VG.stream
  
  -- | @since 1.5.2.0
  {-# INLINE scanl1M' #-}
  scanl1M' :: (HasCallStack, PrimMonad m, VG.Vector v a) => (a -> a -> m a) -> v a -> m (v a)
  scanl1M' f =
    unstreamPrimM
      . scanl1MB' f
      . VG.stream
  
  {-# INLINE prescanlMB #-}
  prescanlMB :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  prescanlMB f x0 = BundleM.prescanlM f x0 . Bundle.lift
  
  {-# INLINE prescanlMB' #-}
  prescanlMB' :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  prescanlMB' f x0 = BundleM.prescanlM' f x0 . Bundle.lift
  
  {-# INLINE postscanlMB #-}
  postscanlMB :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  postscanlMB f x0 = BundleM.postscanlM f x0 . Bundle.lift
  
  {-# INLINE postscanlMB' #-}
  postscanlMB' :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  postscanlMB' f x0 = BundleM.postscanlM' f x0 . Bundle.lift
  
  {-# INLINE scanlMB #-}
  scanlMB :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  scanlMB f x0 = BundleM.scanlM f x0 . Bundle.lift
  
  {-# INLINE scanlMB' #-}
  scanlMB' :: (Monad m) => (a -> b -> m a) -> a -> Bundle.Bundle v b -> BundleM.Bundle m v a
  scanlMB' f x0 = BundleM.scanlM' f x0 . Bundle.lift
  
  {-# INLINE scanl1MB #-}
  scanl1MB :: (Monad m) => (a -> a -> m a) -> Bundle.Bundle v a -> BundleM.Bundle m v a
  scanl1MB f = BundleM.scanl1M f . Bundle.lift
  
  {-# INLINE scanl1MB' #-}
  scanl1MB' :: (Monad m) => (a -> a -> m a) -> Bundle.Bundle v a -> BundleM.Bundle m v a
  scanl1MB' f = BundleM.scanl1M' f . Bundle.lift
  