{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables        #-}

module Data.Record.Anon.Internal.Util.StrictArray (
    StrictArray -- opaque
    -- * Array index
  , ArrayIndex(..)
  , ZeroBasedIndex(..)
  , ReverseIndex(..)
    -- * Reads
  , (!)
    -- * Conversion
  , fromList
  , fromListN
  , fromLazy
  , toLazy
    -- * Non-monadic combinators
  , (//)
  , update
  , backpermute
  , zipWith
    -- * Monadic combinators
  , mapM
  , zipWithM
  ) where

import Prelude hiding (mapM, zipWith)

import Control.Monad (forM_)
import Control.Monad.ST
import Data.Primitive.SmallArray hiding (writeSmallArray, indexSmallArray)

import qualified Control.Monad             as Monad
import qualified Data.Foldable             as Foldable
import qualified Data.Primitive.SmallArray as SmallArray

#ifdef DEBUG
import GHC.Stack
#endif

{-------------------------------------------------------------------------------
  Definition
-------------------------------------------------------------------------------}

-- | Strict vector
--
-- Implemented as a wrapper around a 'SmallArray'.
--
-- NOTE: The operations on 'Vector' do bounds checking only if the @debug@ flag
-- is enabled.
--
-- NOTE: 'Vector' is implemented as a newtype around 'SmallArray', which in turn
-- is defined as
--
-- > data SmallArray a = SmallArray (SmallArray# a)
--
-- Furthermore, 'Canonical' is a newtype around 'Vector', which is then used in
-- 'Record' as
--
-- > data Record (f :: k -> Type) (r :: Row k) = Record {
-- >       recordCanon :: {-# UNPACK #-} !(Canonical f)
-- >     , ..
-- >     }
--
-- This means that 'Record' will have /direct/ access (no pointers) to the
-- 'SmallArray#'.
newtype StrictArray i a = WrapLazy { forall i a. StrictArray i a -> SmallArray a
unwrapLazy :: SmallArray a }
  deriving newtype (Int -> StrictArray i a -> ShowS
[StrictArray i a] -> ShowS
StrictArray i a -> String
(Int -> StrictArray i a -> ShowS)
-> (StrictArray i a -> String)
-> ([StrictArray i a] -> ShowS)
-> Show (StrictArray i a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall i a. Show a => Int -> StrictArray i a -> ShowS
forall i a. Show a => [StrictArray i a] -> ShowS
forall i a. Show a => StrictArray i a -> String
$cshowsPrec :: forall i a. Show a => Int -> StrictArray i a -> ShowS
showsPrec :: Int -> StrictArray i a -> ShowS
$cshow :: forall i a. Show a => StrictArray i a -> String
show :: StrictArray i a -> String
$cshowList :: forall i a. Show a => [StrictArray i a] -> ShowS
showList :: [StrictArray i a] -> ShowS
Show, StrictArray i a -> StrictArray i a -> Bool
(StrictArray i a -> StrictArray i a -> Bool)
-> (StrictArray i a -> StrictArray i a -> Bool)
-> Eq (StrictArray i a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall i a. Eq a => StrictArray i a -> StrictArray i a -> Bool
$c== :: forall i a. Eq a => StrictArray i a -> StrictArray i a -> Bool
== :: StrictArray i a -> StrictArray i a -> Bool
$c/= :: forall i a. Eq a => StrictArray i a -> StrictArray i a -> Bool
/= :: StrictArray i a -> StrictArray i a -> Bool
Eq, (forall m. Monoid m => StrictArray i m -> m)
-> (forall m a. Monoid m => (a -> m) -> StrictArray i a -> m)
-> (forall m a. Monoid m => (a -> m) -> StrictArray i a -> m)
-> (forall a b. (a -> b -> b) -> b -> StrictArray i a -> b)
-> (forall a b. (a -> b -> b) -> b -> StrictArray i a -> b)
-> (forall b a. (b -> a -> b) -> b -> StrictArray i a -> b)
-> (forall b a. (b -> a -> b) -> b -> StrictArray i a -> b)
-> (forall a. (a -> a -> a) -> StrictArray i a -> a)
-> (forall a. (a -> a -> a) -> StrictArray i a -> a)
-> (forall a. StrictArray i a -> [a])
-> (forall a. StrictArray i a -> Bool)
-> (forall a. StrictArray i a -> Int)
-> (forall a. Eq a => a -> StrictArray i a -> Bool)
-> (forall a. Ord a => StrictArray i a -> a)
-> (forall a. Ord a => StrictArray i a -> a)
-> (forall a. Num a => StrictArray i a -> a)
-> (forall a. Num a => StrictArray i a -> a)
-> Foldable (StrictArray i)
forall a. Eq a => a -> StrictArray i a -> Bool
forall a. Num a => StrictArray i a -> a
forall a. Ord a => StrictArray i a -> a
forall m. Monoid m => StrictArray i m -> m
forall a. StrictArray i a -> Bool
forall a. StrictArray i a -> Int
forall a. StrictArray i a -> [a]
forall a. (a -> a -> a) -> StrictArray i a -> a
forall i a. Eq a => a -> StrictArray i a -> Bool
forall i a. Num a => StrictArray i a -> a
forall i a. Ord a => StrictArray i a -> a
forall m a. Monoid m => (a -> m) -> StrictArray i a -> m
forall i m. Monoid m => StrictArray i m -> m
forall i a. StrictArray i a -> Bool
forall i a. StrictArray i a -> Int
forall i a. StrictArray i a -> [a]
forall b a. (b -> a -> b) -> b -> StrictArray i a -> b
forall a b. (a -> b -> b) -> b -> StrictArray i a -> b
forall i a. (a -> a -> a) -> StrictArray i a -> a
forall i m a. Monoid m => (a -> m) -> StrictArray i a -> m
forall i b a. (b -> a -> b) -> b -> StrictArray i a -> b
forall i a b. (a -> b -> b) -> b -> StrictArray i a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall i m. Monoid m => StrictArray i m -> m
fold :: forall m. Monoid m => StrictArray i m -> m
$cfoldMap :: forall i m a. Monoid m => (a -> m) -> StrictArray i a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> StrictArray i a -> m
$cfoldMap' :: forall i m a. Monoid m => (a -> m) -> StrictArray i a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> StrictArray i a -> m
$cfoldr :: forall i a b. (a -> b -> b) -> b -> StrictArray i a -> b
foldr :: forall a b. (a -> b -> b) -> b -> StrictArray i a -> b
$cfoldr' :: forall i a b. (a -> b -> b) -> b -> StrictArray i a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> StrictArray i a -> b
$cfoldl :: forall i b a. (b -> a -> b) -> b -> StrictArray i a -> b
foldl :: forall b a. (b -> a -> b) -> b -> StrictArray i a -> b
$cfoldl' :: forall i b a. (b -> a -> b) -> b -> StrictArray i a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> StrictArray i a -> b
$cfoldr1 :: forall i a. (a -> a -> a) -> StrictArray i a -> a
foldr1 :: forall a. (a -> a -> a) -> StrictArray i a -> a
$cfoldl1 :: forall i a. (a -> a -> a) -> StrictArray i a -> a
foldl1 :: forall a. (a -> a -> a) -> StrictArray i a -> a
$ctoList :: forall i a. StrictArray i a -> [a]
toList :: forall a. StrictArray i a -> [a]
$cnull :: forall i a. StrictArray i a -> Bool
null :: forall a. StrictArray i a -> Bool
$clength :: forall i a. StrictArray i a -> Int
length :: forall a. StrictArray i a -> Int
$celem :: forall i a. Eq a => a -> StrictArray i a -> Bool
elem :: forall a. Eq a => a -> StrictArray i a -> Bool
$cmaximum :: forall i a. Ord a => StrictArray i a -> a
maximum :: forall a. Ord a => StrictArray i a -> a
$cminimum :: forall i a. Ord a => StrictArray i a -> a
minimum :: forall a. Ord a => StrictArray i a -> a
$csum :: forall i a. Num a => StrictArray i a -> a
sum :: forall a. Num a => StrictArray i a -> a
$cproduct :: forall i a. Num a => StrictArray i a -> a
product :: forall a. Num a => StrictArray i a -> a
Foldable, NonEmpty (StrictArray i a) -> StrictArray i a
StrictArray i a -> StrictArray i a -> StrictArray i a
(StrictArray i a -> StrictArray i a -> StrictArray i a)
-> (NonEmpty (StrictArray i a) -> StrictArray i a)
-> (forall b.
    Integral b =>
    b -> StrictArray i a -> StrictArray i a)
-> Semigroup (StrictArray i a)
forall b. Integral b => b -> StrictArray i a -> StrictArray i a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
forall i a. NonEmpty (StrictArray i a) -> StrictArray i a
forall i a. StrictArray i a -> StrictArray i a -> StrictArray i a
forall i a b. Integral b => b -> StrictArray i a -> StrictArray i a
$c<> :: forall i a. StrictArray i a -> StrictArray i a -> StrictArray i a
<> :: StrictArray i a -> StrictArray i a -> StrictArray i a
$csconcat :: forall i a. NonEmpty (StrictArray i a) -> StrictArray i a
sconcat :: NonEmpty (StrictArray i a) -> StrictArray i a
$cstimes :: forall i a b. Integral b => b -> StrictArray i a -> StrictArray i a
stimes :: forall b. Integral b => b -> StrictArray i a -> StrictArray i a
Semigroup, Semigroup (StrictArray i a)
StrictArray i a
Semigroup (StrictArray i a) =>
StrictArray i a
-> (StrictArray i a -> StrictArray i a -> StrictArray i a)
-> ([StrictArray i a] -> StrictArray i a)
-> Monoid (StrictArray i a)
[StrictArray i a] -> StrictArray i a
StrictArray i a -> StrictArray i a -> StrictArray i a
forall a.
Semigroup a =>
a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall i a. Semigroup (StrictArray i a)
forall i a. StrictArray i a
forall i a. [StrictArray i a] -> StrictArray i a
forall i a. StrictArray i a -> StrictArray i a -> StrictArray i a
$cmempty :: forall i a. StrictArray i a
mempty :: StrictArray i a
$cmappend :: forall i a. StrictArray i a -> StrictArray i a -> StrictArray i a
mappend :: StrictArray i a -> StrictArray i a -> StrictArray i a
$cmconcat :: forall i a. [StrictArray i a] -> StrictArray i a
mconcat :: [StrictArray i a] -> StrictArray i a
Monoid)

{-------------------------------------------------------------------------------
  Array index
-------------------------------------------------------------------------------}

class ArrayIndex i where
  -- | Compute 0-based index from @i@, given the size of the array
  arrayIndex :: Int -> i -> Int

newtype ZeroBasedIndex = ZeroBasedIndex { ZeroBasedIndex -> Int
getZeroBasedIndex :: Int }

instance ArrayIndex ZeroBasedIndex where
  arrayIndex :: Int -> ZeroBasedIndex -> Int
arrayIndex Int
_size = ZeroBasedIndex -> Int
getZeroBasedIndex

-- | Index from the /end/ of the array
--
-- @ReverseIndex 0@ points to the final element.
newtype ReverseIndex = ReverseIndex { ReverseIndex -> Int
getReverseIndex :: Int }

instance ArrayIndex ReverseIndex where
  arrayIndex :: Int -> ReverseIndex -> Int
arrayIndex Int
size ReverseIndex
i = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- ReverseIndex -> Int
getReverseIndex ReverseIndex
i

{-------------------------------------------------------------------------------
  Reads
-------------------------------------------------------------------------------}

(!) :: ArrayIndex i => StrictArray i a -> i -> a
! :: forall i a. ArrayIndex i => StrictArray i a -> i -> a
(!) = SmallArray a -> i -> a
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray (SmallArray a -> i -> a)
-> (StrictArray i a -> SmallArray a) -> StrictArray i a -> i -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StrictArray i a -> SmallArray a
forall i a. StrictArray i a -> SmallArray a
unwrapLazy

{-------------------------------------------------------------------------------
  Conversion
-------------------------------------------------------------------------------}

fromList :: [a] -> StrictArray i a
fromList :: forall a i. [a] -> StrictArray i a
fromList [a]
as = Int -> [a] -> StrictArray i a
forall a i. Int -> [a] -> StrictArray i a
fromListN ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as) [a]
as

fromListN :: Int -> [a] -> StrictArray i a
fromListN :: forall a i. Int -> [a] -> StrictArray i a
fromListN Int
n [a]
as = SmallArray a -> StrictArray i a
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray a -> StrictArray i a)
-> SmallArray a -> StrictArray i a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
r <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
n a
forall a. HasCallStack => a
undefined
    [(Int, a)] -> ((Int, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [a]
as) (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, !a
a) ->
      SmallMutableArray s a -> ZeroBasedIndex -> a -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s a
r (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i) a
a
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s a
r

fromLazy :: forall i a. SmallArray a -> StrictArray i a
fromLazy :: forall i a. SmallArray a -> StrictArray i a
fromLazy SmallArray a
v = Int -> StrictArray i a
go Int
0
  where
    go :: Int -> StrictArray i a
    go :: Int -> StrictArray i a
go Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
v
      = let !_a :: a
_a = SmallArray a -> ZeroBasedIndex -> a
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray a
v (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i)
        in Int -> StrictArray i a
go (Int -> Int
forall a. Enum a => a -> a
succ Int
i)

      | Bool
otherwise
      = SmallArray a -> StrictArray i a
forall i a. SmallArray a -> StrictArray i a
WrapLazy SmallArray a
v

toLazy :: StrictArray i a -> SmallArray a
toLazy :: forall i a. StrictArray i a -> SmallArray a
toLazy = StrictArray i a -> SmallArray a
forall i a. StrictArray i a -> SmallArray a
unwrapLazy

{-------------------------------------------------------------------------------
  Non-monadic combinators
-------------------------------------------------------------------------------}

instance Functor (StrictArray i) where
  fmap :: forall a b. (a -> b) -> StrictArray i a -> StrictArray i b
fmap a -> b
f (WrapLazy SmallArray a
as) = SmallArray b -> StrictArray i b
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray b -> StrictArray i b)
-> SmallArray b -> StrictArray i b
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s b)) -> SmallArray b)
-> (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ do
      SmallMutableArray s b
r <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
newSize b
forall a. HasCallStack => a
undefined
      SmallArray a -> (ZeroBasedIndex -> a -> ST s ()) -> ST s ()
forall (m :: * -> *) a.
Monad m =>
SmallArray a -> (ZeroBasedIndex -> a -> m ()) -> m ()
forArrayM_ SmallArray a
as ((ZeroBasedIndex -> a -> ST s ()) -> ST s ())
-> (ZeroBasedIndex -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ZeroBasedIndex
i a
a -> SmallMutableArray s b -> ZeroBasedIndex -> b -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s b
r ZeroBasedIndex
i (b -> ST s ()) -> b -> ST s ()
forall a b. (a -> b) -> a -> b
$! a -> b
f a
a
      SmallMutableArray s b -> ST s (SmallMutableArray s b)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s b
r
    where
      newSize :: Int
      newSize :: Int
newSize = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as

(//) :: ArrayIndex i => StrictArray i a -> [(i, a)] -> StrictArray i a
// :: forall i a.
ArrayIndex i =>
StrictArray i a -> [(i, a)] -> StrictArray i a
(//) (WrapLazy SmallArray a
as) [(i, a)]
as' = SmallArray a -> StrictArray i a
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray a -> StrictArray i a)
-> SmallArray a -> StrictArray i a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
r <- SmallArray a
-> Int -> Int -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
as Int
0 Int
newSize
    [(i, a)] -> ((i, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(i, a)]
as' (((i, a) -> ST s ()) -> ST s ()) -> ((i, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(i
i, !a
a) -> SmallMutableArray s a -> i -> a -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s a
r i
i a
a
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s a
r
  where
    newSize :: Int
    newSize :: Int
newSize = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as

update ::
     ArrayIndex i
  => StrictArray i a  -- ^ Array to update
  -> [(i, a)]         -- ^ Indices into the original array and their new value
                      --   (the order of this list is irrelevant)
  -> StrictArray i a
update :: forall i a.
ArrayIndex i =>
StrictArray i a -> [(i, a)] -> StrictArray i a
update (WrapLazy SmallArray a
as) [(i, a)]
as' = SmallArray a -> StrictArray i a
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray a -> StrictArray i a)
-> SmallArray a -> StrictArray i a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
r <- SmallArray a
-> Int -> Int -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
as Int
0 Int
newSize
    [(i, a)] -> ((i, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(i, a)]
as' (((i, a) -> ST s ()) -> ST s ()) -> ((i, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(i
j, !a
a) -> SmallMutableArray s a -> i -> a -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s a
r i
j a
a
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s a
r
  where
    newSize :: Int
    newSize :: Int
newSize = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as

backpermute ::
     ArrayIndex i
  => StrictArray i a   -- ^ Array to take values from
  -> [i]               -- ^ List of indices into the source array,
                       --   in the order they must appear in the result array
  -> StrictArray i a
backpermute :: forall i a.
ArrayIndex i =>
StrictArray i a -> [i] -> StrictArray i a
backpermute (WrapLazy SmallArray a
as) [i]
is = SmallArray a -> StrictArray i a
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray a -> StrictArray i a)
-> SmallArray a -> StrictArray i a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
r <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
newSize a
forall a. HasCallStack => a
undefined
    [(Int, i)] -> ((Int, i) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [i] -> [(Int, i)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [i]
is) (((Int, i) -> ST s ()) -> ST s ())
-> ((Int, i) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, i
j) ->
      SmallMutableArray s a -> ZeroBasedIndex -> a -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s a
r (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i) (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! SmallArray a -> i -> a
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray a
as i
j
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s a
r
  where
    newSize :: Int
    newSize :: Int
newSize = [i] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [i]
is

zipWith ::
     (a -> b -> c)
  -> StrictArray i a -> StrictArray i b -> StrictArray i c
zipWith :: forall a b c i.
(a -> b -> c)
-> StrictArray i a -> StrictArray i b -> StrictArray i c
zipWith a -> b -> c
f (WrapLazy SmallArray a
as) (WrapLazy SmallArray b
bs) = SmallArray c -> StrictArray i c
forall i a. SmallArray a -> StrictArray i a
WrapLazy (SmallArray c -> StrictArray i c)
-> SmallArray c -> StrictArray i c
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s c)) -> SmallArray c
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s c)) -> SmallArray c)
-> (forall s. ST s (SmallMutableArray s c)) -> SmallArray c
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s c
r <- Int -> c -> ST s (SmallMutableArray (PrimState (ST s)) c)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
newSize c
forall a. HasCallStack => a
undefined
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
newSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      let !c :: c
c = a -> b -> c
f (SmallArray a -> ZeroBasedIndex -> a
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray a
as (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i))
                 (SmallArray b -> ZeroBasedIndex -> b
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray b
bs (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i))
      SmallMutableArray s c -> ZeroBasedIndex -> c -> ST s ()
forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s c
r (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i) c
c
    SmallMutableArray s c -> ST s (SmallMutableArray s c)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return SmallMutableArray s c
r
  where
    newSize :: Int
    newSize :: Int
newSize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as) (SmallArray b -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray b
bs)

{-------------------------------------------------------------------------------
  Applicative combinators

  NOTE: The monadic combinators here do two traversals, first collecting all
  elements of the vector in memory, and then constructing the new vector. The
  alternative is to use 'traverseSmallArrayP', but it is only sound with
  certain monads. Since this restriction would leak out to users of the library
  (through the monadic combinators on 'Record'), we prefer to avoid it.
-------------------------------------------------------------------------------}

mapM :: forall m i a b.
     Applicative m
  => (a -> m b) -> StrictArray i a -> m (StrictArray i b)
mapM :: forall (m :: * -> *) i a b.
Applicative m =>
(a -> m b) -> StrictArray i a -> m (StrictArray i b)
mapM a -> m b
f (WrapLazy SmallArray a
as) =
    Int -> [b] -> StrictArray i b
forall a i. Int -> [a] -> StrictArray i a
fromListN Int
newSize ([b] -> StrictArray i b) -> m [b] -> m (StrictArray i b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (a -> m b) -> [a] -> m [b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse a -> m b
f (SmallArray a -> [a]
forall a. SmallArray a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList SmallArray a
as)
  where
    newSize :: Int
    newSize :: Int
newSize = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as

zipWithM ::
     Applicative m
  => (a -> b -> m c)
  -> StrictArray i a -> StrictArray i b -> m (StrictArray i c)
zipWithM :: forall (m :: * -> *) a b c i.
Applicative m =>
(a -> b -> m c)
-> StrictArray i a -> StrictArray i b -> m (StrictArray i c)
zipWithM a -> b -> m c
f (WrapLazy SmallArray a
as) (WrapLazy SmallArray b
bs) = do
    Int -> [c] -> StrictArray i c
forall a i. Int -> [a] -> StrictArray i a
fromListN Int
newSize ([c] -> StrictArray i c) -> m [c] -> m (StrictArray i c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      (a -> b -> m c) -> [a] -> [b] -> m [c]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
Monad.zipWithM a -> b -> m c
f (SmallArray a -> [a]
forall a. SmallArray a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList SmallArray a
as) (SmallArray b -> [b]
forall a. SmallArray a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList SmallArray b
bs)
  where
    newSize :: Int
    newSize :: Int
newSize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
as) (SmallArray b -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray b
bs)

{-------------------------------------------------------------------------------
  Internal auxiliary
-------------------------------------------------------------------------------}

forArrayM_ :: forall m a.
     Monad m
  => SmallArray a -> (ZeroBasedIndex -> a -> m ()) -> m ()
forArrayM_ :: forall (m :: * -> *) a.
Monad m =>
SmallArray a -> (ZeroBasedIndex -> a -> m ()) -> m ()
forArrayM_ SmallArray a
arr ZeroBasedIndex -> a -> m ()
f = Int -> m ()
go Int
0
  where
    go :: Int -> m ()
    go :: Int -> m ()
go Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr = do
          ZeroBasedIndex -> a -> m ()
f (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i) (SmallArray a -> ZeroBasedIndex -> a
forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray a
arr (Int -> ZeroBasedIndex
ZeroBasedIndex Int
i))
          Int -> m ()
go (Int -> Int
forall a. Enum a => a -> a
succ Int
i)
      | Bool
otherwise =
          () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

{-------------------------------------------------------------------------------
  Interpreting 'ArrayIndex'

  Bounds checking is only enabled when built with the @debug@ flag set.
-------------------------------------------------------------------------------}

indexSmallArray :: ArrayIndex i => SmallArray r -> i -> r
indexSmallArray :: forall i r. ArrayIndex i => SmallArray r -> i -> r
indexSmallArray SmallArray r
arr i
i = SmallArray r -> Int -> r -> r
forall a r. SmallArray a -> Int -> r -> r
boundsCheck SmallArray r
arr Int
i' (r -> r) -> r -> r
forall a b. (a -> b) -> a -> b
$
    SmallArray r -> Int -> r
forall a. SmallArray a -> Int -> a
SmallArray.indexSmallArray SmallArray r
arr Int
i'
  where
    i' :: Int
    i' :: Int
i' = Int -> i -> Int
forall i. ArrayIndex i => Int -> i -> Int
arrayIndex (SmallArray r -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray r
arr) i
i

writeSmallArray :: ArrayIndex i => SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray :: forall i s a.
ArrayIndex i =>
SmallMutableArray s a -> i -> a -> ST s ()
writeSmallArray SmallMutableArray s a
arr i
i a
a = do
    Int
sz <- SmallMutableArray (PrimState (ST s)) a -> ST s Int
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m Int
getSizeofSmallMutableArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr
    let i' :: Int
i' = Int -> i -> Int
forall i. ArrayIndex i => Int -> i -> Int
arrayIndex Int
sz i
i
    SmallMutableArray s a -> Int -> ST s () -> ST s ()
forall s a r. SmallMutableArray s a -> Int -> r -> r
boundsCheckM SmallMutableArray s a
arr Int
i' (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
SmallArray.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr Int
i' a
a

#ifdef DEBUG
boundsCheck :: HasCallStack => SmallArray a -> Int -> r -> r
boundsCheck arr i k =
    if 0 <= i && i < sizeofSmallArray arr
      then k
      else error $ concat [
               "StrictArray: index " ++ show i ++ " out of bounds"
             , " (array size: " ++ show (sizeofSmallArray arr) ++ ")"
             ]
#else
boundsCheck :: SmallArray a -> Int -> r -> r
boundsCheck :: forall a r. SmallArray a -> Int -> r -> r
boundsCheck SmallArray a
_arr Int
_i r
k = r
k
#endif

#ifdef DEBUG
boundsCheckM :: HasCallStack => SmallMutableArray s a -> Int -> r -> r
boundsCheckM arr i k =
    if 0 <= i && i < sizeofSmallMutableArray arr
      then k
      else error $ concat [
               "StrictArray: index " ++ show i ++ " out of bounds"
             , " (array size: " ++ show (sizeofSmallMutableArray arr) ++ ")"
             ]
#else
boundsCheckM :: SmallMutableArray s a -> Int -> r -> r
boundsCheckM :: forall s a r. SmallMutableArray s a -> Int -> r -> r
boundsCheckM SmallMutableArray s a
_arr Int
_i r
k = r
k
#endif

{-------------------------------------------------------------------------------
  Auxiliary: support primitive < 0.9
-------------------------------------------------------------------------------}

#if !MIN_VERSION_primitive(0,9,0)
getSizeofSmallMutableArray :: SmallMutableArray s a -> ST s Int
getSizeofSmallMutableArray = return . sizeofSmallMutableArray
#endif