{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Module:      Data.Chimera.Internal
-- Copyright:   (c) 2018-2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
module Data.Chimera.Internal (
  -- * Chimera
  Chimera,
  VChimera,
  UChimera,

  -- * Construction
  tabulate,
  tabulateFix,
  tabulateFix',
  iterate,
  iterateWithIndex,
  unfoldr,
  cycle,
  fromListWithDef,
  fromVectorWithDef,
  fromInfinite,

  -- * Manipulation
  interleave,
  prependVector,

  -- * Elimination
  index,
  foldr,
  toList,
  toInfinite,

  -- * Monadic construction
  tabulateM,
  tabulateFixM,
  tabulateFixM',
  iterateM,
  iterateWithIndexM,
  unfoldrM,

  -- * Subvectors
  mapSubvectors,
  imapSubvectors,
  traverseSubvectors,
  zipWithSubvectors,
  zipWithMSubvectors,
  sliceSubvectors,
) where

import Control.Applicative
import Control.Monad.Fix
import Control.Monad.Trans.Class
import qualified Control.Monad.Trans.State.Lazy as LazyState
import Control.Monad.Zip
import Data.Bits
import Data.Coerce
import qualified Data.Foldable as F
import Data.Functor.Identity
import Data.List.Infinite (Infinite (..))
import qualified Data.List.Infinite as Inf
import qualified Data.Primitive.Array as A
import Data.Typeable
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import GHC.Exts (fromListN)
import Prelude hiding (Applicative (..), and, cycle, div, drop, foldr, fromIntegral, iterate, not, or, (*), (^))

#ifdef MIN_VERSION_mtl
import Control.Monad.Reader (MonadReader, ask, local)
#endif
#ifdef MIN_VERSION_distributive
import Data.Distributive
#ifdef MIN_VERSION_adjunctions
import qualified Data.Functor.Rep as Rep
#endif
#endif

import Data.Chimera.FromIntegral

-- | Lazy infinite streams with elements from @a@,
-- backed by a 'G.Vector' @v@ (boxed, unboxed, storable, etc.).
-- Use 'tabulate', 'tabulateFix', etc. to create a stream
-- and 'index' to access its arbitrary elements
-- in constant time.
--
-- @since 0.2.0.0
newtype Chimera v a = Chimera {forall {k} (v :: k -> *) (a :: k). Chimera v a -> Array (v a)
unChimera :: A.Array (v a)}
  deriving
    ( (forall a b. (a -> b) -> Chimera v a -> Chimera v b)
-> (forall a b. a -> Chimera v b -> Chimera v a)
-> Functor (Chimera v)
forall a b. a -> Chimera v b -> Chimera v a
forall a b. (a -> b) -> Chimera v a -> Chimera v b
forall (v :: * -> *) a b.
Functor v =>
a -> Chimera v b -> Chimera v a
forall (v :: * -> *) a b.
Functor v =>
(a -> b) -> Chimera v a -> Chimera v b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (v :: * -> *) a b.
Functor v =>
(a -> b) -> Chimera v a -> Chimera v b
fmap :: forall a b. (a -> b) -> Chimera v a -> Chimera v b
$c<$ :: forall (v :: * -> *) a b.
Functor v =>
a -> Chimera v b -> Chimera v a
<$ :: forall a b. a -> Chimera v b -> Chimera v a
Functor
      -- ^ @since 0.2.0.0
    )

-- | Streams backed by boxed vectors.
--
-- @since 0.3.0.0
type VChimera = Chimera V.Vector

-- | Streams backed by unboxed vectors.
--
-- @since 0.3.0.0
type UChimera = Chimera U.Vector

-- | 'pure' creates a constant stream.
--
-- @since 0.2.0.0
instance Applicative (Chimera V.Vector) where
  pure :: forall a. a -> Chimera Vector a
pure a
a =
    Array (Vector a) -> Chimera Vector a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (Vector a) -> Chimera Vector a)
-> Array (Vector a) -> Chimera Vector a
forall a b. (a -> b) -> a -> b
$
      Int -> [Vector a] -> Array (Vector a)
forall a. Int -> [a] -> Array a
A.arrayFromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([Vector a] -> Array (Vector a)) -> [Vector a] -> Array (Vector a)
forall a b. (a -> b) -> a -> b
$
        a -> Vector a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
a Vector a -> [Vector a] -> [Vector a]
forall a. a -> [a] -> [a]
: (Int -> Vector a) -> [Int] -> [Vector a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
k -> Int -> a -> Vector a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k) a
a) [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  <*> :: forall a b.
Chimera Vector (a -> b) -> Chimera Vector a -> Chimera Vector b
(<*>) = (Vector (a -> b) -> Vector a -> Vector b)
-> Chimera Vector (a -> b) -> Chimera Vector a -> Chimera Vector b
forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c.
(Vector u a, Vector v b, Vector w c) =>
(u a -> v b -> w c) -> Chimera u a -> Chimera v b -> Chimera w c
zipWithSubvectors Vector (a -> b) -> Vector a -> Vector b
forall a b. Vector (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>)
  liftA2 :: forall a b c.
(a -> b -> c)
-> Chimera Vector a -> Chimera Vector b -> Chimera Vector c
liftA2 a -> b -> c
f = (Vector a -> Vector b -> Vector c)
-> Chimera Vector a -> Chimera Vector b -> Chimera Vector c
forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c.
(Vector u a, Vector v b, Vector w c) =>
(u a -> v b -> w c) -> Chimera u a -> Chimera v b -> Chimera w c
zipWithSubvectors ((a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
f)

-- | @since 0.3.1.0
instance Monad (Chimera V.Vector) where
  Chimera Vector a
m >>= :: forall a b.
Chimera Vector a -> (a -> Chimera Vector b) -> Chimera Vector b
>>= a -> Chimera Vector b
f = (Word -> b) -> Chimera Vector b
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate ((Word -> b) -> Chimera Vector b)
-> (Word -> b) -> Chimera Vector b
forall a b. (a -> b) -> a -> b
$ \Word
w -> Chimera Vector b -> Word -> b
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index (a -> Chimera Vector b
f (Chimera Vector a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index Chimera Vector a
m Word
w)) Word
w

-- | @since 0.3.1.0
instance MonadFix (Chimera V.Vector) where
  mfix :: forall a. (a -> Chimera Vector a) -> Chimera Vector a
mfix = (Word -> a) -> Chimera Vector a
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate ((Word -> a) -> Chimera Vector a)
-> ((a -> Chimera Vector a) -> Word -> a)
-> (a -> Chimera Vector a)
-> Chimera Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Word -> a) -> Word -> a
forall a. (a -> Word -> a) -> Word -> a
forall (m :: * -> *) a. MonadFix m => (a -> m a) -> m a
mfix ((a -> Word -> a) -> Word -> a)
-> ((a -> Chimera Vector a) -> a -> Word -> a)
-> (a -> Chimera Vector a)
-> Word
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Chimera Vector a -> Word -> a)
-> (a -> Chimera Vector a) -> a -> Word -> a
forall a b. (a -> b) -> (a -> a) -> a -> b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Chimera Vector a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index

-- | @since 0.3.1.0
instance MonadZip (Chimera V.Vector) where
  mzip :: forall a b.
Chimera Vector a -> Chimera Vector b -> Chimera Vector (a, b)
mzip = (Vector a -> Vector b -> Vector (a, b))
-> Chimera Vector a -> Chimera Vector b -> Chimera Vector (a, b)
forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c.
(Vector u a, Vector v b, Vector w c) =>
(u a -> v b -> w c) -> Chimera u a -> Chimera v b -> Chimera w c
zipWithSubvectors Vector a -> Vector b -> Vector (a, b)
forall a b. Vector a -> Vector b -> Vector (a, b)
forall (m :: * -> *) a b. MonadZip m => m a -> m b -> m (a, b)
mzip
  mzipWith :: forall a b c.
(a -> b -> c)
-> Chimera Vector a -> Chimera Vector b -> Chimera Vector c
mzipWith = (Vector a -> Vector b -> Vector c)
-> Chimera Vector a -> Chimera Vector b -> Chimera Vector c
forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c.
(Vector u a, Vector v b, Vector w c) =>
(u a -> v b -> w c) -> Chimera u a -> Chimera v b -> Chimera w c
zipWithSubvectors ((Vector a -> Vector b -> Vector c)
 -> Chimera Vector a -> Chimera Vector b -> Chimera Vector c)
