-- Copyright 2018-2021 Google LLC
-- Copyright 2022-2023 Andrew Pritchard
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

-- Work around <https://ghc.haskell.org/trac/ghc/ticket/14511>
{-# OPTIONS_GHC -fno-float-in #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- Make Haddock prefer to link to Data.Vec.Short rather than here, and not
-- complain about missing docs for package-internal functions.
{-# OPTIONS_HADDOCK not-home, prune #-}

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ViewPatterns #-}

-- | An implementation of short vectors.
--
-- The underlying implementation uses the 'SmallArray#' primitive,
-- which lacks the \"card marking\" of 'GHC.Exts.Array#': the upside being
-- that it avoids the overhead of maintaining the card state, the downside
-- being that the garbage collector must scan through the entire array
-- rather than just the parts marked as having changed since the last GC.
-- Using 'SmallArray#' is typically a win for arrays with fewer than 128
-- elements.

-- TODO(b/109667526): add rewrite rules, and maybe builder and view
-- interfaces along the way.
--
-- TODO(b/109668556): revisit all the inline pragmas.

module Data.Vec.Short.Internal where

import Prelude hiding ((++), concat, iterate)

#if !MIN_VERSION_base(4, 18, 0)
import Control.Applicative (Applicative(..))
#endif
import Control.DeepSeq (NFData(rnf))
import Control.Exception (assert)
import qualified Data.Data as D
import qualified Data.Foldable as F
import Data.Function (on)
import Data.Functor ((<&>))
import Data.Kind (Type)
import qualified Data.List as L (sort, sortBy, sortOn, findIndex)
import Data.Semigroup (All(..), Any(..), Sum(..), Product(..))
import GHC.Exts
         ( Int(I#), Proxy#, State#, SmallMutableArray#, SmallArray#
         , cloneSmallArray#, copySmallArray#, indexSmallArray#, newSmallArray#
         , sizeofSmallArray#, thawSmallArray#, unsafeFreezeSmallArray#
         , writeSmallArray#, proxy#, coerce
         )
import qualified GHC.Exts as GHC (IsList(..))
import GHC.Stack (HasCallStack)
import GHC.ST (ST(..), runST)
import GHC.TypeNats (Nat, KnownNat, type (+), type (*), natVal')

import Data.Default.Class (Default(..))
import Data.Distributive (Distributive(..))
import Data.Foldable.WithIndex (FoldableWithIndex(..))
import Data.Functor.Apply (Apply(..))
import Data.Functor.Bind (Bind(..))
import Data.Functor.Rep (Representable(..))
import Data.Functor.WithIndex (FunctorWithIndex)
import qualified Data.Functor.WithIndex as X (imap)
import Data.Portray (Portray(..), Portrayal(..), strAtom)
import Data.Portray.Diff (Diff(..))
import Data.Traversable.WithIndex (TraversableWithIndex(..))
import qualified Test.QuickCheck as QC

import Data.Fin.Int (Fin, finToInt, unsafeFin)
import Data.SInt (SInt(SI#, unSInt), sintVal, subSIntL, divSIntR, withSInt, addSInt)

#if !MIN_VERSION_base(4,15,0)
import GHC.Natural (naturalToInteger, naturalToInt)
import GHC.Integer (integerToInt)
#endif

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

foldrEnumFin :: SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin :: forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
sn Fin n -> a -> a
c a
n = Int -> a
go Int
0
 where
   go :: Int -> a
go Int
i
     | Int
i forall a. Eq a => a -> a -> Bool
== forall (n :: Nat). SInt n -> Int
unSInt SInt n
sn = a
n
     | Bool
otherwise = Fin n -> a -> a
c (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i) (Int -> a
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
{-# INLINE [0] foldrEnumFin #-}

forMFin_ :: Applicative f => SInt n -> (Fin n -> f a) -> f ()
forMFin_ :: forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
SInt n -> (Fin n -> f a) -> f ()
forMFin_ SInt n
n Fin n -> f a
f = forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
n (\Fin n
i f ()
rest -> Fin n -> f a
f Fin n
i forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> f ()
rest) (forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ())
{-# INLINE forMFin_ #-}

foldMapFin :: Monoid m => SInt n -> (Fin n -> (# m #)) -> m
foldMapFin :: forall m (n :: Nat). Monoid m => SInt n -> (Fin n -> (# m #)) -> m
foldMapFin SInt n
sn Fin n -> (# m #)
f = Int -> m -> m
go Int
0 forall a. Monoid a => a
mempty
 where
  go :: Int -> m -> m
go Int
i m
acc
    | Int
i forall a. Eq a => a -> a -> Bool
== forall (n :: Nat). SInt n -> Int
unSInt SInt n
sn = m
acc
    | Bool
otherwise = case Fin n -> (# m #)
f (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i) of (# m
x #) -> Int -> m -> m
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (m
acc forall a. Semigroup a => a -> a -> a
<> m
x)
{-# INLINE foldMapFin #-}

foldMFin_
  :: Monad m
  => SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ :: forall (m :: Type -> Type) (n :: Nat) a.
Monad m =>
SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ SInt n
n a -> Fin n -> m a
f a
z = forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
n (\Fin n
i a -> m ()
rest a
a -> a -> Fin n -> m a
f a
a Fin n
i forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m ()
rest) (\a
_ -> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()) a
z
{-# INLINE foldMFin_ #-}

forMFin
  :: Applicative f => SInt n -> (Fin n -> (# f a #)) -> f [a]
forMFin :: forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
SInt n -> (Fin n -> (# f a #)) -> f [a]
forMFin SInt n
n Fin n -> (# f a #)
f = forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
n
  (\Fin n
i -> case Fin n -> (# f a #)
f Fin n
i of (# f a
fx #) -> forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (:) f a
fx)
  (forall (f :: Type -> Type) a. Applicative f => a -> f a
pure [])
{-# INLINE forMFin #-}

-- | Wrap stateful primops which don't return a value.
prim_ :: (State# s -> State# s) -> ST s ()
prim_ :: forall s. (State# s -> State# s) -> ST s ()
prim_ State# s -> State# s
f = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s0 -> case State# s -> State# s
f State# s
s0 of State# s
s1 -> (# State# s
s1, () #)
{-# INLINE prim_ #-}

-- Alas, due to a confluence of problems, we cannot define the combinator:
--
-- > ($#) :: forall (a :: #) (b :: Type) (s :: Type)
-- >      .  (a -> b) -> (State# s -> (# State# s, a #)) -> ST s b
-- > ($#) f p = ST $ \s0 -> case p s0 of (# s1, x #) -> (# s1, f x #)
--
-- While we can hack around most of those problems, and get a version
-- that works for all our use cases, for both ghc-8.0.2 and ghc-8.2.1;
-- the final straw is if we want to keep the code lint clean, HLint-2.1
-- cannot parse the necessary combination of DataKinds and KindSignatures.
-- Thus, we're forced to inline this combinator everywhere we want it.


-- [Note TypeUnsafe]: If the type-unsafe primitives are used without
-- validating their implicit premises, then the 'Nat' type-indices
-- of 'Vec'\/'MutableVec' will become out of sync with the actual
-- 'sizeofSmallArray#'; which in turn invalidates all the safety and
-- correctness guarantees we assumed we could rely on those type-indices
-- to provide.

-- [Note MemoryUnsafe]: There are two sorts of memory unsafety introduced
-- by GHC's primops. First is the usual index-out-of-bounds unsafety.
-- In the functions defined here, this sort of unsafety only leaks out as
-- a result of type-safety having been violated. Second is the impurity
-- introduced by 'unsafeFreezeSmallArray#' and 'unsafeThawSmallArray#'
-- if references are used non-linearly; that is, because these primops
-- freeze\/thaw arrays in place, they allow a term which holds the
-- 'MutableVec' view to make mutations which are then visible to terms
-- holding the 'Vec' view, thereby violating the purity of Haskell.

-- Without these annotations GHC will infer that the @n@ parameter is
-- phantom, which opens the door to bugs by allowing folks to coerce it
-- to a different 'Nat'.
type role Vec nominal representational

-- | @'Vec' n a@ is an element-lazy array of @n@ values of type @a@.
--
-- This comes with a fusion framework, so many intermediate vectors are
-- optimized away, and generally the only Vecs that correspond to actual arrays
-- are those stored in data constructors, accessed multiple times, or appearing
-- as inputs or outputs of non-inlinable functions.  Additionally, some
-- operations that rely on building up a vector incrementally (as opposed to
-- computing each index independently of the others) cannot be fused; notably
-- 'fromList', 'traverse', 'iterate', and 'vscanl'; these will always construct
-- real arrays for their results.
--
-- Most operations are access-strict unless otherwise noted, which means that
-- forcing the result (usually a Vec, but possibly something else, like with
-- 'foldMap') eagerly performs all indexing and drops references to any input
-- arrays.
data Vec (n :: Nat) (a :: Type) = V# (SmallArray# a)
    -- Alas, cannot derive a 'Generic' instance:
    -- \"\"V# must not have exotic unlifted or polymorphic arguments\"\"
    -- Nor can we derive 'Data', but at least that one we can give our
    -- own instance for.

type role MutableVec nominal nominal representational
data MutableVec (s :: Type) (n :: Nat) (a :: Type)
    = MV# (SmallMutableArray# s a)

newMV :: SInt n -> a -> ST s (MutableVec s n a)
newMV :: forall (n :: Nat) a s. SInt n -> a -> ST s (MutableVec s n a)
newMV (SI# Int
n) = forall a s (n :: Nat). Int -> a -> ST s (MutableVec s n a)
unsafeNewMV Int
n
{-# INLINE newMV #-}

-- TODO(b/109668129): We should be able to replace most of the remaining
-- calls to this function with 'newMV' if we made use of the various
-- combinators in "Utils.NatMath". Would be nice to get that extra
-- type-safety, supposing it doesn't introduce significant performance
-- regressions.

-- | This function is /type-unsafe/: because it assumes the 'Int'
-- argument is in fact the reflection of @n@.
unsafeNewMV :: Int -> a -> ST s (MutableVec s n a)
unsafeNewMV :: forall a s (n :: Nat). Int -> a -> ST s (MutableVec s n a)
unsafeNewMV (I# Int#
n) a
x =
    forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case forall a d.
Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newSmallArray# Int#
n a
x State# s
s0 of { (# State# s
s1, SmallMutableArray# s a
sma #) ->
    (# State# s
s1, forall s (n :: Nat) a. SmallMutableArray# s a -> MutableVec s n a
MV# SmallMutableArray# s a
sma #) }
{-# INLINE unsafeNewMV #-}

-- Unsafe version of writeMV, using Int.
-- Each use should be vetted for being in bounds.
unsafeWriteMV :: MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV :: forall s (n :: Nat) a. MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV (MV# SmallMutableArray# s a
sma) (I# Int#
i) a
x = forall s. (State# s -> State# s) -> ST s ()
prim_ (forall d a.
SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeSmallArray# SmallMutableArray# s a
sma Int#
i a
x)
{-# INLINE unsafeWriteMV #-}

writeMV :: MutableVec s n a -> Fin n -> a -> ST s ()
writeMV :: forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i = forall s (n :: Nat) a. MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV MutableVec s n a
mv (forall (n :: Nat). Fin n -> Int
finToInt Fin n
i)
{-# INLINE writeMV #-}

-- | This function is /memory-unsafe/: because it freezes the 'MutableVec'
-- in place. See [Note MemoryUnsafe].
unsafeFreezeMV :: MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV :: forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV (MV# SmallMutableArray# s a
sma) =
    forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeSmallArray# SmallMutableArray# s a
sma State# s
s0 of { (# State# s
s1, SmallArray# a
sa #) ->
    (# State# s
s1, forall (n :: Nat) a. SmallArray# a -> Vec n a
V# SmallArray# a
sa #) }
{-# INLINE unsafeFreezeMV #-}

-- | Safely thaw a vector, by allocating a new array and copying the
-- elements over. This is both type-safe and memory-safe.
safeThawMV :: Vec n a -> ST s (MutableVec s n a)
safeThawMV :: forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV (V# SmallArray# a
sa) =
    forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case forall a d.
SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
thawSmallArray# SmallArray# a
sa Int#
0# (forall a. SmallArray# a -> Int#
sizeofSmallArray# SmallArray# a
sa) State# s
s0 of { (# State# s
s1, SmallMutableArray# s a
sma #) ->
    (# State# s
s1, forall s (n :: Nat) a. SmallMutableArray# s a -> MutableVec s n a
MV# SmallMutableArray# s a
sma #) }
{-# INLINE safeThawMV #-}

-- | This function is /type-unsafe/: because it assumes all the integers
-- are in bounds for their respective arrays. It is also /memory-unsafe/,
-- because we don't do any dynamic checks on those integers. We could
-- add such, but have avoided doing so for performance reasons.
-- See [Note TypeUnsafe] and [Note MemoryUnsafe].
--
-- TODO(b/109671227): would assertions /really/ affect the performance
-- significantly?
unsafeCopyVec :: Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec :: forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec (V# SmallArray# a
src) (I# Int#
srcOff) (MV# SmallMutableArray# s a
dst) (I# Int#
dstOff) (I# Int#
len) =
    forall s. (State# s -> State# s) -> ST s ()
prim_ (forall a d.
SmallArray# a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copySmallArray# SmallArray# a
src Int#
srcOff SmallMutableArray# s a
dst Int#
dstOff Int#
len)
{-# INLINE[1] unsafeCopyVec #-}

-- Avoid 0 length copies.
{-# RULES "unsafeCopyVec/0" forall v s m d . unsafeCopyVec v s m d 0 = return () #-}

-- | Return a known-length slice of a given vector.
--
-- Since the type is insufficiently specific to ensure memory-safety on its own
-- (because the offset argument is just 'Int'), this needs to perform runtime
-- bounds checks to ensure memory safety.
sliceVec :: Vec n a -> Int -> SInt m -> Vec m a
sliceVec :: forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec xs :: Vec n a
xs@(V# SmallArray# a
sa) off :: Int
off@(I# Int#
o) (SI# len :: Int
len@(I# Int#
l)) =
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
0 forall a. Ord a => a -> a -> Bool
<= Int
off Bool -> Bool -> Bool
&& Int
0 forall a. Ord a => a -> a -> Bool
<= Int
len Bool -> Bool -> Bool
&& Int
len forall a. Ord a => a -> a -> Bool
<= forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs forall a. Num a => a -> a -> a
- Int
off) forall a b. (a -> b) -> a -> b
$
    forall (n :: Nat) a. SmallArray# a -> Vec n a
V# (forall a. SmallArray# a -> Int# -> Int# -> SmallArray# a
cloneSmallArray# SmallArray# a
sa Int#
o Int#
l)
{-# INLINE sliceVec #-}

{-
-- If we define a @i :<: n@ type whose witnesses are isomorphic to @i@
-- itself, then we can implement these safely by rephrasing the Pi-type
-- @(i :: Fin n) -> t@ as the non-Pi @forall i. i :<: n -> t@. Otherwise
-- we can only do unsafe implementations, like the 'unsafeCopyVec' and
-- 'sliceVec' above. All the ones that use 'Min' will want to add
-- {-# OPTIONS_GHC -fplugin=GHC.TypeLits.Extra.Solver #-} to infer well.

sliceVec
    :: Vec n a
    -> (o :: Fin n)
    -> (m :: Fin (n - o))
    -> ST s (Vec m a)
sliceVec (V# sa) (finToInt -> I# off) (finToInt -> I# len) =
    V# (cloneSmallArray# sa off len)

copyVec
    :: Vec n a
    -> (srcOff :: Fin n)
    -> MutableVec s m a
    -> (dstOff :: Fin m)
    -> (len :: Fin (m - dstOff `Min` n - srcOff))
    -> ST s ()
copyVec (V# src) (finToInt -> I# srcOff) (MV# dst) (finToInt -> I# dstOff) (finToInt -> I# len) =
    prim_ (copySmallArray# src srcOff dst dstOff len)

-- And similarly for 'copySmallMutableArray#', 'cloneSmallArray#',
-- 'cloneSmallMutableArray#', 'freezeSmallArray#', 'thawSmallArray#'.
-}


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

fetch :: Vec n a -> Fin n -> (# a #)
fetch :: forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fetch (V# SmallArray# a
arr) (forall (n :: Nat). Fin n -> Int
finToInt -> I# Int#
i) = forall a. SmallArray# a -> Int# -> (# a #)
indexSmallArray# SmallArray# a
arr Int#
i
{-# INLINE fetch #-}

fusibleFetch :: Vec n a -> Fin n -> (# a #)
fusibleFetch :: forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fusibleFetch = forall (n :: Nat) a. Accessor n a -> Fin n -> (# a #)
_aIndex forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE fusibleFetch #-}

-- | Extract the given index from a 'Vec'.
--
-- This is subject to fusion if this is the only use of its input, so code like
-- @fmap f v ! i@ (which might arise due to inlining) will optimize to
-- @f (v ! i)@.
(!) :: Vec n a -> Fin n -> a
! :: forall (n :: Nat) a. Vec n a -> Fin n -> a
(!) Vec n a
xs Fin n
i = case forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fusibleFetch Vec n a
xs Fin n
i of (# a
x #) -> a
x
{-# INLINE (!) #-}

-- | Eagerly look up the value at a given position, without forcing the
-- value itself.
--
-- One of the problems with @(!)@ is that it will hold onto the underlying
-- array until the @xs!i@ expression is forced; which is a source of
-- space leaks. However, forcing the @xs!i@ expression will force
-- not only the array lookup but also the value itself; which can be
-- undesirably strict, thereby ruining the compositionality benefits
-- of laziness. The 'indexK' function is designed to overcome those
-- limitations. That is, forcing the expression @indexK xs i k@ will
-- force the array lookup and the @r@ value; thereby leaving it up to
-- @k@ to decide whether or not to force the @a@ before returning @r@.
-- So, for example, if we force @indexK xs i Just@ this will force the
-- array lookup, and wrap the unforced element in the 'Just' constructor.
{- HLINT ignore indexK "Eta reduce" -}
indexK :: Vec n a -> Fin n -> (a -> r) -> r
indexK :: forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
v Fin n
i a -> r
k = case forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fetch Vec n a
v Fin n
i of (# a
x #) -> a -> r
k a
x
{-# INLINE indexK #-}

-- | Return the size of a vector as 'SInt'.
svSize :: Vec n a -> SInt n
-- Note this strongly relies on @n@ matching the actual size of the array: if
-- it didn't, we'd be constructing an invalid 'SInt', which manifests
-- unsafety.  So, it's unsafe for a Vec to have the wrong length.
svSize :: forall (n :: Nat) a. Vec n a -> SInt n
svSize (V# SmallArray# a
sa) = forall (n :: Nat). Int -> SInt n
SI# (Int# -> Int
I# (forall a. SmallArray# a -> Int#
sizeofSmallArray# SmallArray# a
sa))
{-# INLINE svSize #-}

-- | Dynamically determine the (actual) size\/length of the vector,
-- as a standard term-level 'Int'. If you'd rather obtain @n@ at the
-- type-level, and\/or to prove that the returned value is indeed the
-- @n@ of the input, see 'svSize' and 'Data.Vec.Short.withSize' instead. If
-- you'd rather obtain @n@ statically, see 'Data.Vec.Short.vLength'.
vSize :: Vec n a -> Int
vSize :: forall (n :: Nat) a. Vec n a -> Int
vSize = forall (n :: Nat). SInt n -> Int
unSInt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> SInt n
svSize
{-# INLINE vSize #-}


--------------------------------------------------------------------------------
uninitialized :: a
uninitialized :: forall a. a
uninitialized = forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Vec: uninitialized"
{-# NOINLINE uninitialized #-}

-- Unsafe version of createVec, with Int instead of SInt.  Each use
-- should be vetted for size == valueOf @n.  Using this rather than writing out
-- 'SI# at each call site means we have a place to insert assertions more
-- easily.
unsafeCreateVec :: Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec :: forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec Int
n = forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec (forall (n :: Nat). Int -> SInt n
SI# Int
n)
{-# INLINE unsafeCreateVec #-}

createVec
  :: SInt n
  -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec :: forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
n forall s. MutableVec s n a -> ST s ()
action = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MutableVec s n a
mv <- forall (n :: Nat) a s. SInt n -> a -> ST s (MutableVec s n a)
newMV SInt n
n forall a. a
uninitialized
    forall s. MutableVec s n a -> ST s ()
action MutableVec s n a
mv
    forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s n a
mv
{-# INLINE createVec #-}

--------------------------------------------------------------------------------
-- Fusion Internals
--------------------------------------------------------------------------------

-- Fusion framework overview:
--
-- Like with list fusion, the goal is to replace actual intermediate Vec
-- objects with a non-materialized representation.  In this case, that's a
-- function for accessing vector elements by their Fin index, paired with a
-- runtime representation of the Vec length; this representation is called
-- 'Accessor'.  Also like with list fusion, we arrange to rewrite the
-- user-facing API functions to a "fusion form", which converts any input Vecs
-- into Accessors (with 'access'), implements the actual logic in terms of the
-- non-materialized representation, and converts any results into actual Vecs
-- with 'materialize'.  Then a rewrite rule deletes opposing
-- 'access'/'materialize' pairs, eliminating the intermediate allocations
-- wherever this happens.
--
-- To avoid duplicating work computing the Vec elements, we have a soft
-- requirement that a particular call to 'access' must not be used to fetch any
-- index more than once.  Violating this would mean the access/materialize rule
-- can reduce sharing.  No current functions break this rule, but we could
-- imagine adding cases where it's the client's responsibility to make sure
-- elements aren't used multiple times, like a "linear" variant of
-- 'backpermute'.
--
-- Since some Vec functions admit more efficient implementations than you'd get
-- by materializing an Accessor of their contents (e.g. implementing 'take_' by
-- 'unsafeCopyVec'), we adapt another trick from the list fusion library: keep
-- relevant operations on 'Accessor's in symbolic form for one extra simplifier
-- phase, and detect when these operations are still left alone and un-fused,
-- to rewrite them back to specialized implementations.  When we need to do
-- this, we have the specialized implementation under a different name from the
-- user-facing function, write the fusion form as the implementation of the
-- user-facing one with an INLINE pragma, and have an extra rule to bring back
-- the specialized form starting in phase 1.  When we don't need to rewrite
-- back to a specialized implementation (e.g. with 'fmap'), there's simply no
-- specialized implementation provided.
--
-- Why the difference from how list fusion does it (namely, by writing the
-- specialized implementation as the user-facing function and rewriting fusion
-- forms back to the original with phase-controlled RULES)?  I've seen GHC's
-- specialization pass do bad things with that approach, seemingly re-applying
-- rules in the wrong order to the output of specialization, resulting in using
-- element-by-element implementations instead of specialized ones.  By having
-- totally different functions for things like 'append_' and 'pureVec_', we
-- can't accidentally mess them up post-specialization with RULES.
--
-- Here's what happens in each of the GHC simplifier phases:
--
-- Phase [gentle]: GHC isn't willing to inline worker-wrapper'd function bodies
-- yet, so if we tried to use this phase to make real progress, we'd miss
-- things that didn't get exposed to RULES until phase 2.  So, we just bide our
-- time.  Some of our rules are enabled, but we don't change anything after
-- [gentle].
--
-- Phase 2: expand ops into their fusion form with RULES or INLINEs, and do the
-- actual fusion of adjacent ops with the "materialize/access" rule.  Also, do
-- some limited single-op fusion by explicitly merging the fusible form of
-- map-of-map with RULES.  This allows more cases to get rewritten back to
-- specialized implementations in phase 1.
--
-- Phase 1: detect cases that are still just a single op with a specialized
-- implementation available, and rewrite them to use it.  That is, when no
-- fusion actually happened, go back into non-fusion land.
--
-- Phase 0: inline everything and let GHC optimize the things that did get
-- subjected to nontrivial fusion.

data Accessor n a = Accessor
  { forall (n :: Nat) a. Accessor n a -> SInt n
_aSize :: SInt n
  , forall (n :: Nat) a. Accessor n a -> Fin n -> (# a #)
_aIndex :: Fin n -> (# a #)
  }

-- | Convert a 'Vec' into its size and an indexing function.
access :: Vec n a -> Accessor n a
access :: forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
v) (forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fetch Vec n a
v)
{-# INLINE CONLIKE [0] access #-}

-- | Construct an actual 'Vec' from an 'Accessor'.
--
-- Strictness properties: forcing the resulting Vec forces all of the unboxed
-- tuple accesses, so you can make Vecs that are strict in whatever you want by
-- controlling what work goes inside/outside the unboxed tuple construction.
-- Generally this is used to force all of the array accessing so that input
-- 'Vec's are no longer retained after the result is forced; but it's also used
-- to implement element-strict variants of some functions.
materialize :: Accessor n a -> Vec n a
materialize :: forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
n forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
SInt n -> (Fin n -> f a) -> f ()
forMFin_ SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
i ->
  case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i a
x
{-# INLINE [0] materialize #-}

{-# RULES

-- Fuses adjacent Vec ops, keeping everything in Accessor form.
"access/materialize" forall va. access (materialize va) = va

-- Eliminates no-op copies of a vector.
"materialize/access" forall v. materialize (access v) = v

  #-}

-- This rule seems to fire incorrectly when there's no coerce present on
-- GHC-9.2.2, leading to an infinite loop of applying the same expansion; so,
-- disable it for now.  I don't yet know why it started firing.
#if __GLASGOW_HASKELL__ < 902
{-# RULES

-- Transports coercions out from between access/materialize pairs so they can
-- fuse.
"access/coerce/materialize"
  forall v. access (coerce v) = mapVA coerce (access v)

  #-}
#endif

pureVA :: SInt n -> a -> Accessor n a
pureVA :: forall (n :: Nat) a. SInt n -> a -> Accessor n a
pureVA SInt n
n a
x = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
_ -> (# a
x #)
{-# INLINE [0] pureVA #-}

mapVA :: (a -> b) -> Accessor n a -> Accessor n b
mapVA :: forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
i -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> (# a -> b
f a
x #)
{-# INLINE [0] mapVA #-}

mapWithPosVA :: (Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA :: forall (n :: Nat) a b.
(Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$
  \Fin n
i -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> (# Fin n -> a -> b
f Fin n
i a
x #)
{-# INLINE [0] mapWithPosVA #-}

-- Make an 'Accessor' force its elements before returning them.
seqVA :: Accessor n a -> Accessor n a
seqVA :: forall (n :: Nat) a. Accessor n a -> Accessor n a
seqVA (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$
  \Fin n
i -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> a
x seq :: forall a b. a -> b -> b
`seq` (# a
x #)
{-# INLINE [0] seqVA #-}

takeVA
  :: SInt m -> Accessor (m + n) a -> Accessor m a
takeVA :: forall (m :: Nat) (n :: Nat) a.
SInt m -> Accessor (m + n) a -> Accessor m a
takeVA SInt m
m (Accessor SInt (m + n)
_ Fin (m + n) -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt m
m (\Fin m
i -> Fin (m + n) -> (# a #)
get (forall (m :: Nat) (n :: Nat). Fin m -> Fin (m + n)
embedPlus Fin m
i))
 where
  embedPlus :: Fin m -> Fin (m + n)
  embedPlus :: forall (m :: Nat) (n :: Nat). Fin m -> Fin (m + n)
embedPlus = forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). Fin n -> Int
finToInt
{-# INLINE [0] takeVA #-}

dropVA :: SInt m -> Accessor (m + n) a -> Accessor n a
dropVA :: forall (m :: Nat) (n :: Nat) a.
SInt m -> Accessor (m + n) a -> Accessor n a
dropVA SInt m
m (Accessor SInt (m + n)
mn Fin (m + n) -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor (forall (n :: Nat). Int -> SInt n
SI# (forall (n :: Nat). SInt n -> Int
unSInt SInt (m + n)
mn forall a. Num a => a -> a -> a
- forall (n :: Nat). SInt n -> Int
unSInt SInt m
m)) forall a b. (a -> b) -> a -> b
$
  \Fin n
i -> Fin (m + n) -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (forall (n :: Nat). Fin n -> Int
finToInt Fin n
i forall a. Num a => a -> a -> a
+ forall (n :: Nat). SInt n -> Int
unSInt SInt m
m))
{-# INLINE [0] dropVA #-}

revVA :: Accessor n a -> Accessor n a
revVA :: forall (n :: Nat) a. Accessor n a -> Accessor n a
revVA (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
i -> Fin n -> (# a #)
get (forall (n :: Nat). Fin n -> Fin n
complementIt Fin n
i)
 where
  !nMinus1 :: Int
nMinus1 = forall (n :: Nat). SInt n -> Int
unSInt SInt n
n forall a. Num a => a -> a -> a
- Int
1

  complementIt :: Fin n -> Fin n
  complementIt :: forall (n :: Nat). Fin n -> Fin n
complementIt = forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
nMinus1 forall a. Num a => a -> a -> a
-) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). Fin n -> Int
finToInt
{-# INLINE [0] revVA #-}

rotVA :: Fin n -> Accessor n a -> Accessor n a
rotVA :: forall (n :: Nat) a. Fin n -> Accessor n a -> Accessor n a
rotVA (forall (n :: Nat). Fin n -> Int
finToInt -> !Int
o) (Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$
  \(forall (n :: Nat). Fin n -> Int
finToInt -> !Int
i) -> Fin n -> (# a #)
get forall a b. (a -> b) -> a -> b
$ forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin forall a b. (a -> b) -> a -> b
$ if Int
i forall a. Ord a => a -> a -> Bool
>= Int
o then Int
i forall a. Num a => a -> a -> a
- Int
o else Int
nmo forall a. Num a => a -> a -> a
+ Int
i
 where
  !nmo :: Int
nmo = forall (n :: Nat). SInt n -> Int
unSInt SInt n
n forall a. Num a => a -> a -> a
- Int
o
{-# INLINE [0] rotVA #-}

liftA2VA :: (a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c
liftA2VA :: forall a b c (n :: Nat).
(a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c
liftA2VA a -> b -> c
f (Accessor SInt n
n Fin n -> (# a #)
getA) (Accessor SInt n
_ Fin n -> (# b #)
getB) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$
  \Fin n
i -> case Fin n -> (# a #)
getA Fin n
i of (# a
a #) -> case Fin n -> (# b #)
getB Fin n
i of (# b
b #) -> (# a -> b -> c
f a
a b
b #)
{-# INLINE [0] liftA2VA #-}

foldMapVA :: Monoid m => (a -> m) -> Accessor n a -> m
foldMapVA :: forall m a (n :: Nat). Monoid m => (a -> m) -> Accessor n a -> m
foldMapVA a -> m
f (Accessor SInt n
n Fin n -> (# a #)
get) =
  forall m (n :: Nat). Monoid m => SInt n -> (Fin n -> (# m #)) -> m
foldMapFin SInt n
n (\Fin n
i -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> (# a -> m
f a
x #))
{-# INLINE [0] foldMapVA #-}

sequenceVA :: Applicative f => Accessor n (f a) -> f (Vec n a)
sequenceVA :: forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
Accessor n (f a) -> f (Vec n a)
sequenceVA (Accessor SInt n
n Fin n -> (# f a #)
get) = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
n forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
SInt n -> (Fin n -> (# f a #)) -> f [a]
forMFin SInt n
n Fin n -> (# f a #)
get
{-# INLINE [0] sequenceVA #-}

-- SInt version of 'splitFin'.  Maybe I'll change the Fin library to provide
-- an SInt API at some point?
splitFinS :: SInt n -> Fin (n + m) -> Either (Fin n) (Fin m)
splitFinS :: forall (n :: Nat) (m :: Nat).
SInt n -> Fin (n + m) -> Either (Fin n) (Fin m)
splitFinS (SI# Int
n) (forall (n :: Nat). Fin n -> Int
finToInt -> Int
i)
  | Int
i forall a. Ord a => a -> a -> Bool
< Int
n     = forall a b. a -> Either a b
Left (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i)
  | Bool
otherwise = forall a b. b -> Either a b
Right (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (Int
i forall a. Num a => a -> a -> a
- Int
n))

addPosSInt :: SInt n -> SInt m -> SInt (n + m)
addPosSInt :: forall (n :: Nat) (m :: Nat). SInt n -> SInt m -> SInt (n + m)
addPosSInt (SI# Int
n) (SI# Int
m) =
  let nm :: Int
nm = Int
n forall a. Num a => a -> a -> a
+ Int
m
  in  if Int
nm forall a. Ord a => a -> a -> Bool
< Int
0
        then forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"addPosSInt: Int overflow"
        else forall (n :: Nat). Int -> SInt n
SI# (Int
n forall a. Num a => a -> a -> a
+ Int
m)
{-# INLINE addPosSInt #-}

appendVA :: Accessor n a -> Accessor m a -> Accessor (n + m) a
appendVA :: forall (n :: Nat) a (m :: Nat).
Accessor n a -> Accessor m a -> Accessor (n + m) a
appendVA (Accessor SInt n
n Fin n -> (# a #)
getN) (Accessor SInt m
m Fin m -> (# a #)
getM) = forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor
  (forall (n :: Nat) (m :: Nat). SInt n -> SInt m -> SInt (n + m)
addPosSInt SInt n
n SInt m
m)
  (\Fin (n + m)
i -> case forall (n :: Nat) (m :: Nat).
SInt n -> Fin (n + m) -> Either (Fin n) (Fin m)
splitFinS SInt n
n Fin (n + m)
i of
    Left Fin n
i' -> Fin n -> (# a #)
getN Fin n
i'
    Right Fin m
i' -> Fin m -> (# a #)
getM Fin m
i')
{-# INLINE [0] appendVA #-}

--------------------------------------------------------------------------------
-- User-facing API with fusion rules
--------------------------------------------------------------------------------

-- Unsafe version of mkVec, with Int instead of SInt.  Each use should be
-- vetted for s == valueOf @n.  Using this rather than writing out 'SI# and
-- 'unsafeFin' at each call site means we have a place to insert assertions
-- more easily.
unsafeMkVec :: Int -> (Int -> a) -> Vec n a
unsafeMkVec :: forall a (n :: Nat). Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
n Int -> a
f = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec (forall (n :: Nat). Int -> SInt n
SI# Int
n) forall a b. (a -> b) -> a -> b
$ \Fin n
i -> Int -> a
f (forall (n :: Nat). Fin n -> Int
finToInt Fin n
i)
{-# INLINE unsafeMkVec #-}

-- | Create a known-length vector using a pure function.
--
-- Note if you already have a 'Vec' of the desired length, you can use 'svSize'
-- to get the 'SInt' parameter.
tabulateVec, mkVec :: SInt n -> (Fin n -> a) -> Vec n a
tabulateVec :: forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec SInt n
n Fin n -> a
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
i -> (# Fin n -> a
f Fin n
i #)
mkVec :: forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec
{-# INLINE tabulateVec #-}
{-# INLINE mkVec #-}

-- | Element-strict form of 'mkVec': elements are forced when forcing the Vec.
tabulateVec', mkVec' :: SInt n -> (Fin n -> a) -> Vec n a
tabulateVec' :: forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec' SInt n
n Fin n -> a
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall a b. (a -> b) -> a -> b
$
  forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n forall a b. (a -> b) -> a -> b
$ \Fin n
i -> let x :: a
x = Fin n -> a
f Fin n
i in a
x seq :: forall a b. a -> b -> b
`seq` (# a
x #)
mkVec' :: forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec' = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec'
{-# INLINE tabulateVec' #-}
{-# INLINE mkVec' #-}

-- | Construct a vector by choosing an element of another vector for each index.
--
-- @
--     backpermute n f v ! i === v ! f i
-- @
backpermute :: SInt m -> (Fin m -> Fin n) -> Vec n a -> Vec m a
-- Take care: backpermute can reference the same index of the input vector
-- multiple times, so if we subjected the input side to fusion, we'd
-- potentially duplicate work.  It might make sense to make a variant of
-- 'backpermute' that assumes either indices are not duplicated or the
-- computation behind each value is cheap enough to duplicate, but we can't
-- just blindly fuse things into all 'backpermute's.
backpermute :: forall (m :: Nat) (n :: Nat) a.
SInt m -> (Fin m -> Fin n) -> Vec n a -> Vec m a
backpermute SInt m
m Fin m -> Fin n
f Vec n a
xs = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt m
m forall a b. (a -> b) -> a -> b
$ \Fin m
i -> forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
fetch Vec n a
xs (Fin m -> Fin n
f Fin m
i)
{-# INLINE backpermute #-}

--------------------------------------------------------------------------------
-- N.B., since the same @KnownNat n@ instance is passed to both 'createVecP'
-- and 'enumFinP', this will be memory-safe even if the @KnownNat n@
-- instance is illegitimate (e.g., by using 'blackMagic' unsafely).
-- An illegitimate 'KnownNat' instance would only compromise type-safety:
-- since it'd mean that the actual length of the resulting 'SmallArray#'
-- differs from the @n@ the 'Vec' claims it has.

-- | Create a vector of the specified length from a list. If @n < length xs@
-- then the suffix of the vector will be filled with 'uninitialized'
-- values. If @n > length xs@ then the suffix of @xs@ won't be included
-- in the vector. Either way, this function is both type-safe and memory-safe.
listVec :: SInt n -> [a] -> Vec n a
listVec :: forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
n [a]
xs = forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
n forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> (forall a b. (a -> b) -> a -> b
$ [a]
xs) forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
n
  (\Fin n
i [a] -> ST s ()
rest [a]
xs' -> case [a]
xs' of
      [] -> forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i forall a. a
uninitialized forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> [a] -> ST s ()
rest []
      (a
x:[a]
xs'') -> forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i a
x forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> [a] -> ST s ()
rest [a]
xs'')
  (\[a]
_ -> forall (m :: Type -> Type) a. Monad m => a -> m a
return ())
{-# INLINE listVec #-}


-- | Convert a list to a vector of the same length.
withVec :: [a] -> (forall n. Vec n a -> r) -> r
withVec :: forall a r. [a] -> (forall (n :: Nat). Vec n a -> r) -> r
withVec [a]
xs forall (n :: Nat). Vec n a -> r
f = forall r.
(?callStack::CallStack) =>
Int -> (forall (n :: Nat). SInt n -> r) -> r
withSInt (forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
xs) forall a b. (a -> b) -> a -> b
$ \SInt n
sn -> forall (n :: Nat). Vec n a -> r
f (forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn [a]
xs)
{-# INLINABLE withVec #-}

-- | Convert a list to a vector, given a hint for the length of the list.
-- If the hint does not match the actual length of the list, then the
-- behavior of this function is left unspecified. If the hint does not
-- match the desired @n@, then we throw an error just like 'fromList'.
-- For a non-errorful version, see 'withVec' instead.
fromListN :: HasCallStack => SInt n -> Int -> [a] -> Vec n a
fromListN :: forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> Int -> [a] -> Vec n a
fromListN SInt n
sn Int
l [a]
xs
    | Int
l forall a. Eq a => a -> a -> Bool
== Int
n    = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn [a]
xs
    | Bool
otherwise = forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Vec.fromListN: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Int
l forall a. Semigroup a => a -> a -> a
<> [Char]
" /= " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Int
n
    where
    !n :: Int
n = forall (n :: Nat). SInt n -> Int
unSInt SInt n
sn
{-# INLINABLE fromListN #-}


-- | Convert a list to a vector, throwing an error if the list has the
-- wrong length.
-- Note: Because this walks @xs@ to check its length, this cannot be
-- used with the list fusion optimization rules.
fromList :: HasCallStack => SInt n -> [a] -> Vec n a
fromList :: forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList SInt n
sn [a]
xs
    | Int
n forall a. Int -> [a] -> Bool
`eqLength` [a]
xs = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn [a]
xs
    | Bool
otherwise       = forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Vec.fromList: length /= " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Int
n
    where
    !n :: Int
n = forall (n :: Nat). SInt n -> Int
unSInt SInt n
sn
{-# INLINABLE fromList #-}


-- TODO(b/109757264): move this out of library into @ListUtils.hs@
-- | An implementation of @n == length xs@ which short-circuits
-- once it can determine the answer, rather than necessarily recursing
-- through the entire list to compute its length.
eqLength :: Int -> [a] -> Bool
eqLength :: forall a. Int -> [a] -> Bool
eqLength Int
0 []     = Bool
True
eqLength Int
0 (a
_:[a]
_)  = Bool
False -- too long
eqLength Int
_ []     = Bool
False -- too short
eqLength Int
n (a
_:[a]
xs) = forall a. Int -> [a] -> Bool
eqLength (Int
n forall a. Num a => a -> a -> a
- Int
1) [a]
xs


-- To support -XOverloadedLists
instance KnownNat n => GHC.IsList (Vec n a) where
    type Item (Vec n a) = a
    fromListN :: Int -> [Item (Vec n a)] -> Vec n a
fromListN = forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> Int -> [a] -> Vec n a
fromListN forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal
    fromList :: [Item (Vec n a)] -> Vec n a
fromList  = forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal  -- Not subject to list fusion optimizations.
    toList :: Vec n a -> [Item (Vec n a)]
toList    = forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList

--------------------------------------------------------------------------------
-- | Claim that 'vecConstr' is the only data-constructor of 'Vec'.
vecDataType :: D.DataType
vecDataType :: DataType
vecDataType = [Char] -> [Constr] -> DataType
D.mkDataType [Char]
"Vec.Vec" [Constr
vecConstr]

-- | Treat the 'fromList' function as a data-constructor for 'Vec'.
vecConstr :: D.Constr
vecConstr :: Constr
vecConstr = DataType -> [Char] -> [[Char]] -> Fixity -> Constr
D.mkConstr DataType
vecDataType [Char]
"fromList" [] Fixity
D.Prefix

-- The 'KnownNat' constraint is necessary for 'fromList'.
instance (KnownNat n, D.Data a) => D.Data (Vec n a) where
    toConstr :: Vec n a -> Constr
toConstr   Vec n a
_ = Constr
vecConstr
    dataTypeOf :: Vec n a -> DataType
dataTypeOf Vec n a
_ = DataType
vecDataType
    gfoldl :: forall (c :: Type -> Type).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Vec n a -> c (Vec n a)
gfoldl  forall d b. Data d => c (d -> b) -> d -> c b
app forall g. g -> c g
pur Vec n a
v = forall g. g -> c g
pur (forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal) forall d b. Data d => c (d -> b) -> d -> c b
`app` forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList Vec n a
v
    gunfold :: forall (c :: Type -> Type).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Vec n a)
gunfold forall b r. Data b => c (b -> r) -> c r
app forall r. r -> c r
pur Constr
c
        | Constr -> Int
D.constrIndex Constr
c forall a. Eq a => a -> a -> Bool
== Int
1 = forall b r. Data b => c (b -> r) -> c r
app (forall r. r -> c r
pur (forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal))
        | Bool
otherwise            = forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"gunfold@Vec: invalid constrIndex"

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

instance Show a => Show (Vec n a) where
    showsPrec :: Int -> Vec n a -> ShowS
showsPrec Int
p Vec n a
xs = Bool -> ShowS -> ShowS
showParen (Int
p forall a. Ord a => a -> a -> Bool
>= Int
precedence)
        forall a b. (a -> b) -> a -> b
$ [Char] -> ShowS
showString [Char]
"fromListN "
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows (forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> ShowS
showString [Char]
" "
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows (forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList Vec n a
xs)

instance (KnownNat n, Read a) => Read (Vec n a) where
    readsPrec :: Int -> ReadS (Vec n a)
readsPrec Int
p = forall a. Bool -> ReadS a -> ReadS a
readParen (Int
p forall a. Ord a => a -> a -> Bool
>= Int
precedence) forall a b. (a -> b) -> a -> b
$ \[Char]
s ->
        [ forall b. Int -> b -> b
assertSize (forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
ls) (forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> Int -> [a] -> Vec n a
fromListN SInt n
n Int
m [a]
ls, [Char]
s''')
        | ([Char]
"fromListN", [Char]
s') <- ReadS [Char]
lex [Char]
s
        , (Int
m, [Char]
s'') <- forall a. Read a => Int -> ReadS a
readsPrec Int
precedence [Char]
s'
        , ([a]
ls, [Char]
s''') <- forall b. Int -> b -> b
assertSize Int
m forall a. Read a => Int -> ReadS a
readsPrec Int
precedence [Char]
s''
        ]
        where
            n :: SInt n
n = forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal @n

            assertSize :: Int -> b -> b
            assertSize :: forall b. Int -> b -> b
assertSize Int
m b
x = if Int
m forall a. Eq a => a -> a -> Bool
/= forall (n :: Nat). SInt n -> Int
unSInt SInt n
n
                then forall a. (?callStack::CallStack) => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Can't read a Vec with " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Int
m forall a. Semigroup a => a -> a -> a
<>
                    [Char]
" elements into a type `Vec " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show SInt n
n forall a. Semigroup a => a -> a -> a
<>
                    [Char]
"`"
                else b
x

-- Vec is being shown as a function application which has precedence 10. Thus,
-- if we are already in function application context or one that binds even
-- tightlier (i.e. with higher precedence) we need to wrap the expression in
-- parantheses.
precedence :: Int
precedence :: Int
precedence = Int
10

instance Portray a => Portray (Vec n a) where
  portray :: Vec n a -> Portrayal
portray Vec n a
xs = Portrayal -> [Portrayal] -> Portrayal
Apply ([Char] -> Portrayal
strAtom [Char]
"fromListN")
    [forall a. Portray a => a -> Portrayal
portray (forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs), forall a. Portray a => a -> Portrayal
portray forall a b. (a -> b) -> a -> b
$ forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList Vec n a
xs]

instance (Portray a, Diff a) => Diff (Vec n a) where
  diff :: Vec n a -> Vec n a -> Maybe Portrayal
diff Vec n a
x Vec n a
y = (forall a. Diff a => a -> a -> Maybe Portrayal
diff forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList) Vec n a
x Vec n a
y forall (f :: Type -> Type) a b. Functor f => f a -> (a -> b) -> f b
<&>
    \Portrayal
d -> Portrayal -> [Portrayal] -> Portrayal
Apply ([Char] -> Portrayal
strAtom [Char]
"fromListN") [forall a. Portray a => a -> Portrayal
portray (forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
x), Portrayal
d]

instance NFData a => NFData (Vec n a) where
    rnf :: Vec n a -> ()
rnf !Vec n a
xs = forall m (n :: Nat). Monoid m => SInt n -> (Fin n -> (# m #)) -> m
foldMapFin (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) forall a b. (a -> b) -> a -> b
$
        \Fin n
i -> case forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i forall a. NFData a => a -> ()
rnf of () -> (# () #)
    {-# INLINE rnf #-}

-- | Point-wise @(<>)@.
instance Semigroup a => Semigroup (Vec n a) where
    <> :: Vec n a -> Vec n a -> Vec n a
(<>) = forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 forall a. Semigroup a => a -> a -> a
(<>)

-- | Point-wise @mempty@.
instance (KnownNat n, Monoid a) => Monoid (Vec n a) where
    mempty :: Vec n a
mempty = forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty

instance forall a n. (QC.Arbitrary a, KnownNat n) => QC.Arbitrary (Vec n a)
    where
    -- While we could get rid of the intermediate list by digging into
    -- how 'Gen' works under the hood, the benefit doesn't seem worth it.
    arbitrary :: Gen (Vec n a)
arbitrary = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Int -> Gen a -> Gen [a]
QC.vectorOf Int
n forall a. Arbitrary a => Gen a
QC.arbitrary
      where
        !sn :: SInt n
sn@(SI# Int
n) = forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal

    -- If @a@ admits too many ways to shrink, we might prefer to
    -- interleave the @shrink(xs!i)@ lists, rather than concatenating
    -- them as list comprehension will.
    shrink :: Vec n a -> [Vec n a]
shrink Vec n a
xs =
        [ forall (n :: Nat) a. Fin n -> Vec n a -> a -> Vec n a
upd Fin n
i Vec n a
xs a
x'
        | Fin n
i <- forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
sn (:) [], a
x' <- forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i forall a. Arbitrary a => a -> [a]
QC.shrink
        ]
      where
        !sn :: SInt n
sn = forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs

-- | Safely construct a new vector that differs only in one element.
-- This is inefficient, and only intended for internal use.
upd :: Fin n -> Vec n a -> a -> Vec n a
upd :: forall (n :: Nat) a. Fin n -> Vec n a -> a -> Vec n a
upd Fin n
i Vec n a
xs a
x = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    MutableVec s n a
mv <- forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV Vec n a
xs
    forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i a
x
    forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s n a
mv
{-# INLINE upd #-}

instance (Show a) => QC.CoArbitrary (Vec n a) where
    coarbitrary :: forall b. Vec n a -> Gen b -> Gen b
coarbitrary = forall a b. Show a => a -> Gen b -> Gen b
QC.coarbitraryShow

instance (KnownNat n, Num a) => Num (Vec n a) where
    + :: Vec n a -> Vec n a -> Vec n a
(+) = forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Num a => a -> a -> a
(+)
    * :: Vec n a -> Vec n a -> Vec n a
(*) = forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Num a => a -> a -> a
(*)
    (-) = forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (-)
    abs :: Vec n a -> Vec n a
abs = forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
abs
    signum :: Vec n a -> Vec n a
signum = forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
signum
    negate :: Vec n a -> Vec n a
negate = forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
negate
    fromInteger :: Integer -> Vec n a
fromInteger = forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => Integer -> a
fromInteger

instance (KnownNat n, Default a) => Default (Vec n a) where
    def :: Vec n a
def = forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a. Default a => a
def

instance Eq a => Eq (Vec n a) where
    Vec n a
xs == :: Vec n a -> Vec n a -> Bool
== Vec n a
ys = All -> Bool
getAll forall a b. (a -> b) -> a -> b
$ forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Bool -> All
All forall a b. (a -> b) -> a -> b
$ forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 forall a. Eq a => a -> a -> Bool
(==) Vec n a
xs Vec n a
ys
    Vec n a
xs /= :: Vec n a -> Vec n a -> Bool
/= Vec n a
ys = Any -> Bool
getAny forall a b. (a -> b) -> a -> b
$ forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Bool -> Any
Any forall a b. (a -> b) -> a -> b
$ forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 forall a. Eq a => a -> a -> Bool
(/=) Vec n a
xs Vec n a
ys

instance Ord a => Ord (Vec n a) where
    compare :: Vec n a -> Vec n a -> Ordering
compare Vec n a
xs Vec n a
ys = forall (t :: Type -> Type) m. (Foldable t, Monoid m) => t m -> m
F.fold forall a b. (a -> b) -> a -> b
$ forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 forall a. Ord a => a -> a -> Ordering
compare Vec n a
xs Vec n a
ys

mapVec :: (a -> b) -> Vec n a -> Vec n b
mapVec :: forall a b (n :: Nat). (a -> b) -> Vec n a -> Vec n b
mapVec a -> b
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE mapVec #-}

{-# RULES

"mapVec/merge" forall f g va. mapVA f (mapVA g va) = mapVA (f . g) va
"mapVec/coerce" [1]  forall v. materialize (mapVA coerce (access v)) = coerce v
"mapVec/id" mapVA (\x -> x) = id

  #-}

instance Functor (Vec n) where
    fmap :: forall a b. (a -> b) -> Vec n a -> Vec n b
fmap = forall a b (n :: Nat). (a -> b) -> Vec n a -> Vec n b
mapVec

    -- This is just 'pureVec' getting its size from an existing 'Vec'.  Since
    -- the output is independent of the values in the input, we can tie into
    -- the fusion framework to get the size of the input Vec, and still
    -- potentially use the specialized form of 'pureVec'.
    a
x <$ :: forall a b. a -> Vec n b -> Vec n a
<$ Vec n b
v = forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec (forall (n :: Nat) a. Accessor n a -> SInt n
_aSize (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n b
v)) a
x
    {-# INLINE (<$) #-}

instance FunctorWithIndex (Fin n) (Vec n) where imap :: forall a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap = forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap

instance KnownNat n => FoldableWithIndex (Fin n) (Vec n) where
  ifoldMap :: forall m a. Monoid m => (Fin n -> a -> m) -> Vec n a -> m
ifoldMap Fin n -> a -> m
f = forall (t :: Type -> Type) m. (Foldable t, Monoid m) => t m -> m
F.fold forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap Fin n -> a -> m
f

instance KnownNat n => TraversableWithIndex (Fin n) (Vec n) where
  itraverse :: forall (f :: Type -> Type) a b.
Applicative f =>
(Fin n -> a -> f b) -> Vec n a -> f (Vec n b)
itraverse Fin n -> a -> f b
f = forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap Fin n -> a -> f b
f

-- | An element-strict version of 'fmap'.
map' :: (a -> b) -> Vec n a -> Vec n b
map' :: forall a b (n :: Nat). (a -> b) -> Vec n a -> Vec n b
map' a -> b
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Accessor n a -> Accessor n a
seqVA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE map' #-}

-- | A variant of 'fmap' that provides the index in addition to the element.
imap :: (Fin n -> a -> b) -> Vec n a -> Vec n b
imap :: forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap Fin n -> a -> b
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a b.
(Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE imap #-}

-- | Pair each element of a 'Vec' with its index.
--
-- You can also use 'imap', but there should be no harm in using this
-- because the fusion framework should optimize away the intermediate Vec.
withPos :: Vec n a -> Vec n (Fin n, a)
withPos :: forall (n :: Nat) a. Vec n a -> Vec n (Fin n, a)
withPos = forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap (,)
{-# INLINE withPos #-}

-- | An element-strict version of 'imap'.
imap' :: (Fin n -> a -> b) -> Vec n a -> Vec n b
imap' :: forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap' Fin n -> a -> b
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Accessor n a -> Accessor n a
seqVA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a b.
(Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE imap' #-}

pureVec_, pureVec :: SInt n -> a -> Vec n a
pureVec_ :: forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec_ SInt n
n a
x = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) a s. SInt n -> a -> ST s (MutableVec s n a)
newMV SInt n
n a
x forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV
{-# NOINLINE pureVec_ #-}

pureVec :: forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec = \SInt n
n a
x -> forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall (n :: Nat) a. SInt n -> a -> Accessor n a
pureVA SInt n
n a
x)
{-# INLINE pureVec #-}

{-# RULES

"pureVec/spec" [1] forall n x. materialize (pureVA n x) = pureVec_ n x

"mapVA/pureVA" forall f n x. mapVA f (pureVA n x) = pureVA n (f x)

  #-}

liftA2Vec :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Vec :: forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Vec a -> b -> c
f Vec n a
as Vec n b
bs = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall a b c (n :: Nat).
(a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c
liftA2VA a -> b -> c
f (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
as) (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n b
bs))
{-# INLINE liftA2Vec #-}

instance Apply (Vec n) where
  liftF2 :: forall a b c. (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftF2 = forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Vec
  {-# INLINE liftF2 #-}

instance KnownNat n => Applicative (Vec n) where
  pure :: forall a. a -> Vec n a
pure = forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal
  {-# INLINE pure #-}

  liftA2 :: forall a b c. (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2 = forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Vec
  {-# INLINE liftA2 #-}

  <*> :: forall a b. Vec n (a -> b) -> Vec n a -> Vec n b
(<*>) = forall a b c (n :: Nat).
(a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Vec forall a b. (a -> b) -> a -> b
($)
  {-# INLINE (<*>) #-}

  Vec n a
_  *> :: forall a b. Vec n a -> Vec n b -> Vec n b
*> Vec n b
ys = Vec n b
ys
  {-# INLINE (*>) #-}

  Vec n a
xs <* :: forall a b. Vec n a -> Vec n b -> Vec n a
<* Vec n b
_  = Vec n a
xs
  {-# INLINE (<*) #-}

instance Bind (Vec n) where
  Vec n a
xs >>- :: forall a b. Vec n a -> (a -> Vec n b) -> Vec n b
>>- a -> Vec n b
f = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (case forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
xs of
    Accessor SInt n
n Fin n -> (# a #)
get -> forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n (\Fin n
i -> case Fin n -> (# a #)
get Fin n
i of
      (# a
x #) -> a -> Vec n b
f a
x forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
`fusibleFetch` Fin n
i))
  {-# INLINE (>>-) #-}

instance KnownNat n => Monad (Vec n) where >>= :: forall a b. Vec n a -> (a -> Vec n b) -> Vec n b
(>>=) = forall (m :: Type -> Type) a b. Bind m => m a -> (a -> m b) -> m b
(>>-)

-- | A truly lazy version of @liftA2@.
--
-- As opposed to the actual @liftA2@ it does not inspect the arguments which
-- makes it possible it to use in code that has lazy knot-tying.
liftA2Lazy :: SInt n -> (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Lazy :: forall (n :: Nat) a b c.
SInt n -> (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftA2Lazy SInt n
sn a -> b -> c
f Vec n a
xs Vec n b
ys = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec SInt n
sn forall a b. (a -> b) -> a -> b
$ \Fin n
i ->
    forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i forall a b. (a -> b) -> a -> b
$ \a
x ->
    forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n b
ys Fin n
i forall a b. (a -> b) -> a -> b
$ \b
y ->
      a -> b -> c
f a
x b
y

--------------------------------------------------------------------------------
-- | > unsafeIndexK xs i === indexK xs (unsafeFin i)
--
-- TODO(b/109672429): try to get rid of all the uses of this function,
-- and other uses of 'unsafeFin' as well.
unsafeIndexK :: Vec n a -> Int -> (a -> r) -> r
unsafeIndexK :: forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK (V# SmallArray# a
sa) (I# Int#
i) a -> r
k = case forall a. SmallArray# a -> Int# -> (# a #)
indexSmallArray# SmallArray# a
sa Int#
i of (# a
x #) -> a -> r
k a
x
{-# INLINE unsafeIndexK #-}

instance Foldable (Vec n) where
    length :: forall a. Vec n a -> Int
length = forall (n :: Nat) a. Vec n a -> Int
vSize
    {-# INLINE length #-}

    null :: forall a. Vec n a -> Bool
null = (Int
0 forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) a. Foldable t => t a -> Int
length
    {-# INLINE null #-}

    foldMap :: forall m a. Monoid m => (a -> m) -> Vec n a -> m
foldMap a -> m
f = forall m a (n :: Nat). Monoid m => (a -> m) -> Accessor n a -> m
foldMapVA a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
    {-# INLINE foldMap #-}

    fold :: forall m. Monoid m => Vec n m -> m
fold = forall m a (n :: Nat). Monoid m => (a -> m) -> Accessor n a -> m
foldMapVA forall a. a -> a
id forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
    {-# INLINE fold #-}

    foldr :: forall a b. (a -> b -> b) -> b -> Vec n a -> b
foldr a -> b -> b
f b
acc0 = \Vec n a
v ->
      let Accessor SInt n
n Fin n -> (# a #)
get = forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v
      in  forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
n
            (\Fin n
i b
acc -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> a -> b -> b
f a
x b
acc)
            b
acc0
    {-# INLINE foldr #-}

    foldr' :: forall a b. (a -> b -> b) -> b -> Vec n a -> b
foldr' a -> b -> b
f b
acc0 = \Vec n a
v ->
      let !(Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v
          go :: Int -> b -> b
go !Int
i !b
acc
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
acc
            | Bool
otherwise =
                case Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i) of (# a
x #) -> Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1) (a -> b -> b
f a
x b
acc)
      in  Int -> b -> b
go (forall (n :: Nat). SInt n -> Int
unSInt SInt n
n forall a. Num a => a -> a -> a
- Int
1) b
acc0
    {-# INLINE foldr' #-}

    foldl :: forall b a. (b -> a -> b) -> b -> Vec n a -> b
foldl b -> a -> b
f b
acc0 = \Vec n a
v ->
      let !(Accessor SInt n
n Fin n -> (# a #)
get) = forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v
          go :: Int -> b
go !Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
acc0
            | Bool
otherwise = case Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i) of (# a
x #) -> b -> a -> b
f (Int -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1)) a
x
      in  Int -> b
go (forall (n :: Nat). SInt n -> Int
unSInt SInt n
n forall a. Num a => a -> a -> a
- Int
1)
    {-# INLINE foldl #-}

    foldl' :: forall b a. (b -> a -> b) -> b -> Vec n a -> b
foldl' b -> a -> b
f b
acc0 = \Vec n a
v ->
      let !(Accessor (forall (n :: Nat). SInt n -> Int
unSInt -> !Int
n) Fin n -> (# a #)
get) = forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v
          go :: Int -> b -> b
go !Int
i !b
acc
            | Int
i forall a. Ord a => a -> a -> Bool
>= Int
n = b
acc
            | Bool
otherwise =
                case Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i) of (# a
a #) -> Int -> b -> b
go (Int
iforall a. Num a => a -> a -> a
+Int
1) (b -> a -> b
f b
acc a
a)
      in  Int -> b -> b
go Int
0 b
acc0
    {-# INLINE foldl' #-}

    foldr1 :: forall a. (a -> a -> a) -> Vec n a -> a
foldr1 a -> a -> a
f = \Vec n a
v ->
      case forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v of
        Accessor (forall a. Num a => a -> a -> a
subtract Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). SInt n -> Int
unSInt -> Int
lMinus1) Fin n -> (# a #)
get
          | Int
lMinus1 forall a. Ord a => a -> a -> Bool
< Int
0 -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"foldr1@Vec: empty list"
          | Bool
otherwise ->
              let !(# a
z #) = Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
lMinus1)
                  go :: Int -> a
go !Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
lMinus1 = a
z
                        | Bool
otherwise =
                            let !(# a
x #) = Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i)
                            in  a -> a -> a
f a
x (Int -> a
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))
              in  Int -> a
go Int
0
    {-# INLINE foldr1 #-}

    foldl1 :: forall a. (a -> a -> a) -> Vec n a -> a
foldl1 a -> a -> a
f = \Vec n a
v ->
      case forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v of
        Accessor (forall a. Num a => a -> a -> a
subtract Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat). SInt n -> Int
unSInt -> Int
lMinus1) Fin n -> (# a #)
get
          | Int
lMinus1 forall a. Ord a => a -> a -> Bool
< Int
0 -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"foldl1@Vec: empty list"
          | Bool
otherwise ->
              let !(# a
z #) = Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (Int
0 :: Int))
                  go :: Int -> a
go !Int
i | Int
i forall a. Ord a => a -> a -> Bool
<= Int
0 = a
z
                        | Bool
otherwise =
                            let !(# a
x #) = Fin n -> (# a #)
get (forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i)
                            in  a -> a -> a
f (Int -> a
go (Int
i forall a. Num a => a -> a -> a
- Int
1)) a
x
              in  Int -> a
go Int
lMinus1
    {-# INLINE foldl1 #-}

    -- The INLINE declarations here are important to fusion: without it, GHC
    -- fully optimizes the default implementation (which is identical to this
    -- one) when compiling this module, and exposes the post-optimized
    -- unfolding.  Then when we use 'sum', the thing that gets inlined is
    -- already post-fusion, and no fusion can happen.
    --
    -- With the INLINE, GHC exposes (the Core desugaring of) this exact term as
    -- the unfolding, and 'sum' can participate in fusion.

    sum :: forall a. Num a => Vec n a -> a
sum = coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. a -> Sum a
Sum
    {-# INLINE sum #-}

    product :: forall a. Num a => Vec n a -> a
product = coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall a. a -> Product a
Product
    {-# INLINE product #-}

    minimum :: forall a. Ord a => Vec n a -> a
minimum = forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 forall a. Ord a => a -> a -> a
min
    {-# INLINE minimum #-}

    maximum :: forall a. Ord a => Vec n a -> a
maximum = forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 forall a. Ord a => a -> a -> a
max
    {-# INLINE maximum #-}

    elem :: forall a. Eq a => a -> Vec n a -> Bool
elem a
x = coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> Any
Any forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Eq a => a -> a -> Bool
== a
x))
    {-# INLINE elem #-}

traverseVec :: Applicative f => (a -> f b) -> Vec n a -> f (Vec n b)
traverseVec :: forall (f :: Type -> Type) a b (n :: Nat).
Applicative f =>
(a -> f b) -> Vec n a -> f (Vec n b)
traverseVec a -> f b
f = forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
Accessor n (f a) -> f (Vec n a)
sequenceVA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> f b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE traverseVec #-}

-- TODO(b/109674103): for linear-use applicatives ('IO', 'Maybe',...) we
-- can give a more efficient definition.
instance Traversable (Vec n) where
    traverse :: forall (f :: Type -> Type) a b.
Applicative f =>
(a -> f b) -> Vec n a -> f (Vec n b)
traverse = forall (f :: Type -> Type) a b (n :: Nat).
Applicative f =>
(a -> f b) -> Vec n a -> f (Vec n b)
traverseVec
    {-# INLINE traverse #-}

-- We don't bother defining 'collect', since anything other than the
-- default instance would be egregiously inefficient if the function
-- actually allocates vectors, since we'd end up allocating @n@ identical
-- copies of every vector in the @f@ structure! If, however, we were
-- to have a bunch of fusion rules for things like @mkVec f ! i = f i@;
-- then, it might be worth defining 'collect' directly, since we could
-- avoid allocating any of the intermediate arrays!
--
-- TODO(b/109674103): for linear-use functors ('IO', 'Maybe',...) we
-- can give a more efficient definition.
instance KnownNat n => Distributive (Vec n) where
    distribute :: forall (f :: Type -> Type) a.
Functor f =>
f (Vec n a) -> Vec n (f a)
distribute f (Vec n a)
fxs = forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate forall a b. (a -> b) -> a -> b
$ \Rep (Vec n)
i -> forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (n :: Nat) a. Vec n a -> Fin n -> a
! Rep (Vec n)
i) f (Vec n a)
fxs
    {-# INLINE distribute #-}

instance KnownNat n => Representable (Vec n) where
    type Rep (Vec n) = Fin n

    tabulate :: forall a. (Rep (Vec n) -> a) -> Vec n a
tabulate = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal
    {-# INLINE tabulate #-}

    index :: forall a. Vec n a -> Rep (Vec n) -> a
index = forall (n :: Nat) a. Vec n a -> Fin n -> a
(!)
    {-# INLINE index #-}

-- | 'Prelude.scanl', for 'Vec'.
vscanl :: (b -> a -> b) -> b -> Vec n a -> Vec (1 + n) b
-- TODO(awpr): we can probably subject the input Vec to fusion here.
vscanl :: forall b a (n :: Nat).
(b -> a -> b) -> b -> Vec n a -> Vec (1 + n) b
vscanl b -> a -> b
f b
b Vec n a
v = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall (m :: Nat) (n :: Nat).
(?callStack::CallStack) =>
SInt m -> SInt n -> SInt (m + n)
`addSInt` forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl b -> a -> b
f b
b forall a b. (a -> b) -> a -> b
$ forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList Vec n a
v


--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
-- List like ops and their isomorphisms.

-- | A zero-length 'Vec' of any element type.
nil :: Vec 0 a
-- Note: in the C-- code, this is a single global thunk with a polymorphic
-- type; we won't re-create it separately for different types.  The NOINLINE
-- serves to ensure we just reference the thunk rather than inlining the code
-- that builds it.  This blocks fusion, but we have another trick: rewrite any
-- materialize @0 to nil.  Then anything that would've fused with this will no
-- longer reference it at all.
nil :: forall a. Vec 0 a
nil = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal (\Fin 0
i -> Fin 0
i seq :: forall a b. a -> b -> b
`seq` forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Vec.nil: the impossible happened")
{-# NOINLINE nil #-}

{-# RULES

-- Note this matches only zero-length vectors because the LHS type is unified
-- with the RHS type.
"materialize@0" forall va. materialize va = nil

  #-}

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

-- | Concatenate two 'Vec's.
infixr 5 ++
append_, (++) :: Vec n a -> Vec m a -> Vec (n + m) a
append_ :: forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
append_ Vec n a
xs Vec m a
ys = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    let !n :: Int
n = forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs
        !m :: Int
m = forall (n :: Nat) a. Vec n a -> Int
vSize Vec m a
ys
    MutableVec s (n + m) a
zs <- forall a s (n :: Nat). Int -> a -> ST s (MutableVec s n a)
unsafeNewMV (Int
n forall a. Num a => a -> a -> a
+ Int
m) forall a. a
uninitialized
    forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
0 MutableVec s (n + m) a
zs Int
0 Int
n
    forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec m a
ys Int
0 MutableVec s (n + m) a
zs Int
n Int
m
    forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s (n + m) a
zs
{-# NOINLINE append_ #-}

++ :: forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
(++) = \Vec n a
xs Vec m a
ys -> forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall (n :: Nat) a (m :: Nat).
Accessor n a -> Accessor m a -> Accessor (n + m) a
appendVA (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
xs) (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec m a
ys))
{-# INLINE (++) #-}

{-# RULES

"++/spec" [1]
  forall xs ys. materialize (appendVA (access xs) (access ys)) = xs `append_` ys

  #-}

-- TODO(b/109675695): may want other variants of this operational-function
-- with different types.
-- TODO: might as well simply call the underlying
-- 'thawSmallArray#' directly, instead of passing through
-- 'sliceVec'... ne? (i.e., to avoid the dynamic bounds checks)
take_ :: SInt m -> Vec (m + n) a -> Vec m a
take_ :: forall (m :: Nat) (n :: Nat) a. SInt m -> Vec (m + n) a -> Vec m a
take_ SInt m
m Vec (m + n) a
xs = forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec (m + n) a
xs Int
0 SInt m
m
{-# NOINLINE take_ #-}
-- TODO(awpr): fusion behaves poorly here because of Core-level casts arising
-- from unifying @m + n@ with some other Nat: we get
-- @materialize (takeVA m (access ((materialize _) `cast` <Co:1>)))@, which
-- can't match the access/materialize rule.  The @cast@ comes from converting
-- the equal-but-not-syntactically-identical types @Vec (m + n) a@ and
-- @Vec _something a@ using the coercion @_something ~N# (m + n)@.  If we
-- define this with inequality constraints instead,
-- @take_ :: (m <= n) => SInt m -> Vec n a -> Vec m a@, then there's no need
-- for a cast there; @n@ and @m@ are just unified away to be identical to the
-- sizes of the input and output vectors.

{-# RULES

"take_/spec" [1] forall m xs. materialize (takeVA m (access xs)) = take_ m xs

  #-}

-- TODO(b/109675695): may want other variants of this operational-function
-- with different types.
drop_ :: forall m n a. SInt m -> Vec (m + n) a -> Vec n a
drop_ :: forall (m :: Nat) (n :: Nat) a. SInt m -> Vec (m + n) a -> Vec n a
drop_ SInt m
m Vec (m + n) a
xs =
  forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec (m + n) a
xs (forall (n :: Nat). SInt n -> Int
unSInt SInt m
m) forall a b. (a -> b) -> a -> b
$
  forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec (m + n) a
xs forall (m :: Nat) (n :: Nat). SInt (m + n) -> SInt m -> SInt n
`subSIntL` SInt m
m
{-# NOINLINE drop_ #-}
-- TODO(awpr): as with 'take_', casts are causing trouble here.  Consider
-- messing with the type signature to avoid them.

{-# RULES

"drop_/spec" [1] forall m xs. materialize (dropVA m (access xs)) = drop_ m xs

  #-}

-- | Split a vector into two shorter vectors at the given index.
split
  :: forall m n a. SInt m -> Vec (m + n) a -> (Vec m a, Vec n a)
-- TODO(b/109675695): may want other variants of this operational-function
-- with different types.
split :: forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec (m + n) a -> (Vec m a, Vec n a)
split SInt m
m Vec (m + n) a
xs =
  let va :: Accessor (m + n) a
va = forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec (m + n) a
xs
  in  (forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall (m :: Nat) (n :: Nat) a.
SInt m -> Accessor (m + n) a -> Accessor m a
takeVA SInt m
m Accessor (m + n) a
va), forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall (m :: Nat) (n :: Nat) a.
SInt m -> Accessor (m + n) a -> Accessor n a
dropVA SInt m
m Accessor (m + n) a
va))
{-# INLINE split #-}

-- TODO(awpr): fusion for 'concat' and 'reshape'?

-- | Concatenate a nested 'Vec' into one longer 'Vec'.
concat :: forall m n a. Vec n (Vec m a) -> Vec (n * m) a
concat :: forall (m :: Nat) (n :: Nat) a. Vec n (Vec m a) -> Vec (n * m) a
concat Vec n (Vec m a)
xs =
  let !n :: Int
n = forall (n :: Nat) a. Vec n a -> Int
vSize Vec n (Vec m a)
xs
  in  if Int
n forall a. Eq a => a -> a -> Bool
== Int
0 then
        -- Outer size is 0, return the empty array
        forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec Int
0 forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n * m) a
_ -> forall (m :: Type -> Type) a. Monad m => a -> m a
return ()
      else
          let !m :: Int
m = forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
0 forall (n :: Nat) a. Vec n a -> Int
vSize  -- We know the size is > 0.
          in  forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
n forall a. Num a => a -> a -> a
* Int
m) forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n * m) a
marr ->
                forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
F.forM_ [Int
0..Int
nforall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \ Int
i ->
                  forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
i forall a b. (a -> b) -> a -> b
$ \ Vec m a
ys ->  -- i is [0..n-1], so in bounds.
                    forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec m a
ys Int
0 MutableVec s (n * m) a
marr (Int
i forall a. Num a => a -> a -> a
* Int
m) Int
m
{-# INLINE concat #-}

-- | Turn a vector into a vector of vector by chunking it.
reshape :: SInt m -> Vec (n * m) a -> Vec n (Vec m a)
reshape :: forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec (n * m) a -> Vec n (Vec m a)
reshape SInt m
m =
  let !m' :: Int
m' = forall (n :: Nat). SInt n -> Int
unSInt SInt m
m
  in  \Vec (n * m) a
xs ->
        forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec (n * m) a
xs forall (m :: Nat) (n :: Nat). SInt (m * n) -> SInt n -> SInt m
`divSIntR` SInt m
m) (\Fin n
i -> forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec (n * m) a
xs (forall (n :: Nat). Fin n -> Int
finToInt Fin n
i forall a. Num a => a -> a -> a
* Int
m') SInt m
m)
{-# INLINE reshape #-}

-- | Map each element of a 'Vec' to a (same-sized) sub-'Vec' of the result.
concatMap :: forall m n a b. (a -> Vec m b) -> Vec n a -> Vec (n * m) b
concatMap :: forall (m :: Nat) (n :: Nat) a b.
(a -> Vec m b) -> Vec n a -> Vec (n * m) b
concatMap a -> Vec m b
f = forall (m :: Nat) (n :: Nat) a. Vec n (Vec m a) -> Vec (n * m) a
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Vec m b
f
{-# INLINE concatMap #-}


-- | Generate a Vec by repeated application of a function.
--
-- > toList (Vec.iterate @n f z) === take (valueOf @n) (Prelude.iterate f z)
iterate :: SInt n -> (a -> a) -> a -> Vec n a
iterate :: forall (n :: Nat) a. SInt n -> (a -> a) -> a -> Vec n a
iterate SInt n
sn a -> a
f a
z =
    forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
sn forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv ->
       forall (m :: Type -> Type) (n :: Nat) a.
Monad m =>
SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ SInt n
sn (\a
x Fin n
i -> a -> a
f a
x forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i a
x) a
z
{-# INLINE iterate #-}


-- | A strict version of 'iterate'.
iterate' :: SInt n -> (a -> a) -> a -> Vec n a
iterate' :: forall (n :: Nat) a. SInt n -> (a -> a) -> a -> Vec n a
iterate' SInt n
sn a -> a
f !a
z =
    forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
sn forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv ->
        forall (m :: Type -> Type) (n :: Nat) a.
Monad m =>
SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ SInt n
sn (\a
x Fin n
i -> a -> a
f a
x forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ (forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i forall a b. (a -> b) -> a -> b
$! a
x)) a
z
{-# INLINE iterate' #-}


-- | Return a copy of the array with elements in the reverse order.
rev :: Vec n a -> Vec n a
rev :: forall (n :: Nat) a. Vec n a -> Vec n a
rev = forall (n :: Nat) a. Accessor n a -> Vec n a
materialize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Accessor n a -> Accessor n a
revVA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE rev #-}


-- | Rotate a vector right by @i@ positions.
--
-- @rot 1 [x, y, z] = [z, x, y]@
rot, rot_ :: Fin n -> Vec n a -> Vec n a
rot_ :: forall (n :: Nat) a. Fin n -> Vec n a -> Vec n a
rot_ Fin n
o' Vec n a
xs = forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> do
  let o :: Int
o = forall (n :: Nat). Fin n -> Int
finToInt Fin n
o'
      nmo :: Int
nmo = forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs forall a. Num a => a -> a -> a
- Int
o
  forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
nmo MutableVec s n a
mv Int
0 Int
o
  forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
0   MutableVec s n a
mv Int
o Int
nmo
{-# NOINLINE rot_ #-}

rot :: forall (n :: Nat) a. Fin n -> Vec n a -> Vec n a
rot Fin n
o = \Vec n a
v -> forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (forall (n :: Nat) a. Fin n -> Accessor n a -> Accessor n a
rotVA Fin n
o (forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
v))
{-# INLINE rot #-}

{-# RULES

"rot/spec" [1] forall o v. materialize (rotVA o (access v)) = rot_ o v

  #-}

-- | Return a vector with all elements of the type in ascending order.
viota :: SInt n -> Vec n (Fin n)
viota :: forall (n :: Nat). SInt n -> Vec n (Fin n)
viota SInt n
sn = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt n
sn forall a. a -> a
id
{-# INLINE viota #-}

-- | One variant of the cross product of two vectors.
cross :: Vec m a -> Vec n b -> Vec n (Vec m (a, b))
cross :: forall (m :: Nat) a (n :: Nat) b.
Vec m a -> Vec n b -> Vec n (Vec m (a, b))
cross Vec m a
xs = forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\b
y -> forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (, b
y) Vec m a
xs)
{-# INLINE cross #-}

-- TODO: the concatenated version.
-- cross :: Vec m a -> Vec n b -> Vec (m * n) (a, b)
--
-- TODO: the transposed nested version.
-- cross :: Vec m a -> Vec n b -> Vec m (Vec n (a, b))

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

-- Element insertion and removal.

-- TODO(awpr): fusion implementations of insert and remove?

-- Unsafe version of insert.  Assumes i < valueOf @(n+1)
-- Statically determined 0 length copies are removed by a RULE.
unsafeInsert :: Int -> a -> Vec n a -> Vec (n+1) a
unsafeInsert :: forall a (n :: Nat). Int -> a -> Vec n a -> Vec (n + 1) a
unsafeInsert Int
i a
xi Vec n a
xs =
  let !n :: Int
n = forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs
  in  forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
nforall a. Num a => a -> a -> a
+Int
1) forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n + 1) a
mv -> do
        -- Explicitly: @mv[0..i-1] := xs[0..i-1]@
        forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
0 MutableVec s (n + 1) a
mv Int
0 Int
i
        -- Explicitly: @mv[i] := xi@
        forall s (n :: Nat) a. MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV MutableVec s (n + 1) a
mv Int
i a
xi
        -- Explicitly: @mv[i+1..n] := xs[i..n-1]@
        forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
i MutableVec s (n + 1) a
mv (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
n forall a. Num a => a -> a -> a
- Int
i)

-- | Insert an element at the given position in a vector.
-- O(n)
insert :: Fin (n+1) -> a -> Vec n a -> Vec (n+1) a
insert :: forall (n :: Nat) a. Fin (n + 1) -> a -> Vec n a -> Vec (n + 1) a
insert Fin (n + 1)
i = forall a (n :: Nat). Int -> a -> Vec n a -> Vec (n + 1) a
unsafeInsert (forall (n :: Nat). Fin n -> Int
finToInt Fin (n + 1)
i)

-- Unsafe version of remove.  Assumes i < valueOf @(n+1)
-- Statically determined 0 length copies are removed by a RULE.
unsafeRemove :: Int -> Vec (n+1) a -> Vec n a
unsafeRemove :: forall (n :: Nat) a. Int -> Vec (n + 1) a -> Vec n a
unsafeRemove Int
i Vec (n + 1) a
xs =
  let !np1 :: Int
np1 = forall (n :: Nat) a. Vec n a -> Int
vSize Vec (n + 1) a
xs
  in  forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
np1 forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$ \ MutableVec s n a
mv -> do
        -- Explicitly: @mv[0..i-1] := xs[0..i-1]@
        forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec (n + 1) a
xs Int
0 MutableVec s n a
mv Int
0 Int
i
        -- Explicitly: @mv[i..n-1] := xs[i+1..n]@
        forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec (n + 1) a
xs (Int
iforall a. Num a => a -> a -> a
+Int
1) MutableVec s n a
mv Int
i (Int
np1 forall a. Num a => a -> a -> a
- Int
1 forall a. Num a => a -> a -> a
- Int
i)

-- | Remove an element at the given position in a vector.
-- O(n)
remove :: Fin (n+1) -> Vec (n+1) a -> Vec n a
remove :: forall (n :: Nat) a. Fin (n + 1) -> Vec (n + 1) a -> Vec n a
remove Fin (n + 1)
i = forall (n :: Nat) a. Int -> Vec (n + 1) a -> Vec n a
unsafeRemove (forall (n :: Nat). Fin n -> Int
finToInt Fin (n + 1)
i)

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

-- | Sort a 'Vec' according to its 'Ord' instance.
vsort :: Ord a => Vec n a -> Vec n a
vsort :: forall a (n :: Nat). Ord a => Vec n a -> Vec n a
vsort Vec n a
xs = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
L.sort forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList forall a b. (a -> b) -> a -> b
$ Vec n a
xs

-- | Sort a 'Vec' with a given comparison function.
vsortBy :: (a -> a -> Ordering) -> Vec n a -> Vec n a
vsortBy :: forall a (n :: Nat). (a -> a -> Ordering) -> Vec n a -> Vec n a
vsortBy a -> a -> Ordering
f Vec n a
xs = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs)forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy a -> a -> Ordering
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList forall a b. (a -> b) -> a -> b
$ Vec n a
xs

-- | Sort a 'Vec' with a given sort-key function.
vsortOn :: Ord b => (a -> b) -> Vec n a -> Vec n a
vsortOn :: forall b a (n :: Nat). Ord b => (a -> b) -> Vec n a -> Vec n a
vsortOn a -> b
f Vec n a
xs = forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs)forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. Ord b => (a -> b) -> [a] -> [a]
L.sortOn a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList forall a b. (a -> b) -> a -> b
$ Vec n a
xs


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

-- | Transpose a vector of vectors.
vtranspose :: SInt m -> Vec n (Vec m a) -> Vec m (Vec n a)
vtranspose :: forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec n (Vec m a) -> Vec m (Vec n a)
vtranspose SInt m
sm Vec n (Vec m a)
xs =
  let !s :: Int
s = forall (n :: Nat) a. Vec n a -> Int
vSize Vec n (Vec m a)
xs
      !t :: Int
t = forall (n :: Nat). SInt n -> Int
unSInt SInt m
sm
  in  forall a (n :: Nat). Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
t forall a b. (a -> b) -> a -> b
$ \ Int
i ->
        -- s is the size of the outer vector, i.e. valueOf @n
        forall a (n :: Nat). Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
s forall a b. (a -> b) -> a -> b
$ \ Int
j ->
          forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
j forall a b. (a -> b) -> a -> b
$ \ Vec m a
v -> forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec m a
v Int
i forall a. a -> a
id

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

-- | Find the index of the first element, if any, that satisfies a predicate.
vfindIndex :: (a -> Bool) -> Vec n a -> Maybe (Fin n)
vfindIndex :: forall a (n :: Nat). (a -> Bool) -> Vec n a -> Maybe (Fin n)
vfindIndex a -> Bool
p = forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> Maybe Int
L.findIndex a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList

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

-- | Create a singleton 'Vec'.
vec1 :: a -> Vec 1 a
vec1 :: forall a. a -> Vec 1 a
vec1 = forall (f :: Type -> Type) a. Applicative f => a -> f a
pure
{-# INLINE vec1 #-}

-- | Create a 'Vec' from two elements.
vec2 :: a -> a -> Vec 2 a
vec2 :: forall a. a -> a -> Vec 2 a
vec2 a
x0 a
x1 = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall a b. (a -> b) -> a -> b
$ \Fin 2
i -> case forall (n :: Nat). Fin n -> Int
finToInt Fin 2
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
_ -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Impossible: Fin out of range"
{-# INLINE vec2 #-}

-- | Create a 'Vec' from three elements.
vec3 :: a -> a -> a -> Vec 3 a
vec3 :: forall a. a -> a -> a -> Vec 3 a
vec3 a
x0 a
x1 a
x2 = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall a b. (a -> b) -> a -> b
$ \Fin 3
i -> case forall (n :: Nat). Fin n -> Int
finToInt Fin 3
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
2 -> a
x2
  Int
_ -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Impossible: Fin out of range"
{-# INLINE vec3 #-}

-- | Create a 'Vec' from four elements.
vec4 :: a -> a -> a -> a -> Vec 4 a
vec4 :: forall a. a -> a -> a -> a -> Vec 4 a
vec4 a
x0 a
x1 a
x2 a
x3 = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall a b. (a -> b) -> a -> b
$ \Fin 4
i -> case forall (n :: Nat). Fin n -> Int
finToInt Fin 4
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
2 -> a
x2
  Int
3 -> a
x3
  Int
_ -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Impossible: Fin out of range"
{-# INLINE vec4 #-}

-- | Create a 'Vec' from six elements.
vec6 :: a -> a -> a -> a -> a -> a ->Vec 6 a
vec6 :: forall a. a -> a -> a -> a -> a -> a -> Vec 6 a
vec6 a
x0 a
x1 a
x2 a
x3 a
x4 a
x5 = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall a b. (a -> b) -> a -> b
$ \Fin 6
i -> case forall (n :: Nat). Fin n -> Int
finToInt Fin 6
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
2 -> a
x2
  Int
3 -> a
x3
  Int
4 -> a
x4
  Int
5 -> a
x5
  Int
_ -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Impossible: Fin out of range"
{-# INLINE vec6 #-}

-- | Create a 'Vec' from eight elements.
vec8 :: a -> a -> a -> a -> a -> a -> a -> a -> Vec 8 a
vec8 :: forall a. a -> a -> a -> a -> a -> a -> a -> a -> Vec 8 a
vec8 a
x0 a
x1 a
x2 a
x3 a
x4 a
x5 a
x6 a
x7 = forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal forall a b. (a -> b) -> a -> b
$ \Fin 8
i -> case forall (n :: Nat). Fin n -> Int
finToInt Fin 8
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
2 -> a
x2
  Int
3 -> a
x3
  Int
4 -> a
x4
  Int
5 -> a
x5
  Int
6 -> a
x6
  Int
7 -> a
x7
  Int
_ -> forall a. (?callStack::CallStack) => [Char] -> a
error [Char]
"Impossible: Fin out of range"
{-# INLINE vec8 #-}

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

-- | Get the value of a statically known natural number.
{-# INLINE valueOf #-}
valueOf :: forall (n :: Nat) (i :: Type) . (KnownNat n, Num i) => i
valueOf :: forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (n :: Nat). KnownNat n => Proxy# n -> Nat
natVal' (forall {k} (a :: k). Proxy# a
proxy# :: Proxy# n)

#if !MIN_VERSION_base(4,15,0)
-- base-4.15.0.0 removed naturalToInt.
{-# RULES "integerToInt . naturalToInteger => naturalToInt"
  forall a. integerToInt (naturalToInteger a) =
      let !(I# i) = naturalToInt a
      in i
  #-}
#endif

-- | Modify the given index of a 'Vec'.
overIx :: Fin n -> (a -> a) -> Vec n a -> Vec n a
overIx :: forall (n :: Nat) a. Fin n -> (a -> a) -> Vec n a -> Vec n a
overIx Fin n
i a -> a
f Vec n a
v = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MutableVec s n a
mv <- forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV Vec n a
v
  forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
v Fin n
i (forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f)
  forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s n a
mv
{-# INLINE overIx #-}