vector-pull-0.1.0.0: lib/Data/Vector/Pull.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -Wno-dodgy-imports #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
{-# HLINT ignore "Redundant lambda" #-}
{-# HLINT ignore "Avoid lambda" #-}
{-# OPTIONS_GHC -Wno-unused-do-bind #-}

-- |
-- Module      : Data.Vector.Pull
-- Copyright   : (c) Michael Ledger 2024-2026
-- License     : MPL-2.0
-- Maintainer  : Michael Ledger <mike@quasimal.com>
--
-- "Pull" vectors are arrays represented as a function from an index to a value.
-- This representation allows for vector operations to be completely fused
-- together with little effort -- other than inlining. It works well in
-- scenarios where your vector operations are all chained together around a
-- single source location, *and* consumption/materialisation also takes place
-- there.
--
-- But beware! If operations do not fuse properly, you can pay for it dearly.
-- Every operation adds another closure which much be evaluated for every single
-- element access.
module Data.Vector.Pull (
  Pull,
  optimise,

  -- * Construction
  fromVector,
  safeFromVector,
  fromMVector,
  mapFromVector,
  fromList,
  empty,
  singleton,
  append,
  enumFromTo,
  enumFromLen,
  replicate,
  generate,
  surround,
  surroundMaybes,
  intersperse,
  intersperseWith,
  intersperseMapWith,
  cons,
  snoc,

  -- ** Manipulation
  take,
  drop,

  -- * Use
  (!),
  (!?),
  set,
  modify,

  -- * Consuming 'Pull' vectors
  toVector,
  toVectorA,
  toVectorM,
  chunkPullM,
  toListM,
  toStreamM,
  toStream,
  writeToVector,
  writeToVectorM,
  mapToVector,
  uncons,
  head,
  last,

  -- ** Concurrent consumers
  concurrentToVectorM,
  concurrentChunkedToVectorM,

  -- * Size
  length,

  -- * Mapping
  map,
  imap,
  zipWith,
  mapWithNext,
  enumerate,
  Enumerated (Enumerated),

  -- * Folds
  foldr,
  foldr',
  foldl,
  foldl',
  folded,

  -- ** Indexed variants
  Data.Vector.Pull.ifoldr,
  Data.Vector.Pull.ifoldr',
  Data.Vector.Pull.ifoldl,
  Data.Vector.Pull.ifoldl',

  -- * Traversals
  traverse,
  mapM,
  mapM_,

  -- ** Indexed variants
  imapM,
  imapM_,

  -- * Traversals
  update,
)
where

import Control.Concurrent (newEmptyMVar, takeMVar, tryPutMVar)
import Control.Concurrent.Counter qualified as Counter
import Control.Monad (ap, liftM2, void, when, (>=>))
import Control.Monad.Catch (MonadCatch, MonadThrow (throwM), catchAll)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Coerce
import Data.Foldable qualified as Foldable
import Data.Functor ((<&>))
import Data.Monoid
import Data.Vector (Vector)
import Data.Vector qualified as V
import Data.Vector.Generic qualified as G
import Data.Vector.Generic.Mutable qualified as GM
import Data.Vector.Unboxed qualified as UV
import GHC.Base (Int#, quotInt#, quotRemInt#, tagToEnum#, (*#), (+#), (-#), (<#), (<=#), (==#), (>#), (>=#))
import GHC.Exts (TYPE)
import GHC.Int (Int (..))
import Optics.AffineTraversal (atraversalVL)
import Optics.At.Core (Index, IxValue, Ixed (ix))
import Optics.Fold (Fold, foldring)
import Optics.Indexed.Core (FoldableWithIndex (..), ifor_)
import Streaming (Of, Stream)
import Streaming.Internal qualified as S
import Streaming.Prelude qualified as S
import Text.Printf (printf)
import Text.Show (Show (show))
import Prelude hiding (
  atomically,
  cons,
  drop,
  empty,
  enumFromTo,
  folded,
  foldl,
  foldl',
  foldr,
  fromList,
  head,
  ifoldl,
  ifoldl',
  ifoldr,
  ifoldr',
  imap,
  last,
  length,
  map,
  mapM,
  mapM_,
  modify,
  newTVarIO,
  replicate,
  set,
  show,
  snoc,
  take,
  toList,
  traverse,
  uncons,
  zipWith,
  (!?),
 )

-- | An index-based data structure. Conceptually it is just a @ Int -> a @. This
-- should have faster indexing, but slower everything-else. When your
-- construction of a 'Pull' is able to inline, the results can be extremely
-- efficient. When inlining is not available, modifications to the generator
-- layer closures on top of eachother.
data Pull (a :: TYPE r) = Pull
  { index# :: Int# -> a
  , length# :: Int#
  }

instance Functor Pull where
  {-# INLINE fmap #-}
  fmap = map

instance Applicative Pull where
  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}
  {-# INLINE liftA2 #-}
  pure = singleton
  liftA2 = liftM2
  (<*>) = ap

instance Monad Pull where
  {-# INLINE (>>=) #-}
  m >>= k = foldr (append . k) empty m

instance Foldable Pull where
  {-# INLINE foldr #-}
  {-# INLINE foldr' #-}
  {-# INLINE foldl #-}
  {-# INLINE foldl' #-}
  {-# INLINE toList #-}
  foldr = Data.Vector.Pull.foldr
  foldr' = Data.Vector.Pull.foldr'
  foldl = Data.Vector.Pull.foldl
  foldl' = Data.Vector.Pull.foldl'
  length = Data.Vector.Pull.length
  toList = Data.Vector.Pull.toList

instance FoldableWithIndex Int Pull where
  ifoldMap = Data.Vector.Pull.ifoldMap
  ifoldMap' = Data.Vector.Pull.ifoldMap'
  ifoldr = Data.Vector.Pull.ifoldr
  ifoldr' = Data.Vector.Pull.ifoldr'
  ifoldl = Data.Vector.Pull.ifoldl
  ifoldl' = Data.Vector.Pull.ifoldl'

instance (Show a) => Show (Pull a) where
  show = show . toVector @V.Vector

instance Semigroup (Pull a) where
  (<>) = append

instance Monoid (Pull a) where
  mempty = empty

type instance Index (Pull a) = Int

type instance IxValue (Pull a) = a

instance Ixed (Pull a) where
  {-# INLINE ix #-}
  ix i = atraversalVL \point f x ->
    if 0 <= i && i < length x
      then f (x ! i) <&> set x i
      else point x

{-# INLINE optimise #-}

-- | /O(N * D)/
optimise :: forall v a. (G.Vector v a) => Pull a -> Pull a
optimise = fromVector @v . toVector

--------------------------------------------------------------------------------
-- Construction

{-# INLINE fromVector #-}
fromVector :: (G.Vector v a) => v a -> Pull a
fromVector !vec =
  Pull
    { length# = case G.length vec of I# a -> a
    , index# = G.unsafeIndex vec .# I#
    }

{-# INLINE safeFromVector #-}
safeFromVector :: (G.Vector v a) => v a -> Pull a
safeFromVector !vec =
  Pull
    { length# = case G.length vec of I# a -> a
    , index# = \i -> vec G.! I# i
    }

{-# INLINE fromMVector #-}
fromMVector :: (G.Vector v a, GM.PrimMonad m) => G.Mutable v (GM.PrimState m) a -> Pull (m a)
fromMVector !vec =
  Pull
    { length# = case GM.length vec of I# a -> a
    , index# = \i -> GM.unsafeRead vec (I# i)
    }

{-# INLINE mapFromVector #-}
mapFromVector :: forall v b a. (G.Vector v a) => (a -> b) -> v a -> Pull b
mapFromVector f = fmap f . fromVector

fromList :: [a] -> Pull a
fromList = fromVector . V.fromList

{-# INLINE empty #-}
empty :: Pull a
empty =
  Pull
    { index# = outOfBoundsError# "empty" 0#
    , length# = 0#
    }

{-# INLINE singleton #-}
singleton :: a -> Pull a
singleton a =
  Pull
    { index# = \i -> case i of
        0# -> a
        _ -> outOfBoundsError# "singleton" 1# i
    , length# = 1#
    }

{-# INLINE cons #-}
cons :: a -> Pull a -> Pull a
cons x Pull {index#, length#} =
  Pull
    { index# = \i# -> case i# ># 0# of
        1# -> index# (i# -# 1#)
        _ -> x
    , length# = length# +# 1#
    }

{-# INLINE snoc #-}
snoc :: Pull a -> a -> Pull a
snoc Pull {index#, length#} x =
  Pull
    { index# = \i# -> case i# <# length# of
        1# -> index# i#
        _ -> x
    , length# = length# +# 1#
    }

{-# INLINE append #-}
append :: Pull a -> Pull a -> Pull a
append a b =
  Pull
    { length# = length# a +# length# b
    , index# =
        \i -> case i <# length# a of
          1# -> index# a i
          _ -> index# b (i -# length# a)
    }

{-# INLINE surround #-}

-- | 'cons' and 'snoc' at the same time
surround :: a -> Pull a -> a -> Pull a
surround l Pull {length#, index#} r =
  Pull
    { length# = length# +# 2#
    , index# = \i# -> case i# <# 1# of
        1# -> l
        _ -> case i# <# (length# +# 1#) of
          1# -> index# (i# -# 1#)
          _ -> r
    }

{-# INLINE surroundMaybes #-}

-- | 'cons' and 'snoc' at the same time with optional elements
surroundMaybes :: Maybe a -> Pull a -> Maybe a -> Pull a
surroundMaybes (Just l) g (Just r) = surround l g r
surroundMaybes (Just l) g _ = cons l g
surroundMaybes _ g (Just r) = snoc g r
surroundMaybes _ g _ = g

{-# INLINE intersperseMapWith #-}
intersperseMapWith
  :: (b -> b -> b)
  -- ^ How to compute midpoints between elements
  -> (a -> b)
  -- ^ Mapping function
  -> Pull a
  -> Pull b
intersperseMapWith avg lift Pull {index#, length#} = case length# <=# 1# of
  1# -> Pull {index# = \i# -> lift (index# i#), length#}
  _ ->
    Pull
      { length# = length# *# 2# -# 1#
      , index# = \j# ->
          case quotRemInt# j# 2# of
            (# i#, 0# #) -> lift (index# i#)
            (# i#, _ #) -> avg (lift (index# i#)) (lift (index# (i# +# 1#)))
      }

{-# INLINE intersperseWith #-}
intersperseWith :: (a -> a -> a) -> Pull a -> Pull a
intersperseWith fn = intersperseMapWith fn id

{-# INLINE intersperse #-}
intersperse :: a -> Pull a -> Pull a
intersperse x = intersperseMapWith (\_ _ -> x) id

{-# INLINE enumFromTo #-}
enumFromTo :: (Enum a) => a -> a -> Pull a
enumFromTo (fromEnum -> I# a) (fromEnum -> I# b) =
  Pull
    { length# = (b -# a) +# 1#
    , index# = \i -> toEnum (I# (a +# i))
    }

{-# INLINE enumFromLen #-}
enumFromLen :: (Enum a) => a -> Int -> Pull a
enumFromLen (fromEnum -> I# a) (I# length#) =
  Pull
    { length#
    , index# = \i -> toEnum (I# (a +# i))
    }

{-# INLINE replicate #-}
replicate :: Int -> a -> Pull a
replicate (I# l) a =
  Pull
    { length# = l
    , index# = const# a
    }

{-# INLINE generate #-}
generate :: Int -> (Int -> a) -> Pull a
generate (I# l) f =
  Pull
    { length# = l
    , index# = \i -> f (I# i)
    }

--------------------------------------------------------------------------------

{-# INLINE take #-}
take :: Int -> Pull a -> Pull a
take (I# newLength#) Pull {index#, length#} =
  Pull {index#, length# = max# 0# (min# length# newLength#)}

{-# INLINE drop #-}
drop :: Int -> Pull a -> Pull a
drop (I# dropAmount#) it@Pull {index#, length#}
  | tagToEnum# (dropAmount# <# 0#) = it
  | tagToEnum# (length# ># dropAmount#) =
      Pull
        { length# = length# -# dropAmount#
        , index# = index# .+# dropAmount#
        }
  | otherwise = empty

--------------------------------------------------------------------------------
-- Updates

{-# INLINE set #-}

-- | /O(N)/ Beware
set :: Pull a -> Int -> a -> Pull a
set Pull {index#, length#} (I# i) a =
  Pull
    { length# = length#
    , index# = \j -> case j ==# i of
        1# -> a
        _ -> index# j
    }

{-# INLINE modify #-}

-- | /O(N)/ Beware
modify :: Pull a -> Int -> (a -> a) -> Pull a
modify Pull {index#, length#} (I# i) f =
  Pull
    { length# = length#
    , index# = \j -> case j ==# i of
        1# -> f (index# j)
        _ -> index# j
    }

{-# INLINE update #-}
update :: (Coercible i Int, UV.Unbox i) => Vector a -> (a -> b) -> UV.Vector i -> (Int -> b) -> Vector b
update v fn indices _indexFn
  | UV.null indices = V.map fn v
update v fn indices indexFn =
  V.unfoldrExactN
    (V.length v)
    step
    UpdState {i = 0, j = coerce (UV.head indices), next = 1}
  where
    numUpdates = UV.length indices
    step UpdState {i, j, next}
      | i == j =
          ( indexFn i
          , if next < numUpdates
              then UpdState (i + 1) (coerce (UV.unsafeIndex indices next)) (next + 1)
              else UpdState (i + 1) (-1) maxBound
          )
      | otherwise = (fn (v `V.unsafeIndex` i), UpdState {i = i + 1, j, next})

data UpdState = UpdState
  { i, j, next :: !Int
  }

--------------------------------------------------------------------------------
-- Materialisation

{-# INLINE toVector #-}

-- | /O(N * D)/
toVector :: forall v a. (G.Vector v a) => Pull a -> v a
toVector = mapToVector id

{-# INLINE toVectorA #-}

-- | /O(N * D)/
toVectorA :: forall v a f. (G.Vector v a, Applicative f) => Pull (f a) -> f (v a)
toVectorA g =
  G.fromListN (length g) <$> foldr (liftA2 (:)) (pure []) g

{-# INLINE traverse #-}
traverse :: (Applicative f) => (a -> f b) -> Pull a -> f (Pull b)
traverse f x = fromVector @Vector <$> toVectorA (map f x)

{-# INLINE toVectorM #-}
toVectorM :: forall v a m. (GM.PrimMonad m, G.Vector v a) => Pull (m a) -> m (v a)
toVectorM = toMutVectorM >=> G.unsafeFreeze

{-# INLINE toMutVectorM #-}
toMutVectorM :: forall v a m. (GM.PrimMonad m, G.Vector v a) => Pull (m a) -> m (G.Mutable v (GM.PrimState m) a)
toMutVectorM Pull {length#, index#} = do
  result <- GM.unsafeNew (I# length#)
  let go i = case i <# length# of
        1# -> do
          GM.unsafeWrite result (I# i) =<< index# i
          go (i +# 1#)
        _ -> pass
  go 0#
  pure result

data Enumerated a = Enumerated# Int# ~a

pattern Enumerated :: Int -> a -> Enumerated a
pattern Enumerated x a <- Enumerated# (I# -> x) a
  where
    Enumerated (I# x) a = Enumerated# x a

{-# COMPLETE Enumerated #-}

{-# INLINE enumerate #-}
enumerate :: Pull a -> Pull (Enumerated a)
enumerate Pull {length#, index#} =
  Pull
    { length#
    , index# = \i# -> Enumerated# i# (index# i#)
    }

{-# INLINE chunkPullM #-}
chunkPullM :: forall a m. (Monad m) => Int -> Pull (m a) -> Pull (Pull (m a))
chunkPullM chunkSize _gen | chunkSize <= 0 = error "chunkSize must be positive"
chunkPullM (I# chunkSize#) Pull {length#, index#} =
  Pull
    { length# =
        case r# of
          0# -> chunks#
          _ -> chunks# +# 1#
    , index# = \chunkIndex# ->
        Pull
          { index# = \j# -> index# ((chunkIndex# *# chunkSize#) +# j#)
          , length# = case chunkIndex# <# chunks# of
              1# -> chunkSize#
              _ -> r#
          }
    }
  where
    !(# chunks#, r# #) = length# `quotRemInt#` chunkSize#

{-# INLINE concurrentChunkedToVectorM #-}
concurrentChunkedToVectorM
  :: forall v a m
   . (G.Vector v a, MonadIO m, MonadCatch m, MonadThrow m, GM.PrimMonad m)
  => (m () -> m ())
  -- ^ Fork function, e.g. @ (void . forkIO) @
  -> Int
  -> Pull (m a)
  -> m (v a)
concurrentChunkedToVectorM _ _ Pull {length# = 0#} = pure G.empty
concurrentChunkedToVectorM forkIO caps gen@Pull {length#, index#} =
  if caps <= 1
    then toVectorM gen
    else do
      dest <- GM.unsafeNew (I# length#)
      signal <- liftIO newEmptyMVar

      let
        !(I# caps#) = caps
        !chunkSize# = max# 1# (length# `quotInt#` caps#)
        !(# chunks#, r# #) = length# `quotRemInt#` chunkSize#

      remaining <- liftIO case r# of
        0# -> Counter.new (I# chunks#)
        _ -> Counter.new (I# (chunks# +# 1#))

      let
        go# i# = case i# <# chunks# of
          1# -> do
            forkIO do
              catchAll
                (goChunk# (i# *# chunkSize#) chunkSize#)
                (liftIO . void . tryPutMVar signal . Just)
            go# (i# +# 1#)
          _ -> case r# of
            0# -> pass
            _ -> forkIO (goChunk# (i# *# chunkSize#) r#)

        goChunk# j# len# =
          case len# of
            0# -> do
              prevRemaining <- liftIO (Counter.sub remaining 1)
              when (prevRemaining <= 1) (void (liftIO (tryPutMVar signal Nothing {- don't block in case the atomic counter failed and somehow multiple threads wrote to the signal -})))
            _ -> do
              x <- index# j#
              GM.unsafeWrite dest (I# j#) x
              goChunk# (j# +# 1#) (len# -# 1#)

      go# 0#

      anyError <- liftIO (takeMVar signal)
      case anyError of
        Just err -> throwM err
        Nothing -> G.unsafeFreeze dest

{-# INLINE concurrentToVectorM #-}
concurrentToVectorM
  :: (G.Vector v a, MonadIO m, MonadCatch m, MonadThrow m, GM.PrimMonad m)
  => (m () -> m ())
  -- ^ Fork function e.g. @ (void . forkIO) @
  -> Pull (m a)
  -> m (v a)
concurrentToVectorM _ Pull {length# = 0#} = pure G.empty
concurrentToVectorM forkIO gen = do
  dest <- GM.unsafeNew (length gen)
  signal <- liftIO newEmptyMVar
  remaining <- liftIO (Counter.new (length gen))
  ifor_ gen \i x -> forkIO do
    catchAll
      ( do
          result <- x
          GM.unsafeWrite dest i result
          prevRemaining <- liftIO (Counter.sub remaining 1)
          when (prevRemaining <= 1) (void (liftIO (tryPutMVar signal Nothing {- don't block in case atomic counter fail -})))
      )
      (void . liftIO . tryPutMVar signal . Just)
  -- re-throwing won't do anything meaningful here

  anyError <- liftIO (takeMVar signal)
  case anyError of
    Just e -> throwM e
    Nothing -> G.unsafeFreeze dest

{-# INLINE toListM #-}
toListM :: (Monad m) => Pull (m a) -> m [a]
toListM Pull {length#, index#} =
  go 0#
  where
    go i = case i <# length# of
      1# -> do
        x <- index# i
        xs <- go (i +# 1#)
        pure (x : xs)
      _ ->
        pure []

{-# INLINE toStream #-}
toStream :: (Monad m) => Pull a -> Stream (Of a) m ()
toStream Pull {length#, index#} = go 0#
  where
    go i = case i <# length# of
      1# -> S.Step (index# i S.:> go (i +# 1#))
      _ -> S.Return ()

{-# INLINE toStreamM #-}
toStreamM :: (Monad m) => Pull (m a) -> Stream (Of a) m ()
toStreamM Pull {length#, index#} = go 0#
  where
    go i = case i <# length# of
      1# ->
        S.Effect
          ( index# i <&> \val ->
              S.Step (val S.:> go (i +# 1#))
          )
      _ -> S.Return ()

{-# INLINE writeToVector #-}
writeToVector :: (GM.PrimMonad m, GM.PrimState m ~ s, G.Vector v a) => G.Mutable v s a -> Pull a -> m ()
writeToVector dest Pull {length#, index#} = do
  let
    !(I# finalLength#) = min (GM.length dest) (I# length#)
    go i = case i <# finalLength# of
      1# -> do
        GM.unsafeWrite dest (I# i) (index# i)
        go (i +# 1#)
      _ -> pass
  go 0#

{-# INLINE writeToVectorM #-}
writeToVectorM :: (GM.PrimMonad m, GM.PrimState m ~ s, G.Vector v a) => G.Mutable v s a -> Pull (m a) -> m ()
writeToVectorM dest Pull {length#, index#} = do
  let
    !(I# finalLength#) = min (GM.length dest) (I# length#)
    go i = case i <# finalLength# of
      1# -> do
        GM.unsafeWrite dest (I# i) =<< index# i
        go (i +# 1#)
      _ -> pass
  go 0#

{-# INLINE toList #-}

-- | /O(N * D)/
toList :: Pull a -> [a]
toList = foldr (:) []

{-# INLINE uncons #-}
uncons :: Pull a -> Maybe (a, Pull a)
uncons Pull {length#, index#} = case length# of
  0# -> Nothing
  _ ->
    Just
      ( index# 0#
      , Pull
          { length# = length# -# 1#
          , index# = \j -> index# (j +# 1#)
          }
      )

--------------------------------------------------------------------------------
-- Size

{-# INLINE length #-}
length :: Pull a -> Int
length Pull {length#} = I# length#

--------------------------------------------------------------------------------
-- Internal utility

{-# INLINE slice# #-}
slice#
  :: Pull a
  -> (Int# -> a -> b -> b)
  -> b
  -> Int#
  -- ^ starting index (inclusive)
  -> Int#
  -- ^ end index (non-inclusive)
  -> b
slice# Pull {index#} f z0 start end = go start
  where
    go i = case i <# end of
      1# -> f i (index# i) (go (i +# 1#))
      _ -> z0

{-# INLINE mapToVector #-}
mapToVector :: forall v b a. (G.Vector v b) => (a -> b) -> Pull a -> v b
mapToVector _ Pull {length# = 0#} = G.empty
mapToVector f Pull {length#, index#} =
  G.create do
    m <- GM.new (I# length#)
    let go i# = case i# <# length# of
          1# -> do
            GM.unsafeWrite m (I# i#) $! f (index# i#)
            go (i# +# 1#)
          _ -> pass
    go 0#
    pure m

--------------------------------------------------------------------------------
-- Indexing

{-# INLINE (!) #-}

-- | /O(D)/
(!) :: Pull a -> Int -> a
(!) Pull {index#, length#} (I# i#)
  | tagToEnum# (i# <# length#) = index# i#
  | otherwise = outOfBoundsError# "(!)" i# length#

{-# INLINE (!?) #-}

-- | /O(D)/
(!?) :: Pull a -> Int -> Maybe a
(!?) Pull {length#, index#} (I# i#)
  | tagToEnum# (i# <# length#) && tagToEnum# (i# >=# 0#) = Just (index# i#)
  | otherwise = Nothing

--------------------------------------------------------------------------------
-- Maps

{-# INLINE map #-}

{-# INLINE imap #-}

{-# INLINE zipWith #-}

{-# INLINE mapWithNext #-}

-- | /O(1)/
map :: (a -> b) -> Pull a -> Pull b
map f Pull {length#, index#} =
  Pull
    { length#
    , index# = \i -> f (index# i)
    }

-- | /O(1)/
imap :: (Int -> a -> b) -> Pull a -> Pull b
imap f Pull {length#, index#} =
  Pull
    { length#
    , index# = \i -> f (I# i) (index# i)
    }

zipWith :: (a -> b -> c) -> Pull a -> Pull b -> Pull c
zipWith f Pull {length# = lenA#, index# = indexA#} Pull {length# = lenB#, index# = indexB#} =
  Pull
    { length# = min# lenA# lenB#
    , index# = \i# -> f (indexA# i#) (indexB# i#)
    }

mapWithNext :: (a -> a -> b) -> Pull a -> a -> Pull b
mapWithNext f Pull {length#, index#} endMarker =
  Pull
    { length#
    , index# = \i# ->
        f
          (index# i#)
          ( if tagToEnum# (i# <# lenMinus1#)
              then index# (i# +# 1#)
              else endMarker
          )
    }
  where
    lenMinus1# = length# -# 1#

--------------------------------------------------------------------------------
-- Folds

{-# INLINE foldr #-}

{-# INLINE ifoldr #-}

{-# INLINE foldl #-}

{-# INLINE ifoldl #-}

{-# INLINE foldl' #-}

{-# INLINE ifoldl' #-}

{-# INLINE ifoldMap #-}

{-# INLINE ifoldMap' #-}

-- | Right-associative fold
foldr :: (a -> b -> b) -> b -> Pull a -> b
foldr f z o@Pull {length#} = slice# o (\_ !x xs -> f x xs) z 0# length#

-- | Right-associative strict fold
foldr' :: (a -> b -> b) -> b -> Pull a -> b
foldr' f z o@Pull {length#} = slice# o (\_ !x !xs -> f x xs) z 0# length#

-- | Right-associative strict fold, with index
ifoldMap :: (Monoid m) => (Int -> a -> m) -> Pull a -> m
ifoldMap f o@Pull {length#} =
  slice#
    o
    (\i x xs -> f (I# i) x <> xs)
    mempty
    0#
    length#

-- | Right-associative strict fold, with index
ifoldMap' :: (Monoid m) => (Int -> a -> m) -> Pull a -> m
ifoldMap' f o@Pull {length#} =
  slice#
    o
    (\i !x !xs -> f (I# i) x <> xs)
    mempty
    0#
    length#

-- | Right-associative strict fold, with index
ifoldr :: (Int -> a -> b -> b) -> b -> Pull a -> b
ifoldr f z o@Pull {length#} = slice# o (\i !x xs -> f (I# i) x xs) z 0# length#

-- | Right-associative strict fold, with index
ifoldr' :: (Int -> a -> b -> b) -> b -> Pull a -> b
ifoldr' f z o@Pull {length#} = slice# o (\i !x !xs -> f (I# i) x xs) z 0# length#

-- | Left-associative fold
foldl :: (b -> a -> b) -> b -> Pull a -> b
foldl f z0 o@Pull {length#} = slice# o (\_ !x k z -> k (f z x)) id 0# length# z0

-- | Strict left-associative fold
foldl' :: (b -> a -> b) -> b -> Pull a -> b
foldl' f z0 o@Pull {length#} = slice# o (\_ !x k z -> k $! f z x) id 0# length# z0

-- | Left-associative fold, with index
ifoldl :: (Int -> b -> a -> b) -> b -> Pull a -> b
ifoldl f z0 o@Pull {length#} = slice# o (\i !x k z -> k (f (I# i) z x)) id 0# length# z0

-- | Strict left-associative fold, with index
ifoldl' :: (Int -> b -> a -> b) -> b -> Pull a -> b
ifoldl' f z0 o@Pull {length#} = slice# o (\i !x k z -> k $! f (I# i) z x) id 0# length# z0

folded :: Fold (Pull a) a
folded = foldring foldr

--------------------------------------------------------------------------------
-- Traversals

{-# INLINE mapM #-}

{-# INLINE imapM #-}

{-# INLINE mapM_ #-}

{-# INLINE imapM_ #-}

-- | /O(N * D)/ optimising
mapM :: (GM.PrimMonad m) => (a -> m b) -> Pull a -> m (Pull b)
mapM f = fmap (fromVector @Vector) . toVectorM . map f

-- | /O(N * D)/ optimising
imapM :: (Monad m) => (Int -> a -> m b) -> Pull a -> m (Pull b)
imapM f = fmap fromVector . V.imapM f . toVector

-- | /O(N * D)/
mapM_ :: (Monad m) => (a -> m ()) -> Pull a -> m ()
mapM_ f Pull {index#, length#} = go 0#
  where
    go i = case i <# length# of
      1# -> f (index# i) >> go (i +# 1#)
      _ -> pass

-- | /O(N * D)/
imapM_ :: (Monad m) => (Int -> a -> m ()) -> Pull a -> m ()
imapM_ f Pull {index#, length#} = go 0#
  where
    go i = case i <# length# of
      1# -> f (I# i) (index# i) >> go (i +# 1#)
      _ -> pass

{-# NOINLINE outOfBoundsError# #-}
outOfBoundsError# :: String -> Int# -> Int# -> a
outOfBoundsError# a len# i# =
  outOfBoundsError a (I# i#) (I# len#)
  where
    {-# INLINE outOfBoundsError #-}
    outOfBoundsError :: String -> Int -> Int -> a
    outOfBoundsError a i len = error (printf "Data.Vector.Pull.%s: out of bounds i=%d len=%d" a i len)

--
-- {-# INLINE isfoldr# #-}
-- isfoldr#
--   :: forall v a b
--    . (G.Vector v a)
--   => v a
--   -> (Int# -> a -> b -> b)
--   -> b
--   -> Int#
--   -> Int#
--   -> b
-- isfoldr# v = \f z start end ->
--   let go i = case i <# end of
--         1# -> f i (G.unsafeIndex v (I# i)) (go (i +# 1#))
--         _ -> z
--   in go start
--
-- {-# INLINE isfoldr2# #-}
-- isfoldr2#
--   :: forall v a b
--    . (G.Vector v a)
--   => v a
--   -> (Int# -> a -> b -> b)
--   -> b
--   -> Int#
--   -> Int#
--   -> Int#
--   -> b
-- isfoldr2# v = \f z fake0 start end ->
--   let go fake i = case fake <# end of
--         1# -> f fake (G.unsafeIndex v (I# i)) (go (fake +# 1#) (i +# 1#))
--         _ -> z
--   in go (start +# fake0) start
--
--
--
-- # at the end of a function name? That means it goes fast!!!1

--------------------------------------------------------------------------------
-- Utilities / internals

{-# INLINE max# #-}
max# :: Int# -> Int# -> Int#
max# a b = if tagToEnum# (a <# b) then b else a

{-# INLINE min# #-}
min# :: Int# -> Int# -> Int#
min# a b = if tagToEnum# (a <# b) then a else b

head :: Pull a -> Maybe a
head Pull {index#, length#} = case length# of
  0# -> Nothing
  _ -> Just (index# 0#)

last :: Pull a -> Maybe a
last Pull {index#, length#} = case length# of
  0# -> Nothing
  _ -> Just (index# (length# -# 1#))

{-# INLINE pass #-}
pass :: (Monad m) => m ()
pass = pure ()

{-# INLINE (.#) #-}
(.#) :: (b -> a) -> (Int# -> b) -> Int# -> a
(.#) f g a = f (g a)

{-# INLINE (.+#) #-}
(.+#) :: (Int# -> a) -> Int# -> Int# -> a
(.+#) f n a = f (n +# a)

{-# INLINE const# #-}
const# :: a -> Int# -> a
const# a _ = a