-- Copyright 2018-2021 Google LLC
--
-- 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)

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

forMFin_ :: Applicative f => SInt n -> (Fin n -> f a) -> f ()
forMFin_ :: SInt n -> (Fin n -> f a) -> f ()
forMFin_ SInt n
n Fin n -> f a
f = SInt n -> (Fin n -> f () -> f ()) -> f () -> 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 f a -> f () -> f ()
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> f ()
rest) (() -> f ()
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ())
{-# INLINE forMFin_ #-}

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

foldMFin_
  :: Monad m
  => SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ :: SInt n -> (a -> Fin n -> m a) -> a -> m ()
foldMFin_ SInt n
n a -> Fin n -> m a
f a
z = SInt n
-> (Fin n -> (a -> m ()) -> a -> m ()) -> (a -> m ()) -> a -> m ()
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 m a -> (a -> m ()) -> m ()
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m ()
rest) (\a
_ -> () -> m ()
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 :: SInt n -> (Fin n -> (# f a #)) -> f [a]
forMFin SInt n
n Fin n -> (# f a #)
f = SInt n -> (Fin n -> f [a] -> f [a]) -> f [a] -> f [a]
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 #) -> (a -> [a] -> [a]) -> f a -> f [a] -> f [a]
forall (f :: Type -> Type) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (:) f a
fx)
  ([a] -> f [a]
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_ :: (State# s -> State# s) -> ST s ()
prim_ State# s -> State# s
f = STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (STRep s () -> ST s ()) -> STRep s () -> ST s ()
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 :: SInt n -> a -> ST s (MutableVec s n a)
newMV (SI# Int
n) = Int -> a -> ST s (MutableVec s n a)
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 :: Int -> a -> ST s (MutableVec s n a)
unsafeNewMV (I# Int#
n) a
x =
    STRep s (MutableVec s n a) -> ST s (MutableVec s n a)
forall s a. STRep s a -> ST s a
ST (STRep s (MutableVec s n a) -> ST s (MutableVec s n a))
-> STRep s (MutableVec s n a) -> ST s (MutableVec s n a)
forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case Int# -> a -> State# s -> (# State# s, SmallMutableArray# s a #)
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, SmallMutableArray# s a -> MutableVec s n a
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 :: MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV (MV# SmallMutableArray# s a
sma) (I# Int#
i) a
x = (State# s -> State# s) -> ST s ()
forall s. (State# s -> State# s) -> ST s ()
prim_ (SmallMutableArray# s a -> Int# -> a -> State# s -> State# s
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 :: MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i = MutableVec s n a -> Int -> a -> ST s ()
forall s (n :: Nat) a. MutableVec s n a -> Int -> a -> ST s ()
unsafeWriteMV MutableVec s n a
mv (Fin n -> Int
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 :: MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV (MV# SmallMutableArray# s a
sma) =
    STRep s (Vec n a) -> ST s (Vec n a)
forall s a. STRep s a -> ST s a
ST (STRep s (Vec n a) -> ST s (Vec n a))
-> STRep s (Vec n a) -> ST s (Vec n a)
forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case SmallMutableArray# s a -> State# s -> (# State# s, SmallArray# a #)
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, SmallArray# a -> Vec n a
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 :: Vec n a -> ST s (MutableVec s n a)
safeThawMV (V# SmallArray# a
sa) =
    STRep s (MutableVec s n a) -> ST s (MutableVec s n a)
forall s a. STRep s a -> ST s a
ST (STRep s (MutableVec s n a) -> ST s (MutableVec s n a))
-> STRep s (MutableVec s n a) -> ST s (MutableVec s n a)
forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case SmallArray# a
-> Int#
-> Int#
-> State# s
-> (# State# s, SmallMutableArray# s a #)
forall a d.
SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
thawSmallArray# SmallArray# a
sa Int#
0# (SmallArray# a -> Int#
forall a. SmallArray# a -> Int#
sizeofSmallArray# SmallArray# a
sa) State# s
s0 of { (# State# s
s1, SmallMutableArray# s a
sma #) ->
    (# State# s
s1, SmallMutableArray# s a -> MutableVec s n a
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 :: 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) =
    (State# s -> State# s) -> ST s ()
forall s. (State# s -> State# s) -> ST s ()
prim_ (SmallArray# a
-> Int#
-> SmallMutableArray# s a
-> Int#
-> Int#
-> State# s
-> State# s
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 :: 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)) =
    Bool -> Vec m a -> Vec m a
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
off Bool -> Bool -> Bool
&& Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
len Bool -> Bool -> Bool
&& Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off) (Vec m a -> Vec m a) -> Vec m a -> Vec m a
forall a b. (a -> b) -> a -> b
$
    SmallArray# a -> Vec m a
forall (n :: Nat) a. SmallArray# a -> Vec n a
V# (SmallArray# a -> Int# -> Int# -> SmallArray# a
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 :: Vec n a -> Fin n -> (# a #)
fetch (V# SmallArray# a
arr) (Fin n -> Int
forall (n :: Nat). Fin n -> Int
finToInt -> I# Int#
i) = SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexSmallArray# SmallArray# a
arr Int#
i
{-# INLINE fetch #-}

fusibleFetch :: Vec n a -> Fin n -> (# a #)
fusibleFetch :: Vec n a -> Fin n -> (# a #)
fusibleFetch = Accessor n a -> Fin n -> (# a #)
forall (n :: Nat) a. Accessor n a -> Fin n -> (# a #)
_aIndex (Accessor n a -> Fin n -> (# a #))
-> (Vec n a -> Accessor n a) -> Vec n a -> Fin n -> (# a #)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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
(!) Vec n a
xs Fin n
i = case Vec n a -> Fin n -> (# a #)
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 :: Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
v Fin n
i a -> r
k = case Vec n a -> Fin n -> (# a #)
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 :: Vec n a -> SInt n
svSize (V# SmallArray# a
sa) = Int -> SInt n
forall (n :: Nat). Int -> SInt n
SI# (Int# -> Int
I# (SmallArray# a -> Int#
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 :: Vec n a -> Int
vSize = SInt n -> Int
forall (n :: Nat). SInt n -> Int
unSInt (SInt n -> Int) -> (Vec n a -> SInt n) -> Vec n a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize
{-# INLINE vSize #-}


--------------------------------------------------------------------------------
uninitialized :: a
uninitialized :: a
uninitialized = [Char] -> a
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 :: Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec Int
n = SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec (Int -> SInt n
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 :: 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 s. ST s (Vec n a)) -> Vec n a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec n a)) -> Vec n a)
-> (forall s. ST s (Vec n a)) -> Vec n a
forall a b. (a -> b) -> a -> b
$ do
    MutableVec s n a
mv <- SInt n -> a -> ST s (MutableVec s n a)
forall (n :: Nat) a s. SInt n -> a -> ST s (MutableVec s n a)
newMV SInt n
n a
forall a. a
uninitialized
    MutableVec s n a -> ST s ()
forall s. MutableVec s n a -> ST s ()
action MutableVec s n a
mv
    MutableVec s n a -> ST s (Vec n a)
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
  { Accessor n a -> SInt n
_aSize :: SInt n
  , 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 :: Vec n a -> Accessor n a
access Vec n a
v = SInt n -> (Fin n -> (# a #)) -> Accessor n a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
v) (Vec n a -> Fin n -> (# a #)
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 :: Accessor n a -> Vec n a
materialize (Accessor SInt n
n Fin n -> (# a #)
get) = SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
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 ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> SInt n -> (Fin n -> ST s ()) -> ST s ()
forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
SInt n -> (Fin n -> f a) -> f ()
forMFin_ SInt n
n ((Fin n -> ST s ()) -> ST s ()) -> (Fin n -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Fin n
i ->
  case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> MutableVec s n a -> Fin n -> a -> ST s ()
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

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

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

  #-}

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

mapVA :: (a -> b) -> Accessor n a -> Accessor n b
mapVA :: (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f (Accessor SInt n
n Fin n -> (# a #)
get) = SInt n -> (Fin n -> (# b #)) -> Accessor n b
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# b #)) -> Accessor n b)
-> (Fin n -> (# b #)) -> Accessor n b
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 :: (Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f (Accessor SInt n
n Fin n -> (# a #)
get) = SInt n -> (Fin n -> (# b #)) -> Accessor n b
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# b #)) -> Accessor n b)
-> (Fin n -> (# b #)) -> Accessor n b
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 :: Accessor n a -> Accessor n a
seqVA (Accessor SInt n
n Fin n -> (# a #)
get) = SInt n -> (Fin n -> (# a #)) -> Accessor n a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# a #)) -> Accessor n a)
-> (Fin n -> (# a #)) -> Accessor n a
forall a b. (a -> b) -> a -> b
$
  \Fin n
i -> case Fin n -> (# a #)
get Fin n
i of (# a
x #) -> a
x a -> (# a #) -> (# a #)
`seq` (# a
x #)
{-# INLINE [0] seqVA #-}

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

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

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

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

rotVA :: Fin n -> Accessor n a -> Accessor n a
rotVA :: Fin n -> Accessor n a -> Accessor n a
rotVA (Fin n -> Int
forall (n :: Nat). Fin n -> Int
finToInt -> !Int
o) (Accessor SInt n
n Fin n -> (# a #)
get) = SInt n -> (Fin n -> (# a #)) -> Accessor n a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# a #)) -> Accessor n a)
-> (Fin n -> (# a #)) -> Accessor n a
forall a b. (a -> b) -> a -> b
$
  \(Fin n -> Int
forall (n :: Nat). Fin n -> Int
finToInt -> !Int
i) -> Fin n -> (# a #)
get (Fin n -> (# a #)) -> Fin n -> (# a #)
forall a b. (a -> b) -> a -> b
$ Int -> Fin n
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (Int -> Fin n) -> Int -> Fin n
forall a b. (a -> b) -> a -> b
$ if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
o then Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o else Int
nmo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
 where
  !nmo :: Int
nmo = SInt n -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt n
n Int -> Int -> Int
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 :: (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) = SInt n -> (Fin n -> (# c #)) -> Accessor n c
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# c #)) -> Accessor n c)
-> (Fin n -> (# c #)) -> Accessor n c
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 :: (a -> m) -> Accessor n a -> m
foldMapVA a -> m
f (Accessor SInt n
n Fin n -> (# a #)
get) =
  SInt n -> (Fin n -> (# m #)) -> m
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 :: Accessor n (f a) -> f (Vec n a)
sequenceVA (Accessor SInt n
n Fin n -> (# f a #)
get) = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
n ([a] -> Vec n a) -> f [a] -> f (Vec n a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> SInt n -> (Fin n -> (# f a #)) -> f [a]
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 :: SInt n -> Fin (n + m) -> Either (Fin n) (Fin m)
splitFinS (SI# Int
n) (Fin (n + m) -> Int
forall (n :: Nat). Fin n -> Int
finToInt -> Int
i)
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n     = Fin n -> Either (Fin n) (Fin m)
forall a b. a -> Either a b
Left (Int -> Fin n
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin Int
i)
  | Bool
otherwise = Fin m -> Either (Fin n) (Fin m)
forall a b. b -> Either a b
Right (Int -> Fin m
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n))

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

appendVA :: Accessor n a -> Accessor m a -> Accessor (n + m) a
appendVA :: 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) = SInt (n + m) -> (Fin (n + m) -> (# a #)) -> Accessor (n + m) a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor
  (SInt n -> SInt m -> SInt (n + m)
forall (n :: Nat) (m :: Nat). SInt n -> SInt m -> SInt (n + m)
addPosSInt SInt n
n SInt m
m)
  (\Fin (n + m)
i -> case SInt n -> Fin (n + m) -> Either (Fin n) (Fin m)
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 :: Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
n Int -> a
f = SInt n -> (Fin n -> a) -> Vec n a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec (Int -> SInt n
forall (n :: Nat). Int -> SInt n
SI# Int
n) ((Fin n -> a) -> Vec n a) -> (Fin n -> a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \Fin n
i -> Int -> a
f (Fin n -> Int
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 :: SInt n -> (Fin n -> a) -> Vec n a
tabulateVec SInt n
n Fin n -> a
f = Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n a -> Vec n a) -> Accessor n a -> Vec n a
forall a b. (a -> b) -> a -> b
$ SInt n -> (Fin n -> (# a #)) -> Accessor n a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# a #)) -> Accessor n a)
-> (Fin n -> (# a #)) -> Accessor n a
forall a b. (a -> b) -> a -> b
$ \Fin n
i -> (# Fin n -> a
f Fin n
i #)
mkVec :: SInt n -> (Fin n -> a) -> Vec n a
mkVec = SInt n -> (Fin n -> a) -> Vec n a
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' :: SInt n -> (Fin n -> a) -> Vec n a
tabulateVec' SInt n
n Fin n -> a
f = Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n a -> Vec n a) -> Accessor n a -> Vec n a
forall a b. (a -> b) -> a -> b
$
  SInt n -> (Fin n -> (# a #)) -> Accessor n a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt n
n ((Fin n -> (# a #)) -> Accessor n a)
-> (Fin n -> (# a #)) -> Accessor n a
forall a b. (a -> b) -> a -> b
$ \Fin n
i -> let x :: a
x = Fin n -> a
f Fin n
i in a
x a -> (# a #) -> (# a #)
`seq` (# a
x #)
mkVec' :: SInt n -> (Fin n -> a) -> Vec n a
mkVec' = SInt n -> (Fin n -> a) -> Vec n a
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 :: 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 = Accessor m a -> Vec m a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor m a -> Vec m a) -> Accessor m a -> Vec m a
forall a b. (a -> b) -> a -> b
$ SInt m -> (Fin m -> (# a #)) -> Accessor m a
forall (n :: Nat) a. SInt n -> (Fin n -> (# a #)) -> Accessor n a
Accessor SInt m
m ((Fin m -> (# a #)) -> Accessor m a)
-> (Fin m -> (# a #)) -> Accessor m a
forall a b. (a -> b) -> a -> b
$ \Fin m
i -> Vec n a -> Fin n -> (# a #)
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 :: SInt n -> [a] -> Vec n a
listVec SInt n
n [a]
xs = SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
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 ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> (([a] -> ST s ()) -> [a] -> ST s ()
forall a b. (a -> b) -> a -> b
$ [a]
xs) (([a] -> ST s ()) -> ST s ()) -> ([a] -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ SInt n
-> (Fin n -> ([a] -> ST s ()) -> [a] -> ST s ())
-> ([a] -> ST s ())
-> [a]
-> ST s ()
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
      [] -> MutableVec s n a -> Fin n -> a -> ST s ()
forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i a
forall a. a
uninitialized ST s () -> ST s () -> ST s ()
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> [a] -> ST s ()
rest []
      (a
x:[a]
xs'') -> MutableVec s n a -> Fin n -> a -> ST s ()
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 ST s () -> ST s () -> ST s ()
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> [a] -> ST s ()
rest [a]
xs'')
  (\[a]
_ -> () -> ST s ()
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 :: [a] -> (forall (n :: Nat). Vec n a -> r) -> r
withVec [a]
xs forall (n :: Nat). Vec n a -> r
f = Int -> (forall (n :: Nat). SInt n -> r) -> r
forall r.
(?callStack::CallStack) =>
Int -> (forall (n :: Nat). SInt n -> r) -> r
withSInt ([a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
xs) ((forall (n :: Nat). SInt n -> r) -> r)
-> (forall (n :: Nat). SInt n -> r) -> r
forall a b. (a -> b) -> a -> b
$ \SInt n
sn -> Vec n a -> r
forall (n :: Nat). Vec n a -> r
f (SInt n -> [a] -> Vec n a
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 :: SInt n -> Int -> [a] -> Vec n a
fromListN SInt n
sn Int
l [a]
xs
    | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n    = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn [a]
xs
    | Bool
otherwise = [Char] -> Vec n a
forall a. (?callStack::CallStack) => [Char] -> a
error ([Char] -> Vec n a) -> [Char] -> Vec n a
forall a b. (a -> b) -> a -> b
$ [Char]
"Vec.fromListN: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
l [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" /= " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
    where
    !n :: Int
n = SInt n -> Int
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 :: SInt n -> [a] -> Vec n a
fromList SInt n
sn [a]
xs
    | Int
n Int -> [a] -> Bool
forall a. Int -> [a] -> Bool
`eqLength` [a]
xs = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn [a]
xs
    | Bool
otherwise       = [Char] -> Vec n a
forall a. (?callStack::CallStack) => [Char] -> a
error ([Char] -> Vec n a) -> [Char] -> Vec n a
forall a b. (a -> b) -> a -> b
$ [Char]
"Vec.fromList: length /= " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
    where
    !n :: Int
n = SInt n -> Int
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 :: 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) = Int -> [a] -> Bool
forall a. Int -> [a] -> Bool
eqLength (Int
n Int -> Int -> Int
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 = SInt n -> Int -> [a] -> Vec n a
forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> Int -> [a] -> Vec n a
fromListN SInt n
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal
    fromList :: [Item (Vec n a)] -> Vec n a
fromList  = SInt n -> [a] -> Vec n a
forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList SInt n
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    = Vec n a -> [Item (Vec n a)]
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 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 = ([a] -> Vec n a) -> c ([a] -> Vec n a)
forall g. g -> c g
pur (SInt n -> [a] -> Vec n a
forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList SInt n
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal) c ([a] -> Vec n a) -> [a] -> c (Vec n a)
forall d b. Data d => c (d -> b) -> d -> c b
`app` Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList Vec n a
v
    gunfold :: (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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = c ([a] -> Vec n a) -> c (Vec n a)
forall b r. Data b => c (b -> r) -> c r
app (([a] -> Vec n a) -> c ([a] -> Vec n a)
forall r. r -> c r
pur (SInt n -> [a] -> Vec n a
forall (n :: Nat) a.
(?callStack::CallStack) =>
SInt n -> [a] -> Vec n a
fromList SInt n
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal))
        | Bool
otherwise            = [Char] -> c (Vec n a)
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 -> [Char] -> [Char]
showsPrec Int
p Vec n a
xs = Bool -> ([Char] -> [Char]) -> [Char] -> [Char]
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
precedence)
        (([Char] -> [Char]) -> [Char] -> [Char])
-> ([Char] -> [Char]) -> [Char] -> [Char]
forall a b. (a -> b) -> a -> b
$ [Char] -> [Char] -> [Char]
showString [Char]
"fromListN "
        ([Char] -> [Char]) -> ([Char] -> [Char]) -> [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Char] -> [Char]
forall a. Show a => a -> [Char] -> [Char]
shows (Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs)
        ([Char] -> [Char]) -> ([Char] -> [Char]) -> [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> [Char] -> [Char]
showString [Char]
" "
        ([Char] -> [Char]) -> ([Char] -> [Char]) -> [Char] -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [Char] -> [Char]
forall a. Show a => a -> [Char] -> [Char]
shows (Vec n a -> [a]
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 = Bool -> ReadS (Vec n a) -> ReadS (Vec n a)
forall a. Bool -> ReadS a -> ReadS a
readParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
precedence) (ReadS (Vec n a) -> ReadS (Vec n a))
-> ReadS (Vec n a) -> ReadS (Vec n a)
forall a b. (a -> b) -> a -> b
$ \[Char]
s ->
        [ Int -> (Vec n a, [Char]) -> (Vec n a, [Char])
forall b. Int -> b -> b
assertSize ([a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
ls) (SInt n -> Int -> [a] -> Vec n a
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'') <- Int -> ReadS Int
forall a. Read a => Int -> ReadS a
readsPrec Int
precedence [Char]
s'
        , ([a]
ls, [Char]
s''') <- Int -> (Int -> ReadS [a]) -> Int -> ReadS [a]
forall b. Int -> b -> b
assertSize Int
m Int -> ReadS [a]
forall a. Read a => Int -> ReadS a
readsPrec Int
precedence [Char]
s''
        ]
        where
            n :: SInt n
n = (?callStack::CallStack, KnownNat n) => SInt n
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal @n

            assertSize :: Int -> b -> b
            assertSize :: Int -> b -> b
assertSize Int
m b
x = if Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= SInt n -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt n
n
                then [Char] -> b
forall a. (?callStack::CallStack) => [Char] -> a
error ([Char] -> b) -> [Char] -> b
forall a b. (a -> b) -> a -> b
$ [Char]
"Can't read a Vec with " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
m [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
                    [Char]
" elements into a type `Vec " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> SInt n -> [Char]
forall a. Show a => a -> [Char]
show SInt n
n [Char] -> [Char] -> [Char]
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")
    [Int -> Portrayal
forall a. Portray a => a -> Portrayal
portray (Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs), [a] -> Portrayal
forall a. Portray a => a -> Portrayal
portray ([a] -> Portrayal) -> [a] -> Portrayal
forall a b. (a -> b) -> a -> b
$ Vec n a -> [a]
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 = ([a] -> [a] -> Maybe Portrayal
forall a. Diff a => a -> a -> Maybe Portrayal
diff ([a] -> [a] -> Maybe Portrayal)
-> (Vec n a -> [a]) -> Vec n a -> Vec n a -> Maybe Portrayal
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList) Vec n a
x Vec n a
y Maybe Portrayal -> (Portrayal -> Portrayal) -> Maybe Portrayal
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") [Int -> Portrayal
forall a. Portray a => a -> Portrayal
portray (Vec n a -> Int
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 = SInt n -> (Fin n -> (# () #)) -> ()
forall m (n :: Nat). Monoid m => SInt n -> (Fin n -> (# m #)) -> m
foldMapFin (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) ((Fin n -> (# () #)) -> ()) -> (Fin n -> (# () #)) -> ()
forall a b. (a -> b) -> a -> b
$
        \Fin n
i -> case Vec n a -> Fin n -> (a -> ()) -> ()
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i a -> ()
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
(<>) = (a -> a -> a) -> 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 a -> a -> a
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 = a -> Vec n a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
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 = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec SInt n
sn ([a] -> Vec n a) -> Gen [a] -> Gen (Vec n a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen a -> Gen [a]
forall a. Int -> Gen a -> Gen [a]
QC.vectorOf Int
n Gen a
forall a. Arbitrary a => Gen a
QC.arbitrary
      where
        !sn :: SInt n
sn@(SI# Int
n) = SInt 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 =
        [ Fin n -> Vec n a -> a -> Vec n a
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 <- SInt n -> (Fin n -> [Fin n] -> [Fin n]) -> [Fin n] -> [Fin n]
forall (n :: Nat) a. SInt n -> (Fin n -> a -> a) -> a -> a
foldrEnumFin SInt n
sn (:) [], a
x' <- Vec n a -> Fin n -> (a -> [a]) -> [a]
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i a -> [a]
forall a. Arbitrary a => a -> [a]
QC.shrink
        ]
      where
        !sn :: SInt n
sn = Vec n a -> SInt n
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 :: Fin n -> Vec n a -> a -> Vec n a
upd Fin n
i Vec n a
xs a
x = (forall s. ST s (Vec n a)) -> Vec n a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec n a)) -> Vec n a)
-> (forall s. ST s (Vec n a)) -> Vec n a
forall a b. (a -> b) -> a -> b
$ do
    MutableVec s n a
mv <- Vec n a -> ST s (MutableVec s n a)
forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV Vec n a
xs
    MutableVec s n a -> Fin n -> a -> ST s ()
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
    MutableVec s n a -> ST s (Vec n a)
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 :: Vec n a -> Gen b -> Gen b
coarbitrary = Vec n a -> Gen b -> Gen b
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
(+) = (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 a -> a -> a
forall a. Num a => a -> a -> a
(+)
    * :: Vec n a -> Vec n a -> Vec n 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 a -> a -> a
forall a. Num a => a -> a -> 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 (-)
    abs :: Vec n a -> Vec n a
abs = (a -> a) -> Vec n a -> Vec n a
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
abs
    signum :: Vec n a -> Vec n a
signum = (a -> a) -> Vec n a -> Vec n a
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
signum
    negate :: Vec n a -> Vec n a
negate = (a -> a) -> Vec n a -> Vec n a
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
negate
    fromInteger :: Integer -> Vec n a
fromInteger = a -> Vec n a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a -> Vec n a) -> (Integer -> a) -> Integer -> Vec n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger

instance (KnownNat n, Default a) => Default (Vec n a) where
    def :: Vec n a
def = a -> Vec n a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
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 (All -> Bool) -> All -> Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> All) -> Vec n Bool -> All
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Bool -> All
All (Vec n Bool -> All) -> Vec n Bool -> All
forall a b. (a -> b) -> a -> b
$ (a -> a -> Bool) -> Vec n a -> Vec n a -> Vec n Bool
forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 a -> a -> Bool
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 (Any -> Bool) -> Any -> Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> Any) -> Vec n Bool -> Any
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Bool -> Any
Any (Vec n Bool -> Any) -> Vec n Bool -> Any
forall a b. (a -> b) -> a -> b
$ (a -> a -> Bool) -> Vec n a -> Vec n a -> Vec n Bool
forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 a -> a -> Bool
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 = Vec n Ordering -> Ordering
forall (t :: Type -> Type) m. (Foldable t, Monoid m) => t m -> m
F.fold (Vec n Ordering -> Ordering) -> Vec n Ordering -> Ordering
forall a b. (a -> b) -> a -> b
$ (a -> a -> Ordering) -> Vec n a -> Vec n a -> Vec n Ordering
forall (f :: Type -> Type) a b c.
Apply f =>
(a -> b -> c) -> f a -> f b -> f c
liftF2 a -> a -> Ordering
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 :: (a -> b) -> Vec n a -> Vec n b
mapVec a -> b
f = Accessor n b -> Vec n b
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n b -> Vec n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Vec n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Accessor n a -> Accessor n b
forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f (Accessor n a -> Accessor n b)
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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 :: (a -> b) -> Vec n a -> Vec n b
fmap = (a -> b) -> Vec n a -> Vec n b
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 <$ :: a -> Vec n b -> Vec n a
<$ Vec n b
v = SInt n -> a -> Vec n a
forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec (Accessor n b -> SInt n
forall (n :: Nat) a. Accessor n a -> SInt n
_aSize (Vec n b -> Accessor n b
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 :: (Fin n -> a -> b) -> Vec n a -> Vec n b
imap = (Fin n -> a -> b) -> Vec n a -> Vec n b
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 :: (Fin n -> a -> m) -> Vec n a -> m
ifoldMap Fin n -> a -> m
f = Vec n m -> m
forall (t :: Type -> Type) m. (Foldable t, Monoid m) => t m -> m
F.fold (Vec n m -> m) -> (Vec n a -> Vec n m) -> Vec n a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Fin n -> a -> m) -> Vec n a -> Vec n m
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 :: (Fin n -> a -> f b) -> Vec n a -> f (Vec n b)
itraverse Fin n -> a -> f b
f = Vec n (f b) -> f (Vec n b)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA (Vec n (f b) -> f (Vec n b))
-> (Vec n a -> Vec n (f b)) -> Vec n a -> f (Vec n b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Fin n -> a -> f b) -> Vec n a -> Vec n (f b)
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' :: (a -> b) -> Vec n a -> Vec n b
map' a -> b
f = Accessor n b -> Vec n b
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n b -> Vec n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Vec n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Accessor n b -> Accessor n b
forall (n :: Nat) a. Accessor n a -> Accessor n a
seqVA (Accessor n b -> Accessor n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Accessor n a -> Accessor n b
forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> b
f (Accessor n a -> Accessor n b)
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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 :: (Fin n -> a -> b) -> Vec n a -> Vec n b
imap Fin n -> a -> b
f = Accessor n b -> Vec n b
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n b -> Vec n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Vec n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Fin n -> a -> b) -> Accessor n a -> Accessor n b
forall (n :: Nat) a b.
(Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f (Accessor n a -> Accessor n b)
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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 :: Vec n a -> Vec n (Fin n, a)
withPos = (Fin n -> a -> (Fin n, a)) -> Vec n a -> Vec n (Fin n, a)
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' :: (Fin n -> a -> b) -> Vec n a -> Vec n b
imap' Fin n -> a -> b
f = Accessor n b -> Vec n b
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n b -> Vec n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Vec n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Accessor n b -> Accessor n b
forall (n :: Nat) a. Accessor n a -> Accessor n a
seqVA (Accessor n b -> Accessor n b)
-> (Vec n a -> Accessor n b) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Fin n -> a -> b) -> Accessor n a -> Accessor n b
forall (n :: Nat) a b.
(Fin n -> a -> b) -> Accessor n a -> Accessor n b
mapWithPosVA Fin n -> a -> b
f (Accessor n a -> Accessor n b)
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
forall (n :: Nat) a. Vec n a -> Accessor n a
access
{-# INLINE imap' #-}

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

pureVec :: SInt n -> a -> Vec n a
pureVec = \SInt n
n a
x -> Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (SInt n -> a -> Accessor n a
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 :: (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 = Accessor n c -> Vec n c
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize ((a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c
forall a b c (n :: Nat).
(a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c
liftA2VA a -> b -> c
f (Vec n a -> Accessor n a
forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
as) (Vec n b -> Accessor n b
forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n b
bs))
{-# INLINE liftA2Vec #-}

instance Apply (Vec n) where
  liftF2 :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
liftF2 = (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
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 :: a -> Vec n a
pure = SInt n -> a -> Vec n a
forall (n :: Nat) a. SInt n -> a -> Vec n a
pureVec SInt n
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal
  {-# INLINE pure #-}

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

  <*> :: Vec n (a -> b) -> Vec n a -> Vec n b
(<*>) = ((a -> b) -> 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 (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)
  {-# INLINE (<*>) #-}

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

  Vec n a
xs <* :: 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 >>- :: Vec n a -> (a -> Vec n b) -> Vec n b
>>- a -> Vec n b
f = Accessor n b -> Vec n b
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (case Vec n a -> Accessor n a
forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
xs of
    Accessor SInt n
n Fin n -> (# a #)
get -> SInt n -> (Fin n -> (# b #)) -> Accessor n b
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 Vec n b -> Fin n -> (# b #)
forall (n :: Nat) a. Vec n a -> Fin n -> (# a #)
`fusibleFetch` Fin n
i))
  {-# INLINE (>>-) #-}

instance KnownNat n => Monad (Vec n) where >>= :: Vec n a -> (a -> Vec n b) -> Vec n 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 :: 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 = SInt n -> (Fin n -> c) -> Vec n c
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
tabulateVec SInt n
sn ((Fin n -> c) -> Vec n c) -> (Fin n -> c) -> Vec n c
forall a b. (a -> b) -> a -> b
$ \Fin n
i ->
    Vec n a -> Fin n -> (a -> c) -> c
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
xs Fin n
i ((a -> c) -> c) -> (a -> c) -> c
forall a b. (a -> b) -> a -> b
$ \a
x ->
    Vec n b -> Fin n -> (b -> c) -> c
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n b
ys Fin n
i ((b -> c) -> c) -> (b -> c) -> c
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 :: Vec n a -> Int -> (a -> r) -> r
unsafeIndexK (V# SmallArray# a
sa) (I# Int#
i) a -> r
k = case SmallArray# a -> Int# -> (# a #)
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 :: Vec n a -> Int
length = Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize
    {-# INLINE length #-}

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

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

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

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

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

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

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

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

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

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

    elem :: a -> Vec n a -> Bool
elem a
x = Any -> Bool
coerce (Any -> Bool) -> (Vec n a -> Any) -> Vec n a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Any) -> Vec n a -> Any
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Bool -> Any
Any (Bool -> Any) -> (a -> Bool) -> a -> Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool
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 :: (a -> f b) -> Vec n a -> f (Vec n b)
traverseVec a -> f b
f = Accessor n (f b) -> f (Vec n b)
forall (f :: Type -> Type) (n :: Nat) a.
Applicative f =>
Accessor n (f a) -> f (Vec n a)
sequenceVA (Accessor n (f b) -> f (Vec n b))
-> (Vec n a -> Accessor n (f b)) -> Vec n a -> f (Vec n b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> Accessor n a -> Accessor n (f b)
forall a b (n :: Nat). (a -> b) -> Accessor n a -> Accessor n b
mapVA a -> f b
f (Accessor n a -> Accessor n (f b))
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n (f b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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 :: (a -> f b) -> Vec n a -> f (Vec n b)
traverse = (a -> f b) -> Vec n a -> f (Vec n b)
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 :: f (Vec n a) -> Vec n (f a)
distribute f (Vec n a)
fxs = (Rep (Vec n) -> f a) -> Vec n (f a)
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate ((Rep (Vec n) -> f a) -> Vec n (f a))
-> (Rep (Vec n) -> f a) -> Vec n (f a)
forall a b. (a -> b) -> a -> b
$ \Rep (Vec n)
i -> (Vec n a -> a) -> f (Vec n a) -> f a
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vec n a -> Fin n -> a
forall (n :: Nat) a. Vec n a -> Fin n -> a
! Rep (Vec n)
Fin n
i) f (Vec n a)
fxs
    {-# INLINE distribute #-}

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

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

    index :: Vec n a -> Rep (Vec n) -> a
index = (!)
    {-# 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 :: (b -> a -> b) -> b -> Vec n a -> Vec (1 + n) b
vscanl b -> a -> b
f b
b Vec n a
v = SInt (1 + n) -> [b] -> Vec (1 + n) b
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (SInt 1
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal SInt 1 -> SInt n -> SInt (1 + n)
forall (m :: Nat) (n :: Nat).
(?callStack::CallStack) =>
SInt m -> SInt n -> SInt (m + n)
`addSInt` Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
v) ([b] -> Vec (1 + n) b) -> ([a] -> [b]) -> [a] -> Vec (1 + n) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a -> b) -> b -> [a] -> [b]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl b -> a -> b
f b
b ([a] -> Vec (1 + n) b) -> [a] -> Vec (1 + n) b
forall a b. (a -> b) -> a -> b
$ Vec n a -> [a]
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 :: Vec 0 a
nil = SInt 0 -> (Fin 0 -> a) -> Vec 0 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 0
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal (\Fin 0
i -> Fin 0
i Fin 0 -> a -> a
`seq` [Char] -> a
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_ :: Vec n a -> Vec m a -> Vec (n + m) a
append_ Vec n a
xs Vec m a
ys = (forall s. ST s (Vec (n + m) a)) -> Vec (n + m) a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec (n + m) a)) -> Vec (n + m) a)
-> (forall s. ST s (Vec (n + m) a)) -> Vec (n + m) a
forall a b. (a -> b) -> a -> b
$ do
    let !n :: Int
n = Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs
        !m :: Int
m = Vec m a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec m a
ys
    MutableVec s (n + m) a
zs <- Int -> a -> ST s (MutableVec s (n + m) a)
forall a s (n :: Nat). Int -> a -> ST s (MutableVec s n a)
unsafeNewMV (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m) a
forall a. a
uninitialized
    Vec n a -> Int -> MutableVec s (n + m) a -> Int -> Int -> ST s ()
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
    Vec m a -> Int -> MutableVec s (n + m) a -> Int -> Int -> ST s ()
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
    MutableVec s (n + m) a -> ST s (Vec (n + m) a)
forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s (n + m) a
zs
{-# NOINLINE append_ #-}

++ :: Vec n a -> Vec m a -> Vec (n + m) a
(++) = \Vec n a
xs Vec m a
ys -> Accessor (n + m) a -> Vec (n + m) a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n a -> Accessor m a -> Accessor (n + m) a
forall (n :: Nat) a (m :: Nat).
Accessor n a -> Accessor m a -> Accessor (n + m) a
appendVA (Vec n a -> Accessor n a
forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec n a
xs) (Vec m a -> Accessor m a
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_ :: SInt m -> Vec (m + n) a -> Vec m a
take_ SInt m
m Vec (m + n) a
xs = Vec (m + n) a -> Int -> SInt m -> Vec m a
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_ :: SInt m -> Vec (m + n) a -> Vec n a
drop_ SInt m
m Vec (m + n) a
xs =
  Vec (m + n) a -> Int -> SInt n -> Vec n a
forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec (m + n) a
xs (SInt m -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt m
m) (SInt n -> Vec n a) -> SInt n -> Vec n a
forall a b. (a -> b) -> a -> b
$
  Vec (m + n) a -> SInt (m + n)
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec (m + n) a
xs SInt (m + n) -> SInt m -> SInt n
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 :: 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 = Vec (m + n) a -> Accessor (m + n) a
forall (n :: Nat) a. Vec n a -> Accessor n a
access Vec (m + n) a
xs
  in  (Accessor m a -> Vec m a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (SInt m -> Accessor (m + n) a -> Accessor m a
forall (m :: Nat) (n :: Nat) a.
SInt m -> Accessor (m + n) a -> Accessor m a
takeVA SInt m
m Accessor (m + n) a
va), Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (SInt m -> Accessor (m + n) a -> Accessor n a
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 :: Vec n (Vec m a) -> Vec (n * m) a
concat Vec n (Vec m a)
xs =
  let !n :: Int
n = Vec n (Vec m a) -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n (Vec m a)
xs
  in  if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then
        -- Outer size is 0, return the empty array
        Int
-> (forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a
forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec Int
0 ((forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a)
-> (forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a
forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n * m) a
_ -> () -> ST s ()
forall (m :: Type -> Type) a. Monad m => a -> m a
return ()
      else
          let !m :: Int
m = Vec n (Vec m a) -> Int -> (Vec m a -> Int) -> Int
forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
0 Vec m a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize  -- We know the size is > 0.
          in  Int
-> (forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a
forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
m) ((forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a)
-> (forall s. MutableVec s (n * m) a -> ST s ()) -> Vec (n * m) a
forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n * m) a
marr ->
                [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: Type -> Type) (m :: Type -> Type) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
F.forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ Int
i ->
                  Vec n (Vec m a) -> Int -> (Vec m a -> ST s ()) -> ST s ()
forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
i ((Vec m a -> ST s ()) -> ST s ())
-> (Vec m a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ Vec m a
ys ->  -- i is [0..n-1], so in bounds.
                    Vec m a -> Int -> MutableVec s (n * m) a -> Int -> Int -> ST s ()
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 Int -> Int -> Int
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 :: SInt m -> Vec (n * m) a -> Vec n (Vec m a)
reshape SInt m
m =
  let !m' :: Int
m' = SInt m -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt m
m
  in  \Vec (n * m) a
xs ->
        SInt n -> (Fin n -> Vec m a) -> Vec n (Vec m a)
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec (Vec (n * m) a -> SInt (n * m)
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec (n * m) a
xs SInt (n * m) -> SInt m -> SInt n
forall (m :: Nat) (n :: Nat). SInt (m * n) -> SInt n -> SInt m
`divSIntR` SInt m
m) (\Fin n
i -> Vec (n * m) a -> Int -> SInt m -> Vec m a
forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec (n * m) a
xs (Fin n -> Int
forall (n :: Nat). Fin n -> Int
finToInt Fin n
i Int -> Int -> Int
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 :: (a -> Vec m b) -> Vec n a -> Vec (n * m) b
concatMap a -> Vec m b
f = Vec n (Vec m b) -> Vec (n * m) b
forall (m :: Nat) (n :: Nat) a. Vec n (Vec m a) -> Vec (n * m) a
concat (Vec n (Vec m b) -> Vec (n * m) b)
-> (Vec n a -> Vec n (Vec m b)) -> Vec n a -> Vec (n * m) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Vec m b) -> Vec n a -> Vec n (Vec m b)
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 :: SInt n -> (a -> a) -> a -> Vec n a
iterate SInt n
sn a -> a
f a
z =
    SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
sn ((forall s. MutableVec s n a -> ST s ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv ->
       SInt n -> (a -> Fin n -> ST s a) -> a -> ST s ()
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 a -> ST s () -> ST s a
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ MutableVec s n a -> Fin n -> a -> ST s ()
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' :: SInt n -> (a -> a) -> a -> Vec n a
iterate' SInt n
sn a -> a
f !a
z =
    SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
sn ((forall s. MutableVec s n a -> ST s ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv ->
        SInt n -> (a -> Fin n -> ST s a) -> a -> ST s ()
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 a -> ST s () -> ST s a
forall (f :: Type -> Type) a b. Functor f => a -> f b -> f a
<$ (MutableVec s n a -> Fin n -> a -> ST s ()
forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i (a -> ST s ()) -> a -> ST s ()
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 :: Vec n a -> Vec n a
rev = Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Accessor n a -> Vec n a)
-> (Vec n a -> Accessor n a) -> Vec n a -> Vec n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Accessor n a -> Accessor n a
forall (n :: Nat) a. Accessor n a -> Accessor n a
revVA (Accessor n a -> Accessor n a)
-> (Vec n a -> Accessor n a) -> Vec n a -> Accessor n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> Accessor n a
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_ :: Fin n -> Vec n a -> Vec n a
rot_ Fin n
o' Vec n a
xs = SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) ((forall s. MutableVec s n a -> ST s ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> do
  let o :: Int
o = Fin n -> Int
forall (n :: Nat). Fin n -> Int
finToInt Fin n
o'
      nmo :: Int
nmo = Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
o
  Vec n a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
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
  Vec n a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
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 :: Fin n -> Vec n a -> Vec n a
rot Fin n
o = \Vec n a
v -> Accessor n a -> Vec n a
forall (n :: Nat) a. Accessor n a -> Vec n a
materialize (Fin n -> Accessor n a -> Accessor n a
forall (n :: Nat) a. Fin n -> Accessor n a -> Accessor n a
rotVA Fin n
o (Vec n a -> Accessor n a
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 :: SInt n -> Vec n (Fin n)
viota SInt n
sn = SInt n -> (Fin n -> Fin n) -> Vec n (Fin n)
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt n
sn Fin n -> Fin n
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 :: Vec m a -> Vec n b -> Vec n (Vec m (a, b))
cross Vec m a
xs = (b -> Vec m (a, b)) -> Vec n b -> Vec n (Vec m (a, b))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (\b
y -> (a -> (a, b)) -> Vec m a -> Vec m (a, b)
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 :: Int -> a -> Vec n a -> Vec (n + 1) a
unsafeInsert Int
i a
xi Vec n a
xs =
  let !n :: Int
n = Vec n a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n a
xs
  in  Int
-> (forall s. MutableVec s (n + 1) a -> ST s ()) -> Vec (n + 1) a
forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ((forall s. MutableVec s (n + 1) a -> ST s ()) -> Vec (n + 1) a)
-> (forall s. MutableVec s (n + 1) a -> ST s ()) -> Vec (n + 1) a
forall a b. (a -> b) -> a -> b
$ \ MutableVec s (n + 1) a
mv -> do
        -- Explicitly: @mv[0..i-1] := xs[0..i-1]@
        Vec n a -> Int -> MutableVec s (n + 1) a -> Int -> Int -> ST s ()
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@
        MutableVec s (n + 1) a -> Int -> a -> ST s ()
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]@
        Vec n a -> Int -> MutableVec s (n + 1) a -> Int -> Int -> ST s ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
n Int -> Int -> Int
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 :: Fin (n + 1) -> a -> Vec n a -> Vec (n + 1) a
insert Fin (n + 1)
i = Int -> a -> Vec n a -> Vec (n + 1) a
forall a (n :: Nat). Int -> a -> Vec n a -> Vec (n + 1) a
unsafeInsert (Fin (n + 1) -> Int
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 :: Int -> Vec (n + 1) a -> Vec n a
unsafeRemove Int
i Vec (n + 1) a
xs =
  let !np1 :: Int
np1 = Vec (n + 1) a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec (n + 1) a
xs
  in  Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
unsafeCreateVec (Int
np1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ((forall s. MutableVec s n a -> ST s ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \ MutableVec s n a
mv -> do
        -- Explicitly: @mv[0..i-1] := xs[0..i-1]@
        Vec (n + 1) a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
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]@
        Vec (n + 1) a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
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
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) MutableVec s n a
mv Int
i (Int
np1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
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 :: Fin (n + 1) -> Vec (n + 1) a -> Vec n a
remove Fin (n + 1)
i = Int -> Vec (n + 1) a -> Vec n a
forall (n :: Nat) a. Int -> Vec (n + 1) a -> Vec n a
unsafeRemove (Fin (n + 1) -> Int
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 :: Vec n a -> Vec n a
vsort Vec n a
xs = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs) ([a] -> Vec n a) -> (Vec n a -> [a]) -> Vec n a -> Vec n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. Ord a => [a] -> [a]
L.sort ([a] -> [a]) -> (Vec n a -> [a]) -> Vec n a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList (Vec n a -> Vec n a) -> Vec n a -> Vec n a
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 :: (a -> a -> Ordering) -> Vec n a -> Vec n a
vsortBy a -> a -> Ordering
f Vec n a
xs = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs)([a] -> Vec n a) -> (Vec n a -> [a]) -> Vec n a -> Vec n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy a -> a -> Ordering
f ([a] -> [a]) -> (Vec n a -> [a]) -> Vec n a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList (Vec n a -> Vec n a) -> Vec n a -> Vec n a
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 :: (a -> b) -> Vec n a -> Vec n a
vsortOn a -> b
f Vec n a
xs = SInt n -> [a] -> Vec n a
forall (n :: Nat) a. SInt n -> [a] -> Vec n a
listVec (Vec n a -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n a
xs)([a] -> Vec n a) -> (Vec n a -> [a]) -> Vec n a -> Vec n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
L.sortOn a -> b
f ([a] -> [a]) -> (Vec n a -> [a]) -> Vec n a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList (Vec n a -> Vec n a) -> Vec n a -> Vec n a
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 :: 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 = Vec n (Vec m a) -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec n (Vec m a)
xs
      !t :: Int
t = SInt m -> Int
forall (n :: Nat). SInt n -> Int
unSInt SInt m
sm
  in  Int -> (Int -> Vec n a) -> Vec m (Vec n a)
forall a (n :: Nat). Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
t ((Int -> Vec n a) -> Vec m (Vec n a))
-> (Int -> Vec n a) -> Vec m (Vec n a)
forall a b. (a -> b) -> a -> b
$ \ Int
i ->
        -- s is the size of the outer vector, i.e. valueOf @n
        Int -> (Int -> a) -> Vec n a
forall a (n :: Nat). Int -> (Int -> a) -> Vec n a
unsafeMkVec Int
s ((Int -> a) -> Vec n a) -> (Int -> a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \ Int
j ->
          Vec n (Vec m a) -> Int -> (Vec m a -> a) -> a
forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec n (Vec m a)
xs Int
j ((Vec m a -> a) -> a) -> (Vec m a -> a) -> a
forall a b. (a -> b) -> a -> b
$ \ Vec m a
v -> Vec m a -> Int -> (a -> a) -> a
forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec m a
v Int
i a -> a
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 :: (a -> Bool) -> Vec n a -> Maybe (Fin n)
vfindIndex a -> Bool
p = (Int -> Fin n) -> Maybe Int -> Maybe (Fin n)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Fin n
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin (Maybe Int -> Maybe (Fin n))
-> (Vec n a -> Maybe Int) -> Vec n a -> Maybe (Fin n)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
L.findIndex a -> Bool
p ([a] -> Maybe Int) -> (Vec n a -> [a]) -> Vec n a -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vec n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
F.toList

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

-- | Create a singleton 'Vec'.
vec1 :: a -> Vec 1 a
vec1 :: a -> Vec 1 a
vec1 = a -> Vec 1 a
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 :: a -> a -> Vec 2 a
vec2 a
x0 a
x1 = SInt 2 -> (Fin 2 -> a) -> Vec 2 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 2
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal ((Fin 2 -> a) -> Vec 2 a) -> (Fin 2 -> a) -> Vec 2 a
forall a b. (a -> b) -> a -> b
$ \Fin 2
i -> case Fin 2 -> Int
forall (n :: Nat). Fin n -> Int
finToInt Fin 2
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
_ -> [Char] -> a
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 :: a -> a -> a -> Vec 3 a
vec3 a
x0 a
x1 a
x2 = SInt 3 -> (Fin 3 -> a) -> Vec 3 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 3
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal ((Fin 3 -> a) -> Vec 3 a) -> (Fin 3 -> a) -> Vec 3 a
forall a b. (a -> b) -> a -> b
$ \Fin 3
i -> case Fin 3 -> Int
forall (n :: Nat). Fin n -> Int
finToInt Fin 3
i of
  Int
0 -> a
x0
  Int
1 -> a
x1
  Int
2 -> a
x2
  Int
_ -> [Char] -> a
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 :: a -> a -> a -> a -> Vec 4 a
vec4 a
x0 a
x1 a
x2 a
x3 = SInt 4 -> (Fin 4 -> a) -> Vec 4 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 4
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal ((Fin 4 -> a) -> Vec 4 a) -> (Fin 4 -> a) -> Vec 4 a
forall a b. (a -> b) -> a -> b
$ \Fin 4
i -> case Fin 4 -> Int
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
_ -> [Char] -> a
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 :: a -> a -> a -> a -> a -> a -> Vec 6 a
vec6 a
x0 a
x1 a
x2 a
x3 a
x4 a
x5 = SInt 6 -> (Fin 6 -> a) -> Vec 6 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 6
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal ((Fin 6 -> a) -> Vec 6 a) -> (Fin 6 -> a) -> Vec 6 a
forall a b. (a -> b) -> a -> b
$ \Fin 6
i -> case Fin 6 -> Int
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
_ -> [Char] -> a
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 :: 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 = SInt 8 -> (Fin 8 -> a) -> Vec 8 a
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt 8
forall (n :: Nat). (?callStack::CallStack, KnownNat n) => SInt n
sintVal ((Fin 8 -> a) -> Vec 8 a) -> (Fin 8 -> a) -> Vec 8 a
forall a b. (a -> b) -> a -> b
$ \Fin 8
i -> case Fin 8 -> Int
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
_ -> [Char] -> a
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 :: i
valueOf = Natural -> i
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural -> i) -> Natural -> i
forall a b. (a -> b) -> a -> b
$ Proxy# n -> Natural
forall (n :: Nat). KnownNat n => Proxy# n -> Natural
natVal' (Proxy# n
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 :: Fin n -> (a -> a) -> Vec n a -> Vec n a
overIx Fin n
i a -> a
f Vec n a
v = (forall s. ST s (Vec n a)) -> Vec n a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec n a)) -> Vec n a)
-> (forall s. ST s (Vec n a)) -> Vec n a
forall a b. (a -> b) -> a -> b
$ do
  MutableVec s n a
mv <- Vec n a -> ST s (MutableVec s n a)
forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV Vec n a
v
  Vec n a -> Fin n -> (a -> ST s ()) -> ST s ()
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
v Fin n
i (MutableVec s n a -> Fin n -> a -> ST s ()
forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mv Fin n
i (a -> ST s ()) -> (a -> a) -> a -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f)
  MutableVec s n a -> ST s (Vec n a)
forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s n a
mv
{-# INLINE overIx #-}