content: -- | Fast modular multiplication by barrett reduction.
  module BenchLib.Vector.MapAccumL
    ( mapAccumL1,
      mapAccumL2,
      mapAccumL3,
      mapAccumL4,
    )
  where
  
  import Control.Monad.Primitive (PrimMonad)
  import Control.Monad.ST (ST, runST)
  import Control.Monad.Trans.State.Strict (StateT, runState, runStateT, state)
  import Data.Vector.Fusion.Bundle qualified as Bundle
  import Data.Vector.Fusion.Bundle.Monadic qualified as BundleM
  import Data.Vector.Fusion.Bundle.Size qualified as Bundle
  import Data.Vector.Fusion.Stream.Monadic qualified as MS
  import Data.Vector.Generic qualified as VG
  import Data.Vector.Generic.Mutable qualified as VGM
  
  {-# INLINE mapAccumL1 #-}
  mapAccumL1 ::
    (VG.Vector v a, VG.Vector v b) =>
    (s -> a -> (s, b)) ->
    s ->
    v a ->
    (s, v b)
  mapAccumL1 f s0 xs = (\(!x, !s) -> (s, x)) $ runState (VG.mapM g xs) s0
    where
      g a =
        state
          ( \s ->
              let (!s', !b) = f s a
               in (b, s')
          )
  
  {-# INLINE mapAccumL2 #-}
  mapAccumL2 ::
    (VG.Vector v a, VG.Vector v b) =>
    (s -> a -> (s, b)) ->
    s ->
    v a ->
    (s, v b)
  mapAccumL2 f s0 xs = runST $ do
    vec <- VGM.unsafeNew (VG.length xs)
    !res <-
      VG.ifoldM'
        ( \ !s i x -> do
            let (!s', !x') = f s x
            VGM.write vec i x'
            pure s'
        )
        s0
        xs
    vec' <- VG.unsafeFreeze vec
    pure (res, vec')
  
  {-# INLINE mapAccumL3 #-}
  mapAccumL3 ::
    forall v s a b.
    (VG.Vector v a, VG.Vector v b) =>
    (s -> a -> (s, b)) ->
    s ->
    v a ->
    (s, v b)
  mapAccumL3 f s0 xs = (\(!x, !s) -> (s, x)) $ runST $ (`runStateT` s0) $ do
    unstreamPrimM
      . BundleM.mapM g
      $ BundleM.fromVector xs
    where
      g :: forall st. a -> StateT s (ST st) b
      g a =
        state
          ( \s ->
              let (!s', !b) = f s a
               in (b, s')
          )
  
  -- | https://github.com/haskell/vector/issues/416
  unstreamPrimM :: (PrimMonad m, VG.Vector v a) => BundleM.Bundle m u a -> m (v a)
  {-# INLINE [1] unstreamPrimM #-}
  unstreamPrimM s = VGM.munstream s >>= VG.unsafeFreeze
  
  -- @cojna/iota. Note that it does not return state value.
  {-# INLINE mapAccumL4 #-}
  mapAccumL4 ::
    (VG.Vector v a, VG.Vector v b) =>
    (s -> a -> (s, b)) ->
    s ->
    v a ->
    v b
  mapAccumL4 f x =
    VG.unstream
      . Bundle.inplace
        (streamAccumM (\s a -> pure (f s a)) x)
        Bundle.toMax
      . VG.stream
  
  -- @cojna/iota
  streamAccumM :: (Monad m) => (s -> a -> m (s, b)) -> s -> MS.Stream m a -> MS.Stream m b
  streamAccumM f s0 (MS.Stream step x0) = MS.Stream step' (s0, x0)
    where
      step' (!s, x) = do
        r <- step x
        case r of
          MS.Yield a x' -> do
            (s', b) <- f s a
            return $ MS.Yield b (s', x')
          MS.Skip x' -> return $ MS.Skip (s, x')
          MS.Done -> return MS.Done
      {-# INLINE [0] step' #-}
  {-# INLINE [1] streamAccumM #-}
  