-> ((a -> b -> c) -> Vector a -> Vector b -> Vector c)
-> (a -> b -> c)
-> Chimera Vector a
-> Chimera Vector b
-> Chimera Vector c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall (m :: * -> *) a b c.
MonadZip m =>
(a -> b -> c) -> m a -> m b -> m c
mzipWith

#ifdef MIN_VERSION_mtl
-- | @since 0.3.1.0
instance MonadReader Word (Chimera V.Vector) where
  ask :: Chimera Vector Word
ask = (Word -> Word) -> Chimera Vector Word
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate Word -> Word
forall a. a -> a
id
  local :: forall a. (Word -> Word) -> Chimera Vector a -> Chimera Vector a
local = (Chimera Vector a -> (Word -> Word) -> Chimera Vector a)
-> (Word -> Word) -> Chimera Vector a -> Chimera Vector a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Chimera Vector a -> (Word -> Word) -> Chimera Vector a)
 -> (Word -> Word) -> Chimera Vector a -> Chimera Vector a)
-> (Chimera Vector a -> (Word -> Word) -> Chimera Vector a)
-> (Word -> Word)
-> Chimera Vector a
-> Chimera Vector a
forall a b. (a -> b) -> a -> b
$ ((Word -> a) -> Chimera Vector a
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate ((Word -> a) -> Chimera Vector a)
-> ((Word -> Word) -> Word -> a)
-> (Word -> Word)
-> Chimera Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) (((Word -> Word) -> Word -> a)
 -> (Word -> Word) -> Chimera Vector a)
-> (Chimera Vector a -> (Word -> Word) -> Word -> a)
-> Chimera Vector a
-> (Word -> Word)
-> Chimera Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> a) -> (Word -> Word) -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) ((Word -> a) -> (Word -> Word) -> Word -> a)
-> (Chimera Vector a -> Word -> a)
-> Chimera Vector a
-> (Word -> Word)
-> Word
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Chimera Vector a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index
#endif

#ifdef MIN_VERSION_distributive
-- | @since 0.3.1.0
instance Distributive (Chimera V.Vector) where
  distribute :: forall (f :: * -> *) a.
Functor f =>
f (Chimera Vector a) -> Chimera Vector (f a)
distribute = (Word -> f a) -> Chimera Vector (f a)
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate ((Word -> f a) -> Chimera Vector (f a))
-> (f (Chimera Vector a) -> Word -> f a)
-> f (Chimera Vector a)
-> Chimera Vector (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> f (Chimera Vector a) -> f a)
-> f (Chimera Vector a) -> Word -> f a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Chimera Vector a -> a) -> f (Chimera Vector a) -> f a
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Chimera Vector a -> a) -> f (Chimera Vector a) -> f a)
-> (Word -> Chimera Vector a -> a)
-> Word
-> f (Chimera Vector a)
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Chimera Vector a -> Word -> a) -> Word -> Chimera Vector a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Chimera Vector a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index)
  collect :: forall (f :: * -> *) a b.
Functor f =>
(a -> Chimera Vector b) -> f a -> Chimera Vector (f b)
collect a -> Chimera Vector b
f = (Word -> f b) -> Chimera Vector (f b)
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate ((Word -> f b) -> Chimera Vector (f b))
-> (f a -> Word -> f b) -> f a -> Chimera Vector (f b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> f a -> f b) -> f a -> Word -> f b
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
(<$>) ((a -> b) -> f a -> f b) -> (Word -> a -> b) -> Word -> f a -> f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Chimera Vector b -> b) -> (a -> Chimera Vector b) -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Chimera Vector b
f) ((Chimera Vector b -> b) -> a -> b)
-> (Word -> Chimera Vector b -> b) -> Word -> a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Chimera Vector b -> Word -> b) -> Word -> Chimera Vector b -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Chimera Vector b -> Word -> b
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index)

#ifdef MIN_VERSION_adjunctions
-- | @since 0.3.1.0
instance Rep.Representable (Chimera V.Vector) where
  type Rep (Chimera V.Vector) = Word
  tabulate :: forall a. (Rep (Chimera Vector) -> a) -> Chimera Vector a
tabulate = (Word -> a) -> Chimera Vector a
(Rep (Chimera Vector) -> a) -> Chimera Vector a
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate
  index :: forall a. Chimera Vector a -> Rep (Chimera Vector) -> a
index = Chimera Vector a -> Word -> a
Chimera Vector a -> Rep (Chimera Vector) -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index
#endif
#endif

bits :: Int
bits :: Int
bits = Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word
0 :: Word)

