{-# LANGUAGE RecordWildCards #-}

-- | Fast modular multiplication by barrett reduction.
module BenchLib.Vector.ConcatMapM
  ( primConcatMapM1,
    concatMapM1,
    concatMapM2,
    primIConcatMapM1,
    iconcatMapM1,
    iconcatMapM2,
  )
where

import Control.Monad.Primitive (PrimMonad)
import Data.Vector.Fusion.Bundle.Monadic qualified as BundleM
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM

-- I expected these functions do not perform well with State monad, but
-- they're working as fast as IO and ST?

-- | https://github.com/haskell/vector/issues/416
--
-- With this wrapper
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

-- | PrimMonad only. This addes more constraints, but performs well with complex monads.
{-# INLINE primConcatMapM1 #-}
primConcatMapM1 ::
  (PrimMonad m, VG.Vector v a, VG.Vector v b) =>
  (a -> m (v b)) ->
  v a ->
  m (v b)
primConcatMapM1 f =
  unstreamPrimM
    . BundleM.concatVectors
    . BundleM.mapM f
    . BundleM.fromVector

{-# INLINE concatMapM1 #-}
concatMapM1 ::
  (Monad m, VG.Vector v a, VG.Vector v b) =>
  (a -> m (v b)) ->
  v a ->
  m (v b)
concatMapM1 f =
  VG.unstreamM
    . BundleM.concatVectors
    . BundleM.mapM f
    . BundleM.fromVector

{-# INLINE concatMapM2 #-}
concatMapM2 ::
  (Monad m, VG.Vector v a, VG.Vector v b) =>
  (a -> m (v b)) ->
  v a ->
  m (v b)
concatMapM2 f xs =
  VG.unstreamM
    . BundleM.concatVectors
    $ BundleM.generateM (VG.length xs) (f . VG.unsafeIndex xs)

-- | PrimMonad only. This addes more constraints, but performs well with complex monads.
{-# INLINE primIConcatMapM1 #-}
primIConcatMapM1 ::
  (PrimMonad m, VG.Vector v a, VG.Vector v b) =>
  (Int -> a -> m (v b)) ->
  v a ->
  m (v b)
primIConcatMapM1 f =
  unstreamPrimM
    . BundleM.concatVectors
    . BundleM.mapM (uncurry f)
    . BundleM.indexed
    . BundleM.fromVector

{-# INLINE iconcatMapM1 #-}
iconcatMapM1 ::
  (Monad m, VG.Vector v a, VG.Vector v b) =>
  (Int -> a -> m (v b)) ->
  v a ->
  m (v b)
iconcatMapM1 f =
  VG.unstreamM
    . BundleM.concatVectors
    . BundleM.mapM (uncurry f)
    . BundleM.indexed
    . BundleM.fromVector

{-# INLINE iconcatMapM2 #-}
iconcatMapM2 ::
  (Monad m, VG.Vector v a, VG.Vector v b) =>
  (Int -> a -> m (v b)) ->
  v a ->
  m (v b)
iconcatMapM2 f xs =
  VG.unstreamM
    . BundleM.concatVectors
    $ BundleM.generateM (VG.length xs) (\i -> f i (VG.unsafeIndex xs i))