-- | Create a stream of values of a given function.
-- Once created it can be accessed via 'index' or 'toList'.
--
-- >>> ch = tabulate (^ 2) :: UChimera Word
-- >>> index ch 9
-- 81
-- >>> take 10 (toList ch)
-- [0,1,4,9,16,25,36,49,64,81]
--
-- Note that @a@ could be a function type itself,
-- so one can tabulate a function of multiple arguments
-- as a nested 'Chimera' of 'Chimera's.
--
-- @since 0.2.0.0
tabulate :: G.Vector v a => (Word -> a) -> Chimera v a
tabulate :: forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate Word -> a
f = Identity (Chimera v a) -> Chimera v a
forall a. Identity a -> a
runIdentity (Identity (Chimera v a) -> Chimera v a)
-> Identity (Chimera v a) -> Chimera v a
forall a b. (a -> b) -> a -> b
$ (Word -> Identity a) -> Identity (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(Word -> m a) -> m (Chimera v a)
tabulateM ((Word -> a) -> Word -> Identity a
forall a b. Coercible a b => a -> b
coerce Word -> a
f)
{-# INLINEABLE tabulate #-}

-- | Similar to 'V.generateM', but for raw arrays.
generateArrayM :: Monad m => Int -> (Int -> m a) -> m (A.Array a)
generateArrayM :: forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Array a)
generateArrayM Int
n Int -> m a
f = Int -> [a] -> Array a
forall a. Int -> [a] -> Array a
A.arrayFromListN Int
n ([a] -> Array a) -> m [a] -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> m a) -> [Int] -> m [a]
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 Int -> m a
f [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Monadic version of 'tabulate'.
--
-- @since 0.2.0.0
tabulateM
  :: (Monad m, G.Vector v a)
  => (Word -> m a)
  -> m (Chimera v a)
tabulateM :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(Word -> m a) -> m (Chimera v a)
tabulateM Word -> m a
f = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a) -> m (Array (v a)) -> m (Chimera v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> (Int -> m (v a)) -> m (Array (v a))
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Array a)
generateArrayM (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> m (v a)
tabulateSubVector
  where
    tabulateSubVector :: Int -> m (v a)
tabulateSubVector Int
0 = a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton (a -> v a) -> m a -> m (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word -> m a
f Word
0
    tabulateSubVector Int
i = Int -> (Int -> m a) -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> (Int -> m a) -> m (v a)
G.generateM Int
ii (\Int
j -> Word -> m a
f (Int -> Word
int2word (Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)))
      where
        ii :: Int
ii = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINEABLE tabulateM #-}
{-# SPECIALIZE tabulateM :: G.Vector v a => (Word -> Identity a) -> Identity (Chimera v a) #-}

-- | For a given @f@ create a stream of values of a recursive function 'Data.Function.fix' @f@.
-- Once created it can be accessed via 'index' or 'toList'.
--
-- For example, imagine that we want to tabulate
-- <https://en.wikipedia.org/wiki/Catalan_number Catalan numbers>:
--
-- >>> catalan n = if n == 0 then 1 else sum [ catalan i * catalan (n - 1 - i) | i <- [0 .. n - 1] ]
--
-- Can we find @catalanF@ such that @catalan@ = 'Data.Function.fix' @catalanF@?
-- Just replace all recursive calls to @catalan@ with @f@:
--
-- >>> catalanF f n = if n == 0 then 1 else sum [ f i * f (n - 1 - i) | i <- [0 .. n - 1] ]
--
-- Now we are ready to use 'tabulateFix':
--
-- >>> ch = tabulateFix catalanF :: VChimera Integer
-- >>> index ch 9
-- 4862
-- >>> take 10 (toList ch)
-- [1,1,2,5,14,42,132,429,1430,4862]
--
-- __Note__: Only recursive function calls with decreasing arguments are memoized.
-- If full memoization is desired, use 'tabulateFix'' instead.
--
-- Using unboxed \/ storable \/ primitive vectors with 'tabulateFix' is not always a win:
-- the internal memoizing routine necessarily uses boxed vectors to achieve
-- a certain degree of laziness, so converting to 'UChimera' is extra work.
-- This could pay off in a long run by reducing memory residence though.
--
-- @since 0.2.0.0
tabulateFix :: (G.Vector v a, Typeable v) => ((Word -> a) -> Word -> a) -> Chimera v a
tabulateFix :: forall (v :: * -> *) a.
(Vector v a, Typeable v) =>
((Word -> a) -> Word -> a) -> Chimera v a
tabulateFix (Word -> a) -> Word -> a
uf = Identity (Chimera v a) -> Chimera v a
forall a. Identity a -> a
runIdentity (Identity (Chimera v a) -> Chimera v a)
-> Identity (Chimera v a) -> Chimera v a
forall a b. (a -> b) -> a -> b
$ ((Word -> Identity a) -> Word -> Identity a)
-> Identity (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM (((Word -> a) -> Word -> a)
-> (Word -> Identity a) -> Word -> Identity a
forall a b. Coercible a b => a -> b
coerce (Word -> a) -> Word -> a
uf)
{-# INLINEABLE tabulateFix #-}

-- | Fully memoizing version of 'tabulateFix'.
-- This function will tabulate every recursive call,
-- but might allocate a lot of memory in doing so.
-- For example, the following piece of code calculates the
-- highest number reached by the
-- <https://en.wikipedia.org/wiki/Collatz_conjecture#Statement_of_the_problem Collatz sequence>
-- of a given number, but also allocates tens of gigabytes of memory,
-- because the Collatz sequence will spike to very high numbers.
--
-- >>> collatzF :: (Word -> Word) -> (Word -> Word)
-- >>> collatzF _ 0 = 0
-- >>> collatzF f n = if n <= 2 then 4 else n `max` f (if even n then n `quot` 2 else 3 * n + 1)
-- >>>
-- >>> maximumBy (comparing $ index $ tabulateFix' collatzF) [0..1000000]
-- ...
--
-- Using 'Data.Chimera.memoizeFix' instead fixes the problem:
--
-- >>> maximumBy (comparing $ memoizeFix collatzF) [0..1000000]
-- 56991483520
--
-- Since 'tabulateFix'' memoizes all recursive calls, even with increasing argument,
-- you most likely do not want to use it with anything else than boxed vectors ('VChimera').
--
-- @since 0.3.2.0
tabulateFix' :: (G.Vector v a, Typeable v) => ((Word -> a) -> Word -> a) -> Chimera v a
tabulateFix' :: forall (v :: * -> *) a.
(Vector v a, Typeable v) =>
((Word -> a) -> Word -> a) -> Chimera v a
tabulateFix' (Word -> a) -> Word -> a
uf = Identity (Chimera v a) -> Chimera v a
forall a. Identity a -> a
runIdentity (Identity (Chimera v a) -> Chimera v a)
-> Identity (Chimera v a) -> Chimera v a
forall a b. (a -> b) -> a -> b
$ ((Word -> Identity a) -> Word -> Identity a)
-> Identity (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM' (((Word -> a) -> Word -> a)
-> (Word -> Identity a) -> Word -> Identity a
forall a b. Coercible a b => a -> b
coerce (Word -> a) -> Word -> a
uf)
{-# INLINEABLE tabulateFix' #-}

-- | Monadic version of 'tabulateFix'.
-- There are no particular guarantees about the order of recursive calls:
-- they may be executed more than once or executed in different order.
-- That said, monadic effects must be idempotent and commutative.
--
-- @since 0.2.0.0
tabulateFixM
  :: (Monad m, G.Vector v a, Typeable v)
  => ((Word -> m a) -> Word -> m a)
  -> m (Chimera v a)
tabulateFixM :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM = Strategy -> ((Word -> m a) -> Word -> m a) -> m (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
Strategy -> ((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM_ Strategy
Downwards
{-# INLINEABLE tabulateFixM #-}
{-# SPECIALIZE tabulateFixM :: (G.Vector v a, Typeable v) => ((Word -> Identity a) -> Word -> Identity a) -> Identity (Chimera v a) #-}

-- | Monadic version of 'tabulateFix''.
--
-- @since 0.3.3.0
tabulateFixM'
  :: forall m v a
   . (Monad m, G.Vector v a, Typeable v)
  => ((Word -> m a) -> Word -> m a)
  -> m (Chimera v a)
tabulateFixM' :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM' = Strategy -> ((Word -> m a) -> Word -> m a) -> m (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
Strategy -> ((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM_ Strategy
Full
{-# INLINEABLE tabulateFixM' #-}
{-# SPECIALIZE tabulateFixM' :: (G.Vector v a, Typeable v) => ((Word -> Identity a) -> Word -> Identity a) -> Identity (Chimera v a) #-}

-- | Memoization strategy, only used by @tabulateFixM_@.
data Strategy = Full | Downwards

-- | Internal implementation for 'tabulateFixM' and 'tabulateFixM''.
tabulateFixM_
  :: forall m v a
   . (Monad m, G.Vector v a, Typeable v)
  => Strategy
  -> ((Word -> m a) -> Word -> m a)
  -> m (Chimera v a)
tabulateFixM_ :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a, Typeable v) =>
Strategy -> ((Word -> m a) -> Word -> m a) -> m (Chimera v a)
tabulateFixM_ Strategy
strat (Word -> m a) -> Word -> m a
f = m (Chimera v a)
result
  where
    result :: m (Chimera v a)
    result :: m (Chimera v a)
result = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a) -> m (Array (v a)) -> m (Chimera v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> (Int -> m (v a)) -> m (Array (v a))
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Array a)
generateArrayM (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> m (v a)
tabulateSubVector

    tabulateSubVector :: Int -> m (v a)
    tabulateSubVector :: Int -> m (v a)
tabulateSubVector Int
0 =
      a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton (a -> v a) -> m a -> m (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Strategy
strat of
        Strategy
Downwards -> ((Word -> m a) -> Word -> m a) -> Word -> m a
forall a. (a -> a) -> a
fix (Word -> m a) -> Word -> m a
f Word
0
        Strategy
Full -> (Word -> m a) -> Word -> m a
f (\Word
k -> (Chimera v a -> Word -> a) -> Word -> Chimera v a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Chimera v a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index Word
k (Chimera v a -> a) -> m (Chimera v a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Chimera v a)
result) Word
0
    tabulateSubVector Int
i = m (v a)
subResult
      where
        subResult :: m (v a)
subResult = Vector a -> v a
forall (v :: * -> *) a. (Vector v a, Typeable v) => Vector a -> v a
fromBoxedVector (Vector a -> v a) -> m (Vector a) -> m (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Vector a)
subResultBoxed
        subResultBoxed :: m (Vector a)
subResultBoxed = Int -> (Int -> m a) -> m (Vector a)
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM Int
ii (\Int
j -> (Word -> m a) -> Word -> m a
f Word -> m a
fixF (Int -> Word
int2word (Int
ii Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)))
        ii :: Int
ii = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

        fixF :: Word -> m a
        fixF :: Word -> m a
fixF Word
k
          | Word
k Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Word
int2word Int
ii =
              (Chimera v a -> Word -> a) -> Word -> Chimera v a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Chimera v a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index Word
k (Chimera v a -> a) -> m (Chimera v a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Chimera v a)
result
          | Word
k Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word
int2word Int
ii Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` Int
1 Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1 =
              (Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` (Word -> Int
word2int Word
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ii)) (Vector a -> a) -> m (Vector a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Vector a)
subResultBoxed
          | Bool
otherwise =
              case Strategy
strat of
                Strategy
Downwards -> (Word -> m a) -> Word -> m a
f Word -> m a
fixF Word
k
                Strategy
Full -> (Chimera v a -> Word -> a) -> Word -> Chimera v a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Chimera v a -> Word -> a
forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index Word
k (Chimera v a -> a) -> m (Chimera v a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Chimera v a)
result
-- It's crucial to inline into tabulateFixM and tabulateFixM'.
{-# INLINE tabulateFixM_ #-}

fromBoxedVector :: forall v a. (G.Vector v a, Typeable v) => V.Vector a -> v a
fromBoxedVector :: forall (v :: * -> *) a. (Vector v a, Typeable v) => Vector a -> v a
fromBoxedVector = case forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall (a :: * -> *) (b :: * -> *).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @V.Vector @v of
  Just Vector :~: v
Refl -> Vector a -> v a
Vector a -> Vector a
forall a. a -> a
id
  Maybe (Vector :~: v)
Nothing -> Vector a -> v a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
G.convert

-- | 'iterate' @f@ @x@ returns an infinite stream, generated by
-- repeated applications of @f@ to @x@.
--
-- It holds that 'index' ('iterate' @f@ @x@) 0 is equal to @x@.
--
-- >>> ch = iterate (+ 1) 0 :: UChimera Int
-- >>> take 10 (toList ch)
-- [0,1,2,3,4,5,6,7,8,9]
--
-- @since 0.3.0.0
iterate :: G.Vector v a => (a -> a) -> a -> Chimera v a
iterate :: forall (v :: * -> *) a. Vector v a => (a -> a) -> a -> Chimera v a
iterate a -> a
f = Identity (Chimera v a) -> Chimera v a
forall a. Identity a -> a
runIdentity (Identity (Chimera v a) -> Chimera v a)
-> (a -> Identity (Chimera v a)) -> a -> Chimera v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Identity a) -> a -> Identity (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(a -> m a) -> a -> m (Chimera v a)
iterateM ((a -> a) -> a -> Identity a
forall a b. Coercible a b => a -> b
coerce a -> a
f)

-- | Similar to 'G.iterateNM'.
iterateListNM :: forall a m. Monad m => Int -> (a -> m a) -> a -> m [a]
iterateListNM :: forall a (m :: * -> *). Monad m => Int -> (a -> m a) -> a -> m [a]
iterateListNM Int
n a -> m a
f = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 then m [a] -> a -> m [a]
forall a b. a -> b -> a
const ([a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []) else Int -> a -> m [a]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    go :: Int -> a -> m [a]
    go :: Int -> a -> m [a]
go Int
0 a
s = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a
s]
    go Int
k a
s = do
      fs <- a -> m a
f a
s
      (s :) <$> go (k - 1) fs

-- | Monadic version of 'iterate'.
--
-- @since 0.3.0.0
iterateM :: (Monad m, G.Vector v a) => (a -> m a) -> a -> m (Chimera v a)
iterateM :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(a -> m a) -> a -> m (Chimera v a)
iterateM a -> m a
f a
seed = do
  nextSeed <- a -> m a
f a
seed
  let z = a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
seed
  zs <- iterateListNM bits go (G.singleton nextSeed)
  pure $ Chimera $ fromListN (bits + 1) (z : zs)
  where
    go :: v a -> m (v a)
go v a
vec = do
      nextSeed <- a -> m a
f (v a -> a
forall (v :: * -> *) a. Vector v a => v a -> a
G.unsafeLast v a
vec)
      G.iterateNM (G.length vec `shiftL` 1) f nextSeed
{-# SPECIALIZE iterateM :: G.Vector v a => (a -> Identity a) -> a -> Identity (Chimera v a) #-}

-- | 'unfoldr' @f@ @x@ returns an infinite stream, generated by
-- repeated applications of @f@ to @x@, similar to `Data.List.unfoldr`.
--
-- >>> ch = unfoldr (\acc -> (acc * acc, acc + 1)) 0 :: UChimera Int
-- >>> take 10 (toList ch)
-- [0,1,4,9,16,25,36,49,64,81]
--
-- @since 0.3.3.0
unfoldr :: G.Vector v b => (a -> (b, a)) -> a -> Chimera v b
unfoldr :: forall (v :: * -> *) b a.
Vector v b =>
(a -> (b, a)) -> a -> Chimera v b
unfoldr a -> (b, a)
f = Identity (Chimera v b) -> Chimera v b
forall a. Identity a -> a
runIdentity (Identity (Chimera v b) -> Chimera v b)
-> (a -> Identity (Chimera v b)) -> a -> Chimera v b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Identity (b, a)) -> a -> Identity (Chimera v b)
forall (m :: * -> *) (v :: * -> *) b a.
(Monad m, Vector v b) =>
(a -> m (b, a)) -> a -> m (Chimera v b)
unfoldrM ((a -> (b, a)) -> a -> Identity (b, a)
forall a b. Coercible a b => a -> b
coerce a -> (b, a)
f)

-- | This is not quite satisfactory, see https://github.com/haskell/vector/issues/447
unfoldrExactVecNM :: forall m a b v. (Monad m, G.Vector v b) => Int -> (a -> m (b, a)) -> a -> m (v b, a)
unfoldrExactVecNM :: forall (m :: * -> *) a b (v :: * -> *).
(Monad m, Vector v b) =>
Int -> (a -> m (b, a)) -> a -> m (v b, a)
unfoldrExactVecNM Int
n a -> m (b, a)
f a
s = (StateT a m (v b, a) -> a -> m (v b, a))
-> a -> StateT a m (v b, a) -> m (v b, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT a m (v b, a) -> a -> m (v b, a)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
LazyState.evalStateT a
s (StateT a m (v b, a) -> m (v b, a))
-> StateT a m (v b, a) -> m (v b, a)
forall a b. (a -> b) -> a -> b
$ do
  vec <- Int -> StateT a m b -> StateT a m (v b)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> m a -> m (v a)
G.replicateM Int
n StateT a m b
f'
  seed <- LazyState.get
  pure (vec, seed)
  where
    f' :: LazyState.StateT a m b
    f' :: StateT a m b
f' = do
      seed <- StateT a m a
forall (m :: * -> *) s. Monad m => StateT s m s
LazyState.get
      (value, newSeed) <- lift (f seed)
      LazyState.put newSeed
      pure value

-- | Monadic version of 'unfoldr'.
--
-- @since 0.3.3.0
unfoldrM :: (Monad m, G.Vector v b) => (a -> m (b, a)) -> a -> m (Chimera v b)
unfoldrM :: forall (m :: * -> *) (v :: * -> *) b a.
(Monad m, Vector v b) =>
(a -> m (b, a)) -> a -> m (Chimera v b)
unfoldrM a -> m (b, a)
f a
seed = do
  let go :: Int -> a -> m [v b]
go Int
n a
s =
        if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bits
          then [v b] -> m [v b]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
          else do
            (vec, s') <- Int -> (a -> m (b, a)) -> a -> m (v b, a)
forall (m :: * -> *) a b (v :: * -> *).
(Monad m, Vector v b) =>
Int -> (a -> m (b, a)) -> a -> m (v b, a)
unfoldrExactVecNM (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
n) a -> m (b, a)
f a
s
            rest <- go (n + 1) s'
            pure $ vec : rest
  (z, seed') <- Int -> (a -> m (b, a)) -> a -> m (v b, a)
forall (m :: * -> *) a b (v :: * -> *).
(Monad m, Vector v b) =>
Int -> (a -> m (b, a)) -> a -> m (v b, a)
unfoldrExactVecNM Int
1 a -> m (b, a)
f a
seed
  zs <- go 0 seed'
  pure $ Chimera $ fromListN (bits + 1) (z : zs)
{-# SPECIALIZE unfoldrM :: G.Vector v b => (a -> Identity (b, a)) -> a -> Identity (Chimera v b) #-}

-- | 'iterateWithIndex' @f@ @x@ returns an infinite stream, generated by
-- applications of @f@ to a current index and previous value, starting from @x@.
--
-- It holds that 'index' ('iterateWithIndex' @f@ @x@) 0 is equal to @x@.
--
-- >>> ch = iterateWithIndex (+) 100 :: UChimera Word
-- >>> take 10 (toList ch)
-- [100,101,103,106,110,115,121,128,136,145]
--
-- @since 0.3.4.0
iterateWithIndex :: G.Vector v a => (Word -> a -> a) -> a -> Chimera v a
iterateWithIndex :: forall (v :: * -> *) a.
Vector v a =>
(Word -> a -> a) -> a -> Chimera v a
iterateWithIndex Word -> a -> a
f = Identity (Chimera v a) -> Chimera v a
forall a. Identity a -> a
runIdentity (Identity (Chimera v a) -> Chimera v a)
-> (a -> Identity (Chimera v a)) -> a -> Chimera v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> a -> Identity a) -> a -> Identity (Chimera v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(Word -> a -> m a) -> a -> m (Chimera v a)
iterateWithIndexM ((Word -> a -> a) -> Word -> a -> Identity a
forall a b. Coercible a b => a -> b
coerce Word -> a -> a
f)

iterateWithIndexExactVecNM :: forall m a v. (Monad m, G.Vector v a) => Int -> (Word -> a -> m a) -> a -> m (v a)
iterateWithIndexExactVecNM :: forall (m :: * -> *) a (v :: * -> *).
(Monad m, Vector v a) =>
Int -> (Word -> a -> m a) -> a -> m (v a)
iterateWithIndexExactVecNM Int
n Word -> a -> m a
f a
s = Int -> ((Word, a) -> m (a, (Word, a))) -> (Word, a) -> m (v a)
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
Int -> (b -> m (a, b)) -> b -> m (v a)
G.unfoldrExactNM Int
n (Word, a) -> m (a, (Word, a))
go (Int -> Word
int2word Int
n, a
s)
  where
    go :: (Word, a) -> m (a, (Word, a))
    go :: (Word, a) -> m (a, (Word, a))
go (Word
i, a
x) = do
      x' <- Word -> a -> m a
f Word
i a
x
      pure (x', (i + 1, x'))

-- | Monadic version of 'iterateWithIndex'.
--
-- @since 0.3.4.0
iterateWithIndexM :: (Monad m, G.Vector v a) => (Word -> a -> m a) -> a -> m (Chimera v a)
iterateWithIndexM :: forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
(Word -> a -> m a) -> a -> m (Chimera v a)
iterateWithIndexM Word -> a -> m a
f a
seed = do
  nextSeed <- Word -> a -> m a
f Word
1 a
seed
  let z = a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
seed
  zs <- iterateListNM bits go (G.singleton nextSeed)
  pure $ Chimera $ fromListN (bits + 1) (z : zs)
  where
    go :: v a -> m (v a)
go v a
vec =
      Int -> (Word -> a -> m a) -> a -> m (v a)
forall (m :: * -> *) a (v :: * -> *).
(Monad m, Vector v a) =>
Int -> (Word -> a -> m a) -> a -> m (v a)
iterateWithIndexExactVecNM (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Word -> a -> m a
f (v a -> a
forall (v :: * -> *) a. Vector v a => v a -> a
G.unsafeLast v a
vec)
{-# SPECIALIZE iterateWithIndexM :: G.Vector v a => (Word -> a -> Identity a) -> a -> Identity (Chimera v a) #-}

interleaveVec :: G.Vector v a => v a -> v a -> v a
interleaveVec :: forall (v :: * -> *) a. Vector v a => v a -> v a -> v a
interleaveVec v a
as v a
bs =
  Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate
    (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
as Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1)
    (\Int
n -> (if Int -> Bool
forall a. Integral a => a -> Bool
even Int
n then v a
as else v a
bs) v a -> Int -> a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
G.! (Int
n Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
1))

-- | Intertleave two streams, sourcing even elements from the first one
-- and odd elements from the second one.
--
-- >>> ch = interleave (tabulate id) (tabulate (+ 100)) :: UChimera Word
-- >>> take 10 (toList ch)
-- [0,100,1,101,2,102,3,103,4,104]
--
-- @since 0.3.3.0
interleave :: G.Vector v a => Chimera v a -> Chimera v a -> Chimera v a
interleave :: forall (v :: * -> *) a.
Vector v a =>
Chimera v a -> Chimera v a -> Chimera v a
interleave (Chimera Array (v a)
as) (Chimera Array (v a)
bs) = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a) -> Array (v a) -> Chimera v a
forall a b. (a -> b) -> a -> b
$ Int -> [v a] -> Array (v a)
forall a. Int -> [a] -> Array a
A.arrayFromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [v a]
vecs
  where
    vecs :: [v a]
vecs =
      Array (v a) -> Int -> v a
forall a. Array a -> Int -> a
A.indexArray Array (v a)
as Int
0
        v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Array (v a) -> Int -> v a
forall a. Array a -> Int -> a
A.indexArray Array (v a)
bs Int
0
        v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: (Int -> v a) -> [Int] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> v a -> v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a -> v a
interleaveVec (Array (v a) -> Int -> v a
forall a. Array a -> Int -> a
A.indexArray Array (v a)
as Int
i) (Array (v a) -> Int -> v a
forall a. Array a -> Int -> a
A.indexArray Array (v a)
bs Int
i)) [Int
1 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Index a stream in a constant time.
--
-- >>> ch = tabulate (^ 2) :: UChimera Word
-- >>> index ch 9
-- 81
--
-- @since 0.2.0.0
index :: G.Vector v a => Chimera v a -> Word -> a
index :: forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
index (Chimera Array (v a)
vs) Word
i =
  (Array (v a)
vs Array (v a) -> Int -> v a
forall a. Array a -> Int -> a
`A.indexArray` (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lz))
    v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
`G.unsafeIndex` Word -> Int
word2int (Word
i Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word -> Word
forall a. Bits a => a -> a
complement ((Word
1 Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
lz))
  where
    lz :: Int
    !lz :: Int
lz = Word -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Word
i
{-# INLINE index #-}

-- | Convert a stream to an infinite list.
--
-- >>> ch = tabulate (^ 2) :: UChimera Word
-- >>> take 10 (toList ch)
-- [0,1,4,9,16,25,36,49,64,81]
--
-- @since 0.3.0.0
toList :: G.Vector v a => Chimera v a -> [a]
toList :: forall (v :: * -> *) a. Vector v a => Chimera v a -> [a]
toList (Chimera Array (v a)
vs) = (v a -> [a]) -> Array (v a) -> [a]
forall m a. Monoid m => (a -> m) -> Array a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
G.toList Array (v a)
vs

-- | Convert a stream to a proper infinite list.
--
-- @since 0.3.4.0
toInfinite :: G.Vector v a => Chimera v a -> Infinite a
toInfinite :: forall (v :: * -> *) a. Vector v a => Chimera v a -> Infinite a
toInfinite = (a -> Infinite a -> Infinite a) -> Chimera v a -> Infinite a
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> Chimera v a -> b
foldr a -> Infinite a -> Infinite a
forall a. a -> Infinite a -> Infinite a
(:<)

-- | Right-associative fold, necessarily lazy in the accumulator.
-- Any unconditional attempt to force the accumulator even to WHNF
-- will hang the computation. E. g., the following definition isn't productive:
--
-- > import Data.List.NonEmpty (NonEmpty(..))
-- > toNonEmpty = foldr (\a (x :| xs) -> a :| x : xs) :: VChimera a -> NonEmpty a
--
-- One should use lazy patterns, e. g.,
--
-- > toNonEmpty = foldr (\a ~(x :| xs) -> a :| x : xs)
foldr :: G.Vector v a => (a -> b -> b) -> Chimera v a -> b
foldr :: forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> Chimera v a -> b
foldr a -> b -> b
f (Chimera Array (v a)
vs) = (v a -> b -> b) -> b -> Array (v a) -> b
forall a b. (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr ((b -> v a -> b) -> v a -> b -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((b -> v a -> b) -> v a -> b -> b)
-> (b -> v a -> b) -> v a -> b -> b
forall a b. (a -> b) -> a -> b
$ (a -> b -> b) -> b -> v a -> b
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> b) -> b -> v a -> b
G.foldr a -> b -> b
f) b
forall a. HasCallStack => a
undefined Array (v a)
vs

measureOff :: Int -> [a] -> Either Int ([a], [a])
measureOff :: forall a. Int -> [a] -> Either Int ([a], [a])
measureOff Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ([a], [a]) -> Either Int ([a], [a])
forall a b. b -> Either a b
Right (([a], [a]) -> Either Int ([a], [a]))
-> ([a] -> ([a], [a])) -> [a] -> Either Int ([a], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([],)
  | Bool
otherwise = Int -> [a] -> Either Int ([a], [a])
forall {t} {a}. (Eq t, Num t) => t -> [a] -> Either t ([a], [a])
go Int
n
  where
    go :: t -> [a] -> Either t ([a], [a])
go t
m [] = t -> Either t ([a], [a])
forall a b. a -> Either a b
Left t
m
    go t
1 (a
x : [a]
xs) = ([a], [a]) -> Either t ([a], [a])
forall a b. b -> Either a b
Right ([a
x], [a]
xs)
    go t
m (a
x : [a]
xs) = case t -> [a] -> Either t ([a], [a])
go (t
m t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [a]
xs of
      l :: Either t ([a], [a])
l@Left {} -> Either t ([a], [a])
l
      Right ([a]
xs', [a]
xs'') -> ([a], [a]) -> Either t ([a], [a])
forall a b. b -> Either a b
Right (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs', [a]
xs'')

measureOffVector :: G.Vector v a => Int -> v a -> Either Int (v a, v a)
measureOffVector :: forall (v :: * -> *) a.
Vector v a =>
Int -> v a -> Either Int (v a, v a)
measureOffVector Int
n v a
xs
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l = (v a, v a) -> Either Int (v a, v a)
forall a b. b -> Either a b
Right (Int -> v a -> (v a, v a)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
G.splitAt Int
n v a
xs)
  | Bool
otherwise = Int -> Either Int (v a, v a)
forall a b. a -> Either a b
Left (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
  where
    l :: Int
l = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
xs

-- | Create a stream of values from a given prefix, followed by default value
-- afterwards.
--
-- @since 0.3.3.0
fromListWithDef
  :: G.Vector v a
  => a
  -- ^ Default value
  -> [a]
  -- ^ Prefix
  -> Chimera v a
fromListWithDef :: forall (v :: * -> *) a. Vector v a => a -> [a] -> Chimera v a
fromListWithDef a
a = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a)
-> ([a] -> Array (v a)) -> [a] -> Chimera v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Item (Array (v a))] -> Array (v a)
forall l. IsList l => Int -> [Item l] -> l
fromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([v a] -> Array (v a)) -> ([a] -> [v a]) -> [a] -> Array (v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [v a]
go0
  where
    go0 :: [a] -> [v a]
go0 = \case
      [] -> a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
a v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: (Int -> v a) -> [Int] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
k -> Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k) a
a) [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      a
x : [a]
xs -> a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
x v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [v a]
go Int
0 [a]
xs

    go :: Int -> [a] -> [v a]
go Int
k [a]
xs =
      if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
bits
        then []
        else v a
v v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [v a]
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
zs
      where
        kk :: Int
kk = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k
        (v a
v, [a]
zs) =
          case Int -> [a] -> Either Int ([a], [a])
forall a. Int -> [a] -> Either Int ([a], [a])
measureOff Int
kk [a]
xs of
            Left Int
l ->
              ( if Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
kk
                  then Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate Int
kk a
a
                  else Int -> [a] -> v a
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
G.fromListN Int
kk ([a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
l a
a)
              , []
              )
            Right ([a]
ys, [a]
zs') -> (Int -> [a] -> v a
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
G.fromListN Int
kk [a]
ys, [a]
zs')

-- | Create a stream of values from a given infinite list.
--
-- @since 0.3.4.0
fromInfinite
  :: G.Vector v a
  => Infinite a
  -> Chimera v a
fromInfinite :: forall (v :: * -> *) a. Vector v a => Infinite a -> Chimera v a
fromInfinite = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a)
-> (Infinite a -> Array (v a)) -> Infinite a -> Chimera v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Item (Array (v a))] -> Array (v a)
forall l. IsList l => Int -> [Item l] -> l
fromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([v a] -> Array (v a))
-> (Infinite a -> [v a]) -> Infinite a -> Array (v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Infinite a -> [v a]
forall {v :: * -> *} {a}. Vector v a => Infinite a -> [v a]
go0
  where
    go0 :: Infinite a -> [v a]
go0 (a
x :< Infinite a
xs) = a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
x v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> Infinite a -> [v a]
forall {v :: * -> *} {a}. Vector v a => Int -> Infinite a -> [v a]
go Int
0 Infinite a
xs

    go :: Int -> Infinite a -> [v a]
go Int
k Infinite a
xs =
      if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
bits
        then []
        else Int -> [a] -> v a
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
G.fromListN Int
kk [a]
ys v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> Infinite a -> [v a]
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Infinite a
zs
      where
        kk :: Int
kk = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k
        ([a]
ys, Infinite a
zs) = Int -> Infinite a -> ([a], Infinite a)
forall a. Int -> Infinite a -> ([a], Infinite a)
Inf.splitAt Int
kk Infinite a
xs

-- | Create a stream of values from a given prefix, followed by default value
-- afterwards.
--
-- @since 0.3.3.0
fromVectorWithDef
  :: G.Vector v a
  => a
  -- ^ Default value
  -> v a
  -- ^ Prefix
  -> Chimera v a
fromVectorWithDef :: forall (v :: * -> *) a. Vector v a => a -> v a -> Chimera v a
fromVectorWithDef a
a = Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a)
-> (v a -> Array (v a)) -> v a -> Chimera v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Item (Array (v a))] -> Array (v a)
forall l. IsList l => Int -> [Item l] -> l
fromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([v a] -> Array (v a)) -> (v a -> [v a]) -> v a -> Array (v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> [v a]
go0
  where
    go0 :: v a -> [v a]
go0 v a
xs = case v a -> Maybe (a, v a)
forall (v :: * -> *) a. Vector v a => v a -> Maybe (a, v a)
G.uncons v a
xs of
      Maybe (a, v a)
Nothing -> a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
a v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: (Int -> v a) -> [Int] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
k -> Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k) a
a) [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      Just (a
y, v a
ys) -> a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
y v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> v a -> [v a]
go Int
0 v a
ys

    go :: Int -> v a -> [v a]
go Int
k v a
xs = case Int -> v a -> Either Int (v a, v a)
forall (v :: * -> *) a.
Vector v a =>
Int -> v a -> Either Int (v a, v a)
measureOffVector Int
kk v a
xs of
      Left Int
l ->
        (v a
xs v a -> v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a -> v a
G.++ Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate Int
l a
a)
          v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: (Int -> v a) -> [Int] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
n) a
a) [Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      Right (v a
ys, v a
zs) -> v a
ys v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> v a -> [v a]
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v a
zs
      where
        kk :: Int
kk = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
k

-- | Prepend a given vector to a stream of values.
--
-- @since 0.4.0.0
prependVector
  :: forall v a
   . G.Vector v a
  => v a
  -> Chimera v a
  -> Chimera v a
prependVector :: forall (v :: * -> *) a.
Vector v a =>
v a -> Chimera v a -> Chimera v a
prependVector (v a -> Maybe (a, v a)
forall (v :: * -> *) a. Vector v a => v a -> Maybe (a, v a)
G.uncons -> Maybe (a, v a)
Nothing) Chimera v a
ch = Chimera v a
ch
prependVector (v a -> Maybe (a, v a)
forall (v :: * -> *) a. Vector v a => v a -> Maybe (a, v a)
G.uncons -> Just (a
pref0, v a
pref)) (Chimera Array (v a)
as) =
  Array (v a) -> Chimera v a
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v a) -> Chimera v a) -> Array (v a) -> Chimera v a
forall a b. (a -> b) -> a -> b
$
    Int -> [Item (Array (v a))] -> Array (v a)
forall l. IsList l => Int -> [Item l] -> l
fromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([Item (Array (v a))] -> Array (v a))
-> [Item (Array (v a))] -> Array (v a)
forall a b. (a -> b) -> a -> b
$
      ([LazySlice (v a)] -> Item (Array (v a)))
-> [[LazySlice (v a)]] -> [Item (Array (v a))]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [LazySlice (v a)] -> v a
[LazySlice (v a)] -> Item (Array (v a))
forall (v :: * -> *) a. Vector v a => [LazySlice (v a)] -> v a
sliceAndConcat ([[LazySlice (v a)]] -> [Item (Array (v a))])
-> [[LazySlice (v a)]] -> [Item (Array (v a))]
forall a b. (a -> b) -> a -> b
$
        [Word -> Word -> v a -> LazySlice (v a)
forall a. Word -> Word -> a -> LazySlice a
LazySlice Word
0 Word
1 (v a -> LazySlice (v a)) -> v a -> LazySlice (v a)
forall a b. (a -> b) -> a -> b
$ a -> v a
forall (v :: * -> *) a. Vector v a => a -> v a
G.singleton a
pref0] [LazySlice (v a)] -> [[LazySlice (v a)]] -> [[LazySlice (v a)]]
forall a. a -> [a] -> [a]
: Int -> Word -> Word -> [(Word, v a)] -> [[LazySlice (v a)]]
forall t. Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
go Int
0 Word
1 Word
0 [(Word, v a)]
inputs
  where
    inputs :: [(Word, v a)]
    inputs :: [(Word, v a)]
inputs =
      (Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
pref, v a
pref)
        (Word, v a) -> [(Word, v a)] -> [(Word, v a)]
forall a. a -> [a] -> [a]
: [Word] -> [v a] -> [(Word, v a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Word
1 Word -> [Word] -> [Word]
forall a. a -> [a] -> [a]
: (Int -> Word) -> [Int] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Word
1 Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftL`) [Int
0 .. Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) (Array (v a) -> [v a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList Array (v a)
as)

    go :: Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
    go :: forall t. Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
go Int
_ Word
_ Word
_ [] = []
    go Int
n Word
need Word
off orig :: [(Word, t)]
orig@((Word
lt, t
t) : [(Word, t)]
rest)
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bits = []
      | Bool
otherwise = case Word -> Word -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
need) Word
lt of
          Ordering
LT -> [Word -> Word -> t -> LazySlice t
forall a. Word -> Word -> a -> LazySlice a
LazySlice Word
off Word
need t
t] [LazySlice t] -> [[LazySlice t]] -> [[LazySlice t]]
forall a. a -> [a] -> [a]
: Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
forall t. Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word
1 Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
need) [(Word, t)]
orig
          Ordering
EQ -> [Word -> Word -> t -> LazySlice t
forall a. Word -> Word -> a -> LazySlice a
LazySlice Word
off Word
need t
t] [LazySlice t] -> [[LazySlice t]] -> [[LazySlice t]]
forall a. a -> [a] -> [a]
: Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
forall t. Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word
1 Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`shiftL` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Word
0 [(Word, t)]
rest
          Ordering
GT -> case Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
forall t. Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
go Int
n (Word
off Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
need Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
lt) Word
0 [(Word, t)]
rest of
            [] -> [Char] -> [[LazySlice t]]
forall a. HasCallStack => [Char] -> a
error [Char]
"prependVector: the stream should not get exhausted prematurely"
            [LazySlice t]
hd : [[LazySlice t]]
tl -> (Word -> Word -> t -> LazySlice t
forall a. Word -> Word -> a -> LazySlice a
LazySlice Word
off (Word
lt Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
off) t
t LazySlice t -> [LazySlice t] -> [LazySlice t]
forall a. a -> [a] -> [a]
: [LazySlice t]
hd) [LazySlice t] -> [[LazySlice t]] -> [[LazySlice t]]
forall a. a -> [a] -> [a]
: [[LazySlice t]]
tl

data LazySlice a = LazySlice !Word !Word a

sliceAndConcat :: G.Vector v a => [LazySlice (v a)] -> v a
sliceAndConcat :: forall (v :: * -> *) a. Vector v a => [LazySlice (v a)] -> v a
sliceAndConcat =
  [v a] -> v a
forall (v :: * -> *) a. Vector v a => [v a] -> v a
G.concat
    ([v a] -> v a)
-> ([LazySlice (v a)] -> [v a]) -> [LazySlice (v a)] -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LazySlice (v a) -> v a) -> [LazySlice (v a)] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map (\(LazySlice Word
from Word
len v a
vec) -> Int -> Int -> v a -> v a
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
G.slice (Word -> Int
word2int Word
from) (Word -> Int
word2int Word
len) v a
vec)

-- | Return an infinite repetition of a given vector.
-- Throw an error on an empty vector.
--
-- >>> ch = cycle (Data.Vector.fromList [4, 2]) :: VChimera Int
-- >>> take 10 (toList ch)
-- [4,2,4,2,4,2,4,2,4,2]
--
-- @since 0.3.0.0
cycle :: G.Vector v a => v a -> Chimera v a
cycle :: forall (v :: * -> *) a. Vector v a => v a -> Chimera v a
cycle v a
vec = case Word
l of
  Word
0 -> [Char] -> Chimera v a
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Chimera.cycle: empty list"
  Word
_ -> (Word -> a) -> Chimera v a
forall (v :: * -> *) a. Vector v a => (Word -> a) -> Chimera v a
tabulate (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
vec (Int -> a) -> (Word -> Int) -> Word -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Int
word2int (Word -> Int) -> (Word -> Word) -> Word -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word -> Word -> Word
forall a. Integral a => a -> a -> a
`rem` Word
l))
  where
    l :: Word
l = Int -> Word
int2word (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
vec

-- | Map subvectors of a stream, using a given length-preserving function.
--
-- @since 0.3.0.0
mapSubvectors
  :: (G.Vector u a, G.Vector v b)
  => (u a -> v b)
  -> Chimera u a
  -> Chimera v b
mapSubvectors :: forall (u :: * -> *) a (v :: * -> *) b.
(Vector u a, Vector v b) =>
(u a -> v b) -> Chimera u a -> Chimera v b
mapSubvectors u a -> v b
f = Identity (Chimera v b) -> Chimera v b
forall a. Identity a -> a
runIdentity (Identity (Chimera v b) -> Chimera v b)
-> (Chimera u a -> Identity (Chimera v b))
-> Chimera u a
-> Chimera v b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (u a -> Identity (v b)) -> Chimera u a -> Identity (Chimera v b)
forall (u :: * -> *) a (v :: * -> *) b (m :: * -> *).
(Vector u a, Vector v b, Applicative m) =>
(u a -> m (v b)) -> Chimera u a -> m (Chimera v b)
traverseSubvectors ((u a -> v b) -> u a -> Identity (v b)
forall a b. Coercible a b => a -> b
coerce u a -> v b
f)

-- | Map subvectors of a stream, using a given length-preserving function.
-- The first argument of the function is the index of the first element of subvector
-- in the 'Chimera'.
--
-- @since 0.4.0.0
imapSubvectors
  :: (G.Vector u a, G.Vector v b)
  => (Word -> u a -> v b)
  -> Chimera u a
  -> Chimera v b
imapSubvectors :: forall (u :: * -> *) a (v :: * -> *) b.
(Vector u a, Vector v b) =>
(Word -> u a -> v b) -> Chimera u a -> Chimera v b
imapSubvectors Word -> u a -> v b
f (Chimera Array (u a)
bs) = Array (v b) -> Chimera v b
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v b) -> Chimera v b) -> Array (v b) -> Chimera v b
forall a b. (a -> b) -> a -> b
$ (Int -> u a -> v b) -> Array Int -> Array (u a) -> Array (v b)
forall a b c. (a -> b -> c) -> Array a -> Array b -> Array c
forall (m :: * -> *) a b c.
MonadZip m =>
(a -> b -> c) -> m a -> m b -> m c
mzipWith Int -> u a -> v b
safeF (Int -> [Item (Array Int)] -> Array Int
forall l. IsList l => Int -> [Item l] -> l
fromListN (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int
Item (Array Int)
0 .. Int
Item (Array Int)
bits]) Array (u a)
bs
  where
    -- Computing vector length is cheap, so let's check that @f@ preserves length.
    safeF :: Int -> u a -> v b
safeF Int
i u a
x =
      if Int
xLen Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v b -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v b
fx
        then v b
fx
        else [Char] -> v b
forall a. HasCallStack => [Char] -> a
error [Char]
"imapSubvectors: the function is not length-preserving"
      where
        xLen :: Int
xLen = u a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length u a
x
        fx :: v b
fx = Word -> u a -> v b
f (if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Word
0 else Word
1 Word -> Int -> Word
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) u a
x

-- | Traverse subvectors of a stream, using a given length-preserving function.
--
-- Be careful, because similar to 'tabulateM', only lazy monadic effects can
-- be executed in a finite time: lazy state monad is fine, but strict one is
-- not.
--
-- @since 0.3.3.0
traverseSubvectors
  :: (G.Vector u a, G.Vector v b, Applicative m)
  => (u a -> m (v b))
  -> Chimera u a
  -> m (Chimera v b)
traverseSubvectors :: forall (u :: * -> *) a (v :: * -> *) b (m :: * -> *).
(Vector u a, Vector v b, Applicative m) =>
(u a -> m (v b)) -> Chimera u a -> m (Chimera v b)
traverseSubvectors u a -> m (v b)
f (Chimera Array (u a)
bs) = Array (v b) -> Chimera v b
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (v b) -> Chimera v b) -> m (Array (v b)) -> m (Chimera v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (u a -> m (v b)) -> Array (u a) -> m (Array (v 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) -> Array a -> f (Array b)
traverse u a -> m (v b)
safeF Array (u a)
bs
  where
    -- Computing vector length is cheap, so let's check that @f@ preserves length.
    safeF :: u a -> m (v b)
safeF u a
x =
      ( \v b
fx ->
          if u a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length u a
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v b -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v b
fx
            then v b
fx
            else [Char] -> v b
forall a. HasCallStack => [Char] -> a
error [Char]
"traverseSubvectors: the function is not length-preserving"
      )
        (v b -> v b) -> m (v b) -> m (v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> u a -> m (v b)
f u a
x
{-# SPECIALIZE traverseSubvectors :: (G.Vector u a, G.Vector v b) => (u a -> Identity (v b)) -> Chimera u a -> Identity (Chimera v b) #-}

-- | Zip subvectors from two streams, using a given length-preserving function.
--
-- @since 0.3.3.0
zipWithSubvectors
  :: (G.Vector u a, G.Vector v b, G.Vector w c)
  => (u a -> v b -> w c)
  -> Chimera u a
  -> Chimera v b
  -> Chimera w c
zipWithSubvectors :: forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c.
(Vector u a, Vector v b, Vector w c) =>
(u a -> v b -> w c) -> Chimera u a -> Chimera v b -> Chimera w c
zipWithSubvectors u a -> v b -> w c
f = (Identity (Chimera w c) -> Chimera w c
forall a. Identity a -> a
runIdentity (Identity (Chimera w c) -> Chimera w c)
-> (Chimera v b -> Identity (Chimera w c))
-> Chimera v b
-> Chimera w c
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Chimera v b -> Identity (Chimera w c))
 -> Chimera v b -> Chimera w c)
-> (Chimera u a -> Chimera v b -> Identity (Chimera w c))
-> Chimera u a
-> Chimera v b
-> Chimera w c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (u a -> v b -> Identity (w c))
-> Chimera u a -> Chimera v b -> Identity (Chimera w c)
forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c
       (m :: * -> *).
(Vector u a, Vector v b, Vector w c, Applicative m) =>
(u a -> v b -> m (w c))
-> Chimera u a -> Chimera v b -> m (Chimera w c)
zipWithMSubvectors ((u a -> v b -> w c) -> u a -> v b -> Identity (w c)
forall a b. Coercible a b => a -> b
coerce u a -> v b -> w c
f)

-- | Zip subvectors from two streams, using a given monadic length-preserving function.
-- Caveats for 'tabulateM' and 'traverseSubvectors' apply.
--
-- @since 0.3.3.0
zipWithMSubvectors
  :: (G.Vector u a, G.Vector v b, G.Vector w c, Applicative m)
  => (u a -> v b -> m (w c))
  -> Chimera u a
  -> Chimera v b
  -> m (Chimera w c)
zipWithMSubvectors :: forall (u :: * -> *) a (v :: * -> *) b (w :: * -> *) c
       (m :: * -> *).
(Vector u a, Vector v b, Vector w c, Applicative m) =>
(u a -> v b -> m (w c))
-> Chimera u a -> Chimera v b -> m (Chimera w c)
zipWithMSubvectors u a -> v b -> m (w c)
f (Chimera Array (u a)
bs1) (Chimera Array (v b)
bs2) = Array (w c) -> Chimera w c
forall {k} (v :: k -> *) (a :: k). Array (v a) -> Chimera v a
Chimera (Array (w c) -> Chimera w c) -> m (Array (w c)) -> m (Chimera w c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array (m (w c)) -> m (Array (w c))
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a. Applicative f => Array (f a) -> f (Array a)
sequenceA ((u a -> v b -> m (w c))
-> Array (u a) -> Array (v b) -> Array (m (w c))
forall a b c. (a -> b -> c) -> Array a -> Array b -> Array c
forall (m :: * -> *) a b c.
MonadZip m =>
(a -> b -> c) -> m a -> m b -> m c
mzipWith u a -> v b -> m (w c)
safeF Array (u a)
bs1 Array (v b)
bs2)
  where
    -- Computing vector length is cheap, so let's check that @f@ preserves length.
    safeF :: u a -> v b -> m (w c)
safeF u a
x v b
y =
      ( \w c
fx ->
          if u a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length u a
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== w c -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length w c
fx
            then w c
fx
            else [Char] -> w c
forall a. HasCallStack => [Char] -> a
error [Char]
"traverseSubvectors: the function is not length-preserving"
      )
        (w c -> w c) -> m (w c) -> m (w c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> u a -> v b -> m (w c)
f u a
x v b
y
{-# SPECIALIZE zipWithMSubvectors :: (G.Vector u a, G.Vector v b, G.Vector w c) => (u a -> v b -> Identity (w c)) -> Chimera u a -> Chimera v b -> Identity (Chimera w c) #-}

-- | Take a slice of 'Chimera', represented as a list on consecutive subvectors.
--
-- @since 0.3.3.0
sliceSubvectors
  :: G.Vector v a
  => Int
  -- ^ How many initial elements to drop?
  -> Int
  -- ^ How many subsequent elements to take?
  -> Chimera v a
  -> [v a]
sliceSubvectors :: forall (v :: * -> *) a.
Vector v a =>
Int -> Int -> Chimera v a -> [v a]
sliceSubvectors Int
off Int
len = Int -> [v a] -> [v a]
forall {v :: * -> *} {a}. Vector v a => Int -> [v a] -> [v a]
doTake Int
len ([v a] -> [v a]) -> (Chimera v a -> [v a]) -> Chimera v a -> [v a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [v a] -> [v a]
forall {v :: * -> *} {a}. Vector v a => Int -> [v a] -> [v a]
doDrop Int
off ([v a] -> [v a]) -> (Chimera v a -> [v a]) -> Chimera v a -> [v a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (v a) -> [v a]
forall a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (Array (v a) -> [v a])
-> (Chimera v a -> Array (v a)) -> Chimera v a -> [v a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Chimera v a -> Array (v a)
forall {k} (v :: k -> *) (a :: k). Chimera v a -> Array (v a)
unChimera
  where
    doTake :: Int -> [v a] -> [v a]
doTake !Int
_ [] = []
    doTake Int
n (v a
x : [v a]
xs)
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = []
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l = v a
x v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: Int -> [v a] -> [v a]
doTake (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) [v a]
xs
      | Bool
otherwise = [Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.take Int
n v a
x]
      where
        l :: Int
l = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
x

    doDrop :: Int -> [v a] -> [v a]
doDrop !Int
_ [] = []
    doDrop Int
n (v a
x : [v a]
xs)
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = v a
x v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: [v a]
xs
      | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n = Int -> [v a] -> [v a]
doDrop (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) [v a]
xs
      | Bool
otherwise = Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.drop Int
n v a
x v a -> [v a] -> [v a]
forall a. a -> [a] -> [a]
: [v a]
xs
      where
        l :: Int
l = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
x