{-# Language GADTs, DataKinds, TypeOperators, BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# Language PatternGuards #-}
{-# Language PolyKinds #-}
{-# Language TypeApplications, ScopedTypeVariables #-}
{-# Language TupleSections #-}
{-# Language Rank2Types, RoleAnnotations #-}
{-# Language CPP #-}
#if __GLASGOW_HASKELL__ >= 805
{-# Language NoStarIsType #-}
#endif
{-|
Copyright        : (c) Galois, Inc 2014-2019

A fixed-size vector of typed elements.

NB: This module contains an orphan instance. It will be included in GHC 8.10,
see https://gitlab.haskell.org/ghc/ghc/merge_requests/273.
-}
module Data.Parameterized.Vector
  ( Vector
    -- * Lists
  , fromList
  , toList

    -- * Assignments
  , fromAssignment
  , toAssignment

    -- * Length
  , length
  , nonEmpty
  , lengthInt

    -- * Indexing
  , elemAt
  , elemAtMaybe
  , elemAtUnsafe

    -- * Indexing with Fin
  , indicesUpTo
  , indicesOf

    -- * Update
  , insertAt
  , insertAtMaybe

    -- * Sub sequences
  , uncons
  , unsnoc
  , slice
  , Data.Parameterized.Vector.take
  , replace
  , mapAt
  , mapAtM

    -- * Zipping
  , zipWith
  , zipWithM
  , zipWithM_
  , interleave

    -- * Reorder
  , shuffle
  , reverse
  , rotateL
  , rotateR
  , shiftL
  , shiftR

    -- * Construction
  , singleton
  , cons
  , snoc
  , generate
  , generateM
  -- ** Unfolding
  , unfoldr
  , unfoldrM
  , unfoldrWithIndex
  , unfoldrWithIndexM
  , iterateN
  , iterateNM

    -- * Splitting and joining
    -- ** General
  , joinWithM
  , joinWith
  , splitWith
  , splitWithA

    -- ** Vectors
  , split
  , join
  , append

  ) where

import qualified Data.Vector as Vector
import Data.Coerce
import Data.Foldable.WithIndex (FoldableWithIndex(ifoldMap))
import Data.Functor.Compose
import Data.Functor.WithIndex (FunctorWithIndex(imap))
import Data.Vector.Mutable (MVector)
import qualified Data.Vector.Mutable as MVector
import Control.Monad.ST
import Data.Functor.Identity
import Data.Parameterized.Fin
import Data.Parameterized.NatRepr
import Data.Parameterized.NatRepr.Internal
import Data.Proxy
import Data.Traversable.WithIndex (TraversableWithIndex(itraverse))
import Prelude hiding (length,reverse,zipWith)
import Numeric.Natural

import qualified Data.Parameterized.Context as Ctx
import Data.Parameterized.Utils.Endian

-- | Fixed-size non-empty vectors.
data Vector n a where
  Vector :: (1 <= n) => !(Vector.Vector a) -> Vector n a

type role Vector nominal representational

instance Eq a => Eq (Vector n a) where
  (Vector Vector a
x) == :: Vector n a -> Vector n a -> Bool
== (Vector Vector a
y) = Vector a
x Vector a -> Vector a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a
y

instance Show a => Show (Vector n a) where
  show :: Vector n a -> String
show (Vector Vector a
x) = Vector a -> String
forall a. Show a => a -> String
show Vector a
x

-- | Get the elements of the vector as a list, lowest index first.
toList :: Vector n a -> [a]
toList :: Vector n a -> [a]
toList (Vector Vector a
v) = Vector a -> [a]
forall a. Vector a -> [a]
Vector.toList Vector a
v
{-# Inline toList #-}

-- NOTE: We are using the raw 'NatRepr' constructor here, which is unsafe.
-- | Length of the vector.
-- @O(1)@
length :: Vector n a -> NatRepr n
length :: Vector n a -> NatRepr n
length (Vector Vector a
xs) = Natural -> NatRepr n
forall (n :: Nat). Natural -> NatRepr n
NatRepr (Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs) :: Natural)
{-# INLINE length #-}

-- | The length of the vector as an "Int".
lengthInt :: Vector n a -> Int
lengthInt :: Vector n a -> Int
lengthInt (Vector Vector a
xs) = Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs
{-# Inline lengthInt #-}

elemAt :: ((i+1) <= n) => NatRepr i -> Vector n a -> a
elemAt :: NatRepr i -> Vector n a -> a
elemAt NatRepr i
n (Vector Vector a
xs) = Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! NatRepr i -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr i
n

-- | Get the element at the given index.
-- @O(1)@
elemAtMaybe :: Int -> Vector n a -> Maybe a
elemAtMaybe :: Int -> Vector n a -> Maybe a
elemAtMaybe Int
n (Vector Vector a
xs) = Vector a
xs Vector a -> Int -> Maybe a
forall a. Vector a -> Int -> Maybe a
Vector.!? Int
n
{-# INLINE elemAt #-}

-- | Get the element at the given index.
-- Raises an exception if the element is not in the vector's domain.
-- @O(1)@
elemAtUnsafe :: Int -> Vector n a -> a
elemAtUnsafe :: Int -> Vector n a -> a
elemAtUnsafe Int
n (Vector Vector a
xs) = Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! Int
n
{-# INLINE elemAtUnsafe #-}

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

indicesUpTo :: NatRepr n -> Vector (n + 1) (Fin (n + 1))
indicesUpTo :: NatRepr n -> Vector (n + 1) (Fin (n + 1))
indicesUpTo NatRepr n
n =
  NatRepr n
-> (Fin (n + 1) -> Fin (n + 1))
-> Fin (n + 1)
-> Vector (n + 1) (Fin (n + 1))
forall (n :: Nat) a. NatRepr n -> (a -> a) -> a -> Vector (n + 1) a
iterateN
    NatRepr n
n
    ((forall (i :: Nat).
 ((i + 1) <= (n + 1)) =>
 NatRepr i -> Fin (n + 1))
-> Fin (n + 1) -> Fin (n + 1)
forall (n :: Nat) r.
(forall (i :: Nat). ((i + 1) <= n) => NatRepr i -> r) -> Fin n -> r
viewFin
      (\NatRepr i
x ->
        case NatRepr (i + 1)
-> NatRepr (n + 1)
-> Either (LeqProof ((i + 1) + 1) (n + 1)) ((i + 1) :~: (n + 1))
forall (m :: Nat) (n :: Nat).
(m <= n) =>
NatRepr m -> NatRepr n -> Either (LeqProof (m + 1) n) (m :~: n)
testStrictLeq (NatRepr i -> NatRepr (i + 1)
forall (n :: Nat). NatRepr n -> NatRepr (n + 1)
incNat NatRepr i
x) (NatRepr n -> NatRepr (n + 1)
forall (n :: Nat). NatRepr n -> NatRepr (n + 1)
incNat NatRepr n
n) of
          Left LeqProof ((i + 1) + 1) (n + 1)
LeqProof -> NatRepr (i + 1) -> Fin (n + 1)
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
mkFin (NatRepr i -> NatRepr (i + 1)
forall (n :: Nat). NatRepr n -> NatRepr (n + 1)
incNat NatRepr i
x)
          Right (i + 1) :~: (n + 1)
Refl -> NatRepr n -> Fin (i + 1)
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
mkFin NatRepr n
n))
    (case NatRepr n -> NatRepr 1 -> LeqProof 1 (n + 1)
forall (f :: Nat -> Type) (m :: Nat) (g :: Nat -> Type) (n :: Nat).
f m -> g n -> LeqProof n (m + n)
addPrefixIsLeq NatRepr n
n (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) of
       LeqProof 1 (n + 1)
LeqProof -> NatRepr 0 -> Fin (n + 1)
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
mkFin (KnownNat 0 => NatRepr 0
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @0))

indicesOf :: Vector n a -> Vector n (Fin n)
indicesOf :: Vector n a -> Vector n (Fin n)
indicesOf v :: Vector n a
v@(Vector Vector a
_) = -- Pattern match to bring 1 <= n into scope
  case NatRepr n -> NatRepr 1 -> ((n - 1) + 1) :~: n
forall (f :: Nat -> Type) (m :: Nat) (g :: Nat -> Type) (n :: Nat).
(n <= m) =>
f m -> g n -> ((m - n) + n) :~: m
minusPlusCancel (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v) (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) of
    ((n - 1) + 1) :~: n
Refl -> NatRepr (n - 1) -> Vector ((n - 1) + 1) (Fin ((n - 1) + 1))
forall (n :: Nat). NatRepr n -> Vector (n + 1) (Fin (n + 1))
indicesUpTo (NatRepr n -> NatRepr (n - 1)
forall (n :: Nat). (1 <= n) => NatRepr n -> NatRepr (n - 1)
decNat (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v))

instance FunctorWithIndex (Fin n) (Vector n) where
  imap :: (Fin n -> a -> b) -> Vector n a -> Vector n b
imap Fin n -> a -> b
f Vector n a
v = (Fin n -> a -> b) -> Vector n (Fin n) -> Vector n a -> Vector n b
forall a b c (n :: Nat).
(a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
zipWith Fin n -> a -> b
f (Vector n a -> Vector n (Fin n)
forall (n :: Nat) a. Vector n a -> Vector n (Fin n)
indicesOf Vector n a
v) Vector n a
v

instance FoldableWithIndex (Fin n) (Vector n) where
  ifoldMap :: (Fin n -> a -> m) -> Vector n a -> m
ifoldMap Fin n -> a -> m
f Vector n a
v = ((Fin n, a) -> m) -> Vector n (Fin n, a) -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Fin n -> a -> m) -> (Fin n, a) -> m
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Fin n -> a -> m
f) ((Fin n -> a -> (Fin n, a)) -> Vector n a -> Vector n (Fin n, a)
forall i (f :: Type -> Type) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap (,) Vector n a
v)

instance TraversableWithIndex (Fin n) (Vector n) where
  itraverse :: (Fin n -> a -> f b) -> Vector n a -> f (Vector n b)
itraverse Fin n -> a -> f b
f Vector n a
v = ((Fin n, a) -> f b) -> Vector n (Fin n, a) -> f (Vector n b)
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Fin n -> a -> f b) -> (Fin n, a) -> f b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Fin n -> a -> f b
f) ((Fin n -> a -> (Fin n, a)) -> Vector n a -> Vector n (Fin n, a)
forall i (f :: Type -> Type) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap (,) Vector n a
v)

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

-- | Insert an element at the given index.
-- @O(n)@.
insertAt :: ((i + 1) <= n) => NatRepr i -> a -> Vector n a -> Vector n a
insertAt :: NatRepr i -> a -> Vector n a -> Vector n a
insertAt NatRepr i
n a
a (Vector Vector a
xs) = Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
Vector.unsafeUpd Vector a
xs [(NatRepr i -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr i
n,a
a)])

-- | Insert an element at the given index.
-- Return 'Nothing' if the element is outside the vector bounds.
-- @O(n)@.
insertAtMaybe :: Int -> a -> Vector n a -> Maybe (Vector n a)
insertAtMaybe :: Int -> a -> Vector n a -> Maybe (Vector n a)
insertAtMaybe Int
n a
a (Vector Vector a
xs)
  | Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs = Vector n a -> Maybe (Vector n a)
forall a. a -> Maybe a
Just (Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
Vector.unsafeUpd Vector a
xs [(Int
n,a
a)]))
  | Bool
otherwise = Maybe (Vector n a)
forall a. Maybe a
Nothing


-- | Proof that the length of this vector is not 0.
nonEmpty :: Vector n a -> LeqProof 1 n
nonEmpty :: Vector n a -> LeqProof 1 n
nonEmpty (Vector Vector a
_) = LeqProof 1 n
forall (m :: Nat) (n :: Nat). (m <= n) => LeqProof m n
LeqProof
{-# Inline nonEmpty #-}


-- | Remove the first element of the vector, and return the rest, if any.
uncons :: forall n a.  Vector n a -> (a, Either (n :~: 1) (Vector (n-1) a))
uncons :: Vector n a -> (a, Either (n :~: 1) (Vector (n - 1) a))
uncons v :: Vector n a
v@(Vector Vector a
xs) = (Vector a -> a
forall a. Vector a -> a
Vector.head Vector a
xs, Either (n :~: 1) (Vector (n - 1) a)
mbTail)
  where
  mbTail :: Either (n :~: 1) (Vector (n - 1) a)
  mbTail :: Either (n :~: 1) (Vector (n - 1) a)
mbTail = case NatRepr 1 -> NatRepr n -> Either (LeqProof (1 + 1) n) (1 :~: n)
forall (m :: Nat) (n :: Nat).
(m <= n) =>
NatRepr m -> NatRepr n -> Either (LeqProof (m + 1) n) (m :~: n)
testStrictLeq (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v) of
             Left LeqProof (1 + 1) n
n2_leq_n ->
               do LeqProof 1 (n - 1)
LeqProof <- LeqProof 1 (n - 1) -> Either (n :~: 1) (LeqProof 1 (n - 1))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (LeqProof 2 n -> LeqProof 1 1 -> LeqProof (2 - 1) (n - 1)
forall (x_l :: Nat) (x_h :: Nat) (y_l :: Nat) (y_h :: Nat).
LeqProof x_l x_h
-> LeqProof y_l y_h -> LeqProof (x_l - y_h) (x_h - y_l)
leqSub2 LeqProof 2 n
LeqProof (1 + 1) n
n2_leq_n (NatRepr 1 -> LeqProof 1 1
forall (f :: Nat -> Type) (n :: Nat). f n -> LeqProof n n
leqRefl (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1)))
                  Vector (n - 1) a -> Either (n :~: 1) (Vector (n - 1) a)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Vector a -> Vector (n - 1) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a -> Vector a
forall a. Vector a -> Vector a
Vector.tail Vector a
xs))
             Right 1 :~: n
Refl    -> (n :~: n) -> Either (n :~: n) (Vector 0 a)
forall a b. a -> Either a b
Left n :~: n
forall k (a :: k). a :~: a
Refl
{-# Inline uncons #-}

-- | Remove the last element of the vector, and return the rest, if any.
unsnoc :: forall n a.  Vector n a -> (a, Either (n :~: 1) (Vector (n-1) a))
unsnoc :: Vector n a -> (a, Either (n :~: 1) (Vector (n - 1) a))
unsnoc v :: Vector n a
v@(Vector Vector a
xs) = (Vector a -> a
forall a. Vector a -> a
Vector.last Vector a
xs, Either (n :~: 1) (Vector (n - 1) a)
mbTail)
  where
  mbTail :: Either (n :~: 1) (Vector (n - 1) a)
  mbTail :: Either (n :~: 1) (Vector (n - 1) a)
mbTail = case NatRepr 1 -> NatRepr n -> Either (LeqProof (1 + 1) n) (1 :~: n)
forall (m :: Nat) (n :: Nat).
(m <= n) =>
NatRepr m -> NatRepr n -> Either (LeqProof (m + 1) n) (m :~: n)
testStrictLeq (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v) of
             Left LeqProof (1 + 1) n
n2_leq_n ->
               do LeqProof 1 (n - 1)
LeqProof <- LeqProof 1 (n - 1) -> Either (n :~: 1) (LeqProof 1 (n - 1))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (LeqProof 2 n -> LeqProof 1 1 -> LeqProof (2 - 1) (n - 1)
forall (x_l :: Nat) (x_h :: Nat) (y_l :: Nat) (y_h :: Nat).
LeqProof x_l x_h
-> LeqProof y_l y_h -> LeqProof (x_l - y_h) (x_h - y_l)
leqSub2 LeqProof 2 n
LeqProof (1 + 1) n
n2_leq_n (NatRepr 1 -> LeqProof 1 1
forall (f :: Nat -> Type) (n :: Nat). f n -> LeqProof n n
leqRefl (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1)))
                  Vector (n - 1) a -> Either (n :~: 1) (Vector (n - 1) a)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Vector a -> Vector (n - 1) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Int -> Int -> Vector a -> Vector a
forall a. Int -> Int -> Vector a -> Vector a
Vector.slice Int
0 (Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Vector a
xs))
             Right 1 :~: n
Refl    -> (n :~: n) -> Either (n :~: n) (Vector 0 a)
forall a b. a -> Either a b
Left n :~: n
forall k (a :: k). a :~: a
Refl
{-# Inline unsnoc #-}


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

-- | Make a vector of the given length and element type.
-- Returns "Nothing" if the input list does not have the right number of
-- elements.
-- @O(n)@.
fromList :: (1 <= n) => NatRepr n -> [a] -> Maybe (Vector n a)
fromList :: NatRepr n -> [a] -> Maybe (Vector n a)
fromList NatRepr n
n [a]
xs
  | NatRepr n -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr n
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
v = Vector n a -> Maybe (Vector n a)
forall a. a -> Maybe a
Just (Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
v)
  | Bool
otherwise                     = Maybe (Vector n a)
forall a. Maybe a
Nothing
  where
  v :: Vector a
v = [a] -> Vector a
forall a. [a] -> Vector a
Vector.fromList [a]
xs
{-# INLINE fromList #-}

-- | Convert a non-empty 'Ctx.Assignment' to a fixed-size 'Vector'.
--
-- This function uses the same ordering convention as 'Ctx.toVector'.
fromAssignment ::
  forall f ctx tp e.
  (forall tp'. f tp' -> e) ->
  Ctx.Assignment f (ctx Ctx.::> tp) ->
  Vector (Ctx.CtxSize (ctx Ctx.::> tp)) e
fromAssignment :: (forall (tp' :: k). f tp' -> e)
-> Assignment f (ctx ::> tp) -> Vector (CtxSize (ctx ::> tp)) e
fromAssignment forall (tp' :: k). f tp' -> e
f Assignment f (ctx ::> tp)
assign =
  case Assignment f (ctx ::> tp) -> AssignView f (ctx ::> tp)
forall k (f :: k -> Type) (ctx :: Ctx k).
Assignment f ctx -> AssignView f ctx
Ctx.viewAssign Assignment f (ctx ::> tp)
assign of
    Ctx.AssignExtend Assignment f ctx
assign' f tp
_ ->
      case LeqProof 1 1
-> NatRepr (CtxSize ctx) -> LeqProof 1 (1 + CtxSize ctx)
forall (f :: Nat -> Type) (m :: Nat) (n :: Nat) (p :: Nat).
LeqProof m n -> f p -> LeqProof m (n + p)
leqAdd (NatRepr 1 -> LeqProof 1 1
forall (f :: Nat -> Type) (n :: Nat). f n -> LeqProof n n
leqRefl (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1)) (Size ctx -> NatRepr (CtxSize ctx)
forall k (items :: Ctx k). Size items -> NatRepr (CtxSize items)
Ctx.sizeToNatRepr (Assignment f ctx -> Size ctx
forall k (f :: k -> Type) (ctx :: Ctx k).
Assignment f ctx -> Size ctx
Ctx.size Assignment f ctx
assign')) of
        LeqProof 1 (1 + CtxSize ctx)
LeqProof -> Vector e -> Vector (1 + CtxSize ctx) e
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Assignment f (ctx ::> tp)
-> (forall (tp' :: k). f tp' -> e) -> Vector e
forall k (f :: k -> Type) (tps :: Ctx k) e.
Assignment f tps -> (forall (tp :: k). f tp -> e) -> Vector e
Ctx.toVector Assignment f (ctx ::> tp)
assign forall (tp' :: k). f tp' -> e
f)

-- | Convert a 'Vector' into a 'Ctx.Assignment'.
--
-- This function uses the same ordering convention as 'Ctx.toVector'.
toAssignment ::
  Ctx.Size ctx ->
  (forall tp. Ctx.Index ctx tp -> e -> f tp) ->
  Vector (Ctx.CtxSize ctx) e ->
  Ctx.Assignment f ctx
toAssignment :: Size ctx
-> (forall (tp :: k). Index ctx tp -> e -> f tp)
-> Vector (CtxSize ctx) e
-> Assignment f ctx
toAssignment Size ctx
sz forall (tp :: k). Index ctx tp -> e -> f tp
g Vector (CtxSize ctx) e
vec =
  -- The unsafe indexing here relies on the safety of the rest of the Vector
  -- API, specifically the inability to construct vectors that have an
  -- underlying size that differs from the size in their type.
  Size ctx
-> (forall (tp :: k). Index ctx tp -> f tp) -> Assignment f ctx
forall k (ctx :: Ctx k) (f :: k -> Type).
Size ctx
-> (forall (tp :: k). Index ctx tp -> f tp) -> Assignment f ctx
Ctx.generate Size ctx
sz (\Index ctx tp
idx -> Index ctx tp -> e -> f tp
forall (tp :: k). Index ctx tp -> e -> f tp
g Index ctx tp
idx (Int -> Vector (CtxSize ctx) e -> e
forall (n :: Nat) a. Int -> Vector n a -> a
elemAtUnsafe (Index ctx tp -> Int
forall k (ctx :: Ctx k) (tp :: k). Index ctx tp -> Int
Ctx.indexVal Index ctx tp
idx) Vector (CtxSize ctx) e
vec))

-- | Extract a subvector of the given vector.
slice :: (i + w <= n, 1 <= w) =>
            NatRepr i {- ^ Start index -} ->
            NatRepr w {- ^ Width of sub-vector -} ->
            Vector n a -> Vector w a
slice :: NatRepr i -> NatRepr w -> Vector n a -> Vector w a
slice NatRepr i
i NatRepr w
w (Vector Vector a
xs) = Vector a -> Vector w a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Int -> Int -> Vector a -> Vector a
forall a. Int -> Int -> Vector a -> Vector a
Vector.slice (NatRepr i -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr i
i) (NatRepr w -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr w
w) Vector a
xs)
{-# INLINE slice #-}

-- | Take the front (lower-indexes) part of the vector.
take :: forall n x a. (1 <= n) => NatRepr n -> Vector (n + x) a -> Vector n a
take :: NatRepr n -> Vector (n + x) a -> Vector n a
take | LeqProof n (n + x)
LeqProof <- LeqProof n (n + x)
prf = NatRepr 0 -> NatRepr n -> Vector (n + x) a -> Vector n a
forall (i :: Nat) (w :: Nat) (n :: Nat) a.
((i + w) <= n, 1 <= w) =>
NatRepr i -> NatRepr w -> Vector n a -> Vector w a
slice (KnownNat 0 => NatRepr 0
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @0)
  where
  prf :: LeqProof n (n + x)
prf = LeqProof n n -> Proxy x -> LeqProof n (n + x)
forall (f :: Nat -> Type) (m :: Nat) (n :: Nat) (p :: Nat).
LeqProof m n -> f p -> LeqProof m (n + p)
leqAdd (Proxy n -> LeqProof n n
forall (f :: Nat -> Type) (n :: Nat). f n -> LeqProof n n
leqRefl (Proxy n
forall k (t :: k). Proxy t
Proxy @n)) (Proxy x
forall k (t :: k). Proxy t
Proxy @x)

-- | Scope a monadic function to a sub-section of the given vector.
mapAtM :: Monad m => (i + w <= n, 1 <= w) =>
            NatRepr i {- ^ Start index -} ->
            NatRepr w {- ^ Section width -} ->
            (Vector w a -> m (Vector w a)) {-^ map for the sub-vector -} ->
            Vector n a -> m (Vector n a)
mapAtM :: NatRepr i
-> NatRepr w
-> (Vector w a -> m (Vector w a))
-> Vector n a
-> m (Vector n a)
mapAtM NatRepr i
i NatRepr w
w Vector w a -> m (Vector w a)
f (Vector Vector a
vn) =
  let
    (Vector a
vhead, Vector a
vtail) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
Vector.splitAt (NatRepr i -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr i
i) Vector a
vn
    (Vector a
vsect, Vector a
vend) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
Vector.splitAt (NatRepr w -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr w
w) Vector a
vtail
  in do
    Vector Vector a
vsect' <- Vector w a -> m (Vector w a)
f (Vector a -> Vector w a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
vsect)
    Vector n a -> m (Vector n a)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Vector n a -> m (Vector n a)) -> Vector n a -> m (Vector n a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a -> Vector n a) -> Vector a -> Vector n a
forall a b. (a -> b) -> a -> b
$ Vector a
vhead Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
Vector.++ Vector a
vsect' Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
Vector.++ Vector a
vend

-- | Scope a function to a sub-section of the given vector.
mapAt :: (i + w <= n, 1 <= w) =>
            NatRepr i {- ^ Start index -} ->
            NatRepr w {- ^ Section width -} ->
            (Vector w a -> Vector w a) {-^ map for the sub-vector -} ->
            Vector n a -> Vector n a
mapAt :: NatRepr i
-> NatRepr w
-> (Vector w a -> Vector w a)
-> Vector n a
-> Vector n a
mapAt NatRepr i
i NatRepr w
w Vector w a -> Vector w a
f Vector n a
vn = Identity (Vector n a) -> Vector n a
forall a. Identity a -> a
runIdentity (Identity (Vector n a) -> Vector n a)
-> Identity (Vector n a) -> Vector n a
forall a b. (a -> b) -> a -> b
$ NatRepr i
-> NatRepr w
-> (Vector w a -> Identity (Vector w a))
-> Vector n a
-> Identity (Vector n a)
forall (m :: Type -> Type) (i :: Nat) (w :: Nat) (n :: Nat) a.
(Monad m, (i + w) <= n, 1 <= w) =>
NatRepr i
-> NatRepr w
-> (Vector w a -> m (Vector w a))
-> Vector n a
-> m (Vector n a)
mapAtM NatRepr i
i NatRepr w
w (Vector w a -> Identity (Vector w a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Vector w a -> Identity (Vector w a))
-> (Vector w a -> Vector w a)
-> Vector w a
-> Identity (Vector w a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector w a -> Vector w a
f) Vector n a
vn

-- | Replace a sub-section of a vector with the given sub-vector.
replace :: (i + w <= n, 1 <= w) =>
              NatRepr i {- ^ Start index -} ->
              Vector w a {- ^ sub-vector -} ->
              Vector n a -> Vector n a
replace :: NatRepr i -> Vector w a -> Vector n a -> Vector n a
replace NatRepr i
i Vector w a
vw Vector n a
vn = NatRepr i
-> NatRepr w
-> (Vector w a -> Vector w a)
-> Vector n a
-> Vector n a
forall (i :: Nat) (w :: Nat) (n :: Nat) a.
((i + w) <= n, 1 <= w) =>
NatRepr i
-> NatRepr w
-> (Vector w a -> Vector w a)
-> Vector n a
-> Vector n a
mapAt NatRepr i
i (Vector w a -> NatRepr w
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector w a
vw) (Vector w a -> Vector w a -> Vector w a
forall a b. a -> b -> a
const Vector w a
vw) Vector n a
vn

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

instance Functor (Vector n) where
  fmap :: (a -> b) -> Vector n a -> Vector n b
fmap a -> b
f (Vector Vector a
xs) = Vector b -> Vector n b
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector ((a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
Vector.map a -> b
f Vector a
xs)
  {-# Inline fmap #-}

instance Foldable (Vector n) where
  foldMap :: (a -> m) -> Vector n a -> m
foldMap a -> m
f (Vector Vector a
xs) = (a -> m) -> Vector a -> m
forall (t :: Type -> Type) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Vector a
xs

instance Traversable (Vector n) where
  traverse :: (a -> f b) -> Vector n a -> f (Vector n b)
traverse a -> f b
f (Vector Vector a
xs) = Vector b -> Vector n b
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector b -> Vector n b) -> f (Vector b) -> f (Vector n b)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Vector a -> f (Vector b)
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Vector a
xs
  {-# Inline traverse #-}

-- | Zip two vectors, potentially changing types.
-- @O(n)@
zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
zipWith a -> b -> c
f (Vector Vector a
xs) (Vector Vector b
ys) = Vector c -> Vector n c
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector ((a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
Vector.zipWith a -> b -> c
f Vector a
xs Vector b
ys)
{-# Inline zipWith #-}

zipWithM :: Monad m => (a -> b -> m c) ->
                       Vector n a -> Vector n b -> m (Vector n c)
zipWithM :: (a -> b -> m c) -> Vector n a -> Vector n b -> m (Vector n c)
zipWithM a -> b -> m c
f (Vector Vector a
xs) (Vector Vector b
ys) = Vector c -> Vector n c
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector c -> Vector n c) -> m (Vector c) -> m (Vector n c)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
forall (m :: Type -> Type) a b c.
Monad m =>
(a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
Vector.zipWithM a -> b -> m c
f Vector a
xs Vector b
ys
{-# Inline zipWithM #-}

zipWithM_ :: Monad m => (a -> b -> m ()) -> Vector n a -> Vector n b -> m ()
zipWithM_ :: (a -> b -> m ()) -> Vector n a -> Vector n b -> m ()
zipWithM_ a -> b -> m ()
f (Vector Vector a
xs) (Vector Vector b
ys) = (a -> b -> m ()) -> Vector a -> Vector b -> m ()
forall (m :: Type -> Type) a b c.
Monad m =>
(a -> b -> m c) -> Vector a -> Vector b -> m ()
Vector.zipWithM_ a -> b -> m ()
f Vector a
xs Vector b
ys
{-# Inline zipWithM_ #-}

{- | Interleave two vectors.  The elements of the first vector are
at even indexes in the result, the elements of the second are at odd indexes. -}
interleave ::
  forall n a. (1 <= n) => Vector n a -> Vector n a -> Vector (2 * n) a
interleave :: Vector n a -> Vector n a -> Vector (2 * n) a
interleave (Vector Vector a
xs) (Vector Vector a
ys)
  | LeqProof 1 (2 * n)
LeqProof <- Proxy 2 -> Proxy n -> LeqProof 1 (2 * n)
forall (p :: Nat -> Type) (q :: Nat -> Type) (x :: Nat) (y :: Nat).
(1 <= x, 1 <= y) =>
p x -> q y -> LeqProof 1 (x * y)
leqMulPos (Proxy 2
forall k (t :: k). Proxy t
Proxy @2) (Proxy n
forall k (t :: k). Proxy t
Proxy @n) = Vector a -> Vector (2 * n) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
zs
  where
  len :: Int
len = Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
ys
  zs :: Vector a
zs  = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
Vector.generate Int
len (\Int
i -> let v :: Vector a
v = if Int -> Bool
forall a. Integral a => a -> Bool
even Int
i then Vector a
xs else Vector a
ys
                                   in Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2))

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

{- | Move the elements around, as specified by the given function.
  * Note: the reindexing function says where each of the elements
          in the new vector come from.
  * Note: it is OK for the same input element to end up in mulitple places
          in the result.
@O(n)@
-}
shuffle :: (Int -> Int) -> Vector n a -> Vector n a
shuffle :: (Int -> Int) -> Vector n a -> Vector n a
shuffle Int -> Int
f (Vector Vector a
xs) = Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
ys
  where
  ys :: Vector a
ys = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
Vector.generate (Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs) (\Int
i -> Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! Int -> Int
f Int
i)
{-# Inline shuffle #-}

-- | Reverse the vector.
reverse :: forall a n. (1 <= n) => Vector n a -> Vector n a
reverse :: Vector n a -> Vector n a
reverse Vector n a
x = (Int -> Int) -> Vector n a -> Vector n a
forall (n :: Nat) a. (Int -> Int) -> Vector n a -> Vector n a
shuffle (\Int
i -> Vector n a -> Int
forall (n :: Nat) a. Vector n a -> Int
lengthInt Vector n a
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Vector n a
x

-- | Rotate "left".  The first element of the vector is on the "left", so
-- rotate left moves all elemnts toward the corresponding smaller index.
-- Elements that fall off the beginning end up at the end.
rotateL :: Int -> Vector n a -> Vector n a
rotateL :: Int -> Vector n a -> Vector n a
rotateL !Int
n Vector n a
xs = (Int -> Int) -> Vector n a -> Vector n a
forall (n :: Nat) a. (Int -> Int) -> Vector n a -> Vector n a
shuffle Int -> Int
rotL Vector n a
xs
  where
  !len :: Int
len   = Vector n a -> Int
forall (n :: Nat) a. Vector n a -> Int
lengthInt Vector n a
xs
  rotL :: Int -> Int
rotL Int
i = (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
len          -- `len` is known to be >= 1
{-# Inline rotateL #-}

-- | Rotate "right".  The first element of the vector is on the "left", so
-- rotate right moves all elemnts toward the corresponding larger index.
-- Elements that fall off the end, end up at the beginning.
rotateR :: Int -> Vector n a -> Vector n a
rotateR :: Int -> Vector n a -> Vector n a
rotateR !Int
n Vector n a
xs = (Int -> Int) -> Vector n a -> Vector n a
forall (n :: Nat) a. (Int -> Int) -> Vector n a -> Vector n a
shuffle Int -> Int
rotR Vector n a
xs
  where
  !len :: Int
len   = Vector n a -> Int
forall (n :: Nat) a. Vector n a -> Int
lengthInt Vector n a
xs
  rotR :: Int -> Int
rotR Int
i = (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
len        -- `len` is known to be >= 1
{-# Inline rotateR #-}

{- | Move all elements towards smaller indexes.
Elements that fall off the front are ignored.
Empty slots are filled in with the given element.
@O(n)@. -}
shiftL :: Int -> a -> Vector n a -> Vector n a
shiftL :: Int -> a -> Vector n a -> Vector n a
shiftL !Int
x a
a (Vector Vector a
xs) = Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
ys
  where
  !len :: Int
len = Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs
  ys :: Vector a
ys   = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
Vector.generate Int
len (\Int
i -> let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
                                    in if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len then a
a else Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! Int
j)
{-# Inline shiftL #-}

{- | Move all elements towards the larger indexes.
Elements that "fall" off the end are ignored.
Empty slots are filled in with the given element.
@O(n)@. -}
shiftR :: Int -> a -> Vector n a -> Vector n a
shiftR :: Int -> a -> Vector n a -> Vector n a
shiftR !Int
x a
a (Vector Vector a
xs) = Vector a -> Vector n a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector Vector a
ys
  where
  !len :: Int
len = Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
xs
  ys :: Vector a
ys   = Int -> (Int -> a) -> Vector a
forall a. Int -> (Int -> a) -> Vector a
Vector.generate Int
len (\Int
i -> let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x
                                    in if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then a
a else Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! Int
j)
{-# Inline shiftR #-}

-------------------------------------------------------------------------------i

-- | Append two vectors. The first one is at lower indexes in the result.
append :: Vector m a -> Vector n a -> Vector (m + n) a
append :: Vector m a -> Vector n a -> Vector (m + n) a
append v1 :: Vector m a
v1@(Vector Vector a
xs) v2 :: Vector n a
v2@(Vector Vector a
ys) =
  case NatRepr m -> NatRepr n -> LeqProof 1 (m + n)
forall (m :: Nat) (n :: Nat) (p :: Nat -> Type) (q :: Nat -> Type).
(1 <= m, 1 <= n) =>
p m -> q n -> LeqProof 1 (m + n)
leqAddPos (Vector m a -> NatRepr m
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector m a
v1) (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v2) of { LeqProof 1 (m + n)
LeqProof ->
    Vector a -> Vector (m + n) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a
xs Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
Vector.++ Vector a
ys)
  }
{-# Inline append #-}

--------------------------------------------------------------------------------
-- Constructing Vectors

-- | Vector with exactly one element
singleton :: forall a. a -> Vector 1 a
singleton :: a -> Vector 1 a
singleton a
a = Vector a -> Vector 1 a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (a -> Vector a
forall a. a -> Vector a
Vector.singleton a
a)

leqLen :: forall n a. Vector n a -> LeqProof 1 (n + 1)
leqLen :: Vector n a -> LeqProof 1 (n + 1)
leqLen Vector n a
v =
  let leqSucc :: forall f z. f z -> LeqProof z (z + 1)
      leqSucc :: f z -> LeqProof z (z + 1)
leqSucc f z
fz = LeqProof z z -> NatRepr 1 -> LeqProof z (z + 1)
forall (f :: Nat -> Type) (m :: Nat) (n :: Nat) (p :: Nat).
LeqProof m n -> f p -> LeqProof m (n + p)
leqAdd (f z -> LeqProof z z
forall (f :: Nat -> Type) (n :: Nat). f n -> LeqProof n n
leqRefl f z
fz :: LeqProof z z) (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1)
  in LeqProof 1 n -> LeqProof n (n + 1) -> LeqProof 1 (n + 1)
forall (m :: Nat) (n :: Nat) (p :: Nat).
LeqProof m n -> LeqProof n p -> LeqProof m p
leqTrans (Vector n a -> LeqProof 1 n
forall (n :: Nat) a. Vector n a -> LeqProof 1 n
nonEmpty Vector n a
v :: LeqProof 1 n) (NatRepr n -> LeqProof n (n + 1)
forall (f :: Nat -> Type) (z :: Nat). f z -> LeqProof z (z + 1)
leqSucc (Vector n a -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector n a
v))

-- | Add an element to the head of a vector
cons :: forall n a. a -> Vector n a -> Vector (n+1) a
cons :: a -> Vector n a -> Vector (n + 1) a
cons a
a v :: Vector n a
v@(Vector Vector a
x) = case Vector n a -> LeqProof 1 (n + 1)
forall (n :: Nat) a. Vector n a -> LeqProof 1 (n + 1)
leqLen Vector n a
v of LeqProof 1 (n + 1)
LeqProof -> (Vector a -> Vector (n + 1) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (a -> Vector a -> Vector a
forall a. a -> Vector a -> Vector a
Vector.cons a
a Vector a
x))

-- | Add an element to the tail of a vector
snoc :: forall n a. Vector n a -> a -> Vector (n+1) a
snoc :: Vector n a -> a -> Vector (n + 1) a
snoc v :: Vector n a
v@(Vector Vector a
x) a
a = case Vector n a -> LeqProof 1 (n + 1)
forall (n :: Nat) a. Vector n a -> LeqProof 1 (n + 1)
leqLen Vector n a
v of LeqProof 1 (n + 1)
LeqProof -> (Vector a -> Vector (n + 1) a
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector (Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
Vector.snoc Vector a
x a
a))

-- | This newtype wraps Vector so that we can curry it in the call to
-- @natRecBounded@. It adds 1 to the length so that the base case is
-- a @Vector@ of non-zero length.
newtype Vector' a n = MkVector' (Vector (n+1) a)

unVector' :: Vector' a n -> Vector (n+1) a
unVector' :: Vector' a n -> Vector (n + 1) a
unVector' (MkVector' Vector (n + 1) a
v) = Vector (n + 1) a
v

generate' :: forall h a
           . NatRepr h
          -> (forall n. (n <= h) => NatRepr n -> a)
          -> Vector' a h
generate' :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> a) -> Vector' a h
generate' NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> a
gen =
  Identity (Vector' a h) -> Vector' a h
forall a. Identity a -> a
runIdentity (Identity (Vector' a h) -> Vector' a h)
-> Identity (Vector' a h) -> Vector' a h
forall a b. (a -> b) -> a -> b
$ NatRepr h
-> (forall (n :: Nat).
    (n <= h) =>
    NatRepr n -> () -> Identity (a, ()))
-> ()
-> Identity (Vector' a h)
forall (m :: Type -> Type) (h :: Nat) a b.
Monad m =>
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector' a h)
unfoldrWithIndexM' NatRepr h
h (\NatRepr n
n ()
_last -> (a, ()) -> Identity (a, ())
forall a. a -> Identity a
Identity (NatRepr n -> a
forall (n :: Nat). (n <= h) => NatRepr n -> a
gen NatRepr n
n, ())) ()

-- | Apply a function to each element in a range starting at zero;
-- return the a vector of values obtained.
-- cf. both @natFromZero@ and @Data.Vector.generate@
generate :: forall h a
          . NatRepr h
         -> (forall n. (n <= h) => NatRepr n -> a)
         -> Vector (h + 1) a
generate :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> a)
-> Vector (h + 1) a
generate NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> a
gen = Vector' a h -> Vector (h + 1) a
forall a (n :: Nat). Vector' a n -> Vector (n + 1) a
unVector' (NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> a) -> Vector' a h
forall (h :: Nat) a.
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> a) -> Vector' a h
generate' NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> a
gen)

-- | Since @Vector@ is traversable, we can pretty trivially sequence
-- @natFromZeroVec@ inside a monad.
generateM :: forall m h a. (Monad m)
          => NatRepr h
          -> (forall n. (n <= h) => NatRepr n -> m a)
          -> m (Vector (h + 1) a)
generateM :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> m a)
-> m (Vector (h + 1) a)
generateM NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> m a
gen = Vector (h + 1) (m a) -> m (Vector (h + 1) a)
forall (t :: Type -> Type) (m :: Type -> Type) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence (Vector (h + 1) (m a) -> m (Vector (h + 1) a))
-> Vector (h + 1) (m a) -> m (Vector (h + 1) a)
forall a b. (a -> b) -> a -> b
$ NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> m a)
-> Vector (h + 1) (m a)
forall (h :: Nat) a.
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> a)
-> Vector (h + 1) a
generate NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> m a
gen

newtype Compose3 m f g a = Compose3 { Compose3 m f g a -> m (f (g a))
getCompose3 :: m (f (g a)) }

unfoldrWithIndexM' :: forall m h a b. (Monad m)
                  => NatRepr h
                  -> (forall n. (n <= h) => NatRepr n -> b -> m (a, b))
                  -> b
                  -> m (Vector' a h)
unfoldrWithIndexM' :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector' a h)
unfoldrWithIndexM' NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b)
gen b
start =
  case NatRepr h -> Either (h :~: 0) (LeqProof 1 h)
forall (n :: Nat). NatRepr n -> Either (n :~: 0) (LeqProof 1 n)
isZeroOrGT1 NatRepr h
h of
    Left h :~: 0
Refl -> (b, Vector' a 0) -> Vector' a 0
forall a b. (a, b) -> b
snd ((b, Vector' a 0) -> Vector' a 0)
-> m (b, Vector' a 0) -> m (Vector' a 0)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Compose3 m ((,) b) (Vector' a) 0 -> m (b, Vector' a 0)
forall k (m :: k -> Type) k (f :: k -> k) k (g :: k -> k) (a :: k).
Compose3 m f g a -> m (f (g a))
getCompose3 Compose3 m ((,) b) (Vector' a) 0
base
    Right LeqProof 1 h
LeqProof ->
      case (NatRepr h -> NatRepr 1 -> ((h - 1) + 1) :~: h
forall (f :: Nat -> Type) (m :: Nat) (g :: Nat -> Type) (n :: Nat).
(n <= m) =>
f m -> g n -> ((m - n) + n) :~: m
minusPlusCancel NatRepr h
h (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) :: h - 1 + 1 :~: h) of { ((h - 1) + 1) :~: h
Refl ->
        (b, Vector' a h) -> Vector' a h
forall a b. (a, b) -> b
snd ((b, Vector' a h) -> Vector' a h)
-> m (b, Vector' a h) -> m (Vector' a h)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Compose3 m ((,) b) (Vector' a) h -> m (b, Vector' a h)
forall k (m :: k -> Type) k (f :: k -> k) k (g :: k -> k) (a :: k).
Compose3 m f g a -> m (f (g a))
getCompose3 (NatRepr (h - 1)
-> NatRepr (h - 1)
-> Compose3 m ((,) b) (Vector' a) 0
-> (forall (n :: Nat).
    (n <= (h - 1)) =>
    NatRepr n
    -> Compose3 m ((,) b) (Vector' a) n
    -> Compose3 m ((,) b) (Vector' a) (n + 1))
-> Compose3 m ((,) b) (Vector' a) ((h - 1) + 1)
forall (m :: Nat) (h :: Nat) (f :: Nat -> Type).
(m <= h) =>
NatRepr m
-> NatRepr h
-> f 0
-> (forall (n :: Nat). (n <= h) => NatRepr n -> f n -> f (n + 1))
-> f (m + 1)
natRecBounded (NatRepr h -> NatRepr (h - 1)
forall (n :: Nat). (1 <= n) => NatRepr n -> NatRepr (n - 1)
decNat NatRepr h
h) (NatRepr h -> NatRepr (h - 1)
forall (n :: Nat). (1 <= n) => NatRepr n -> NatRepr (n - 1)
decNat NatRepr h
h) Compose3 m ((,) b) (Vector' a) 0
base forall (n :: Nat).
(n <= (h - 1)) =>
NatRepr n
-> Compose3 m ((,) b) (Vector' a) n
-> Compose3 m ((,) b) (Vector' a) (n + 1)
forall (p :: Nat).
(1 <= h, p <= (h - 1)) =>
NatRepr p
-> Compose3 m ((,) b) (Vector' a) p
-> Compose3 m ((,) b) (Vector' a) (p + 1)
step)
      }
  where base :: Compose3 m ((,) b) (Vector' a) 0
        base :: Compose3 m ((,) b) (Vector' a) 0
base =
          case LeqProof 0 h
forall (n :: Nat). LeqProof 0 n
leqZero @h of { LeqProof 0 h
LeqProof ->
          m (b, Vector' a 0) -> Compose3 m ((,) b) (Vector' a) 0
forall k k k (m :: k -> Type) (f :: k -> k) (g :: k -> k) (a :: k).
m (f (g a)) -> Compose3 m f g a
Compose3 (m (b, Vector' a 0) -> Compose3 m ((,) b) (Vector' a) 0)
-> m (b, Vector' a 0) -> Compose3 m ((,) b) (Vector' a) 0
forall a b. (a -> b) -> a -> b
$ (\(a
hd, b
b) -> (b
b, Vector (0 + 1) a -> Vector' a 0
forall a (n :: Nat). Vector (n + 1) a -> Vector' a n
MkVector' (a -> Vector 1 a
forall a. a -> Vector 1 a
singleton a
hd))) ((a, b) -> (b, Vector' a 0)) -> m (a, b) -> m (b, Vector' a 0)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> NatRepr 0 -> b -> m (a, b)
forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b)
gen (KnownNat 0 => NatRepr 0
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @0) b
start
          }
        step :: forall p. (1 <= h, p <= h - 1)
             => NatRepr p
             -> Compose3 m ((,) b) (Vector' a) p
             -> Compose3 m ((,) b) (Vector' a) (p + 1)
        step :: NatRepr p
-> Compose3 m ((,) b) (Vector' a) p
-> Compose3 m ((,) b) (Vector' a) (p + 1)
step NatRepr p
p (Compose3 m (b, Vector' a p)
mv) =
          case NatRepr h -> NatRepr 1 -> ((h - 1) + 1) :~: h
forall (f :: Nat -> Type) (m :: Nat) (g :: Nat -> Type) (n :: Nat).
(n <= m) =>
f m -> g n -> ((m - n) + n) :~: m
minusPlusCancel NatRepr h
h (KnownNat 1 => NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @1) :: h - 1 + 1 :~: h of { ((h - 1) + 1) :~: h
Refl ->
          case (LeqProof p (h - 1)
-> LeqProof 1 1 -> LeqProof (p + 1) ((h - 1) + 1)
forall (x_l :: Nat) (x_h :: Nat) (y_l :: Nat) (y_h :: Nat).
LeqProof x_l x_h
-> LeqProof y_l y_h -> LeqProof (x_l + y_l) (x_h + y_h)
leqAdd2 (LeqProof p (h - 1)
forall (m :: Nat) (n :: Nat). (m <= n) => LeqProof m n
LeqProof :: LeqProof p (h-1))
                        (LeqProof 1 1
forall (m :: Nat) (n :: Nat). (m <= n) => LeqProof m n
LeqProof :: LeqProof 1 1) :: LeqProof (p+1) h) of { LeqProof (p + 1) h
LeqProof ->
            m (b, Vector' a (p + 1)) -> Compose3 m ((,) b) (Vector' a) (p + 1)
forall k k k (m :: k -> Type) (f :: k -> k) (g :: k -> k) (a :: k).
m (f (g a)) -> Compose3 m f g a
Compose3 (m (b, Vector' a (p + 1))
 -> Compose3 m ((,) b) (Vector' a) (p + 1))
-> m (b, Vector' a (p + 1))
-> Compose3 m ((,) b) (Vector' a) (p + 1)
forall a b. (a -> b) -> a -> b
$
              do (b
seed, MkVector' Vector (p + 1) a
v) <- m (b, Vector' a p)
mv
                 (a
next, b
nextSeed) <- NatRepr (p + 1) -> b -> m (a, b)
forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b)
gen (NatRepr p -> NatRepr (p + 1)
forall (n :: Nat). NatRepr n -> NatRepr (n + 1)
incNat NatRepr p
p) b
seed
                 (b, Vector' a (p + 1)) -> m (b, Vector' a (p + 1))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ((b, Vector' a (p + 1)) -> m (b, Vector' a (p + 1)))
-> (b, Vector' a (p + 1)) -> m (b, Vector' a (p + 1))
forall a b. (a -> b) -> a -> b
$ (b
nextSeed, Vector ((p + 1) + 1) a -> Vector' a (p + 1)
forall a (n :: Nat). Vector (n + 1) a -> Vector' a n
MkVector' (Vector ((p + 1) + 1) a -> Vector' a (p + 1))
-> Vector ((p + 1) + 1) a -> Vector' a (p + 1)
forall a b. (a -> b) -> a -> b
$ Vector (p + 1) a -> a -> Vector ((p + 1) + 1) a
forall (n :: Nat) a. Vector n a -> a -> Vector (n + 1) a
snoc Vector (p + 1) a
v a
next)
          }}

-- | Monadically unfold a vector, with access to the current index.
--
-- c.f. @Data.Vector.unfoldrExactNM@
unfoldrWithIndexM :: forall m h a b. (Monad m)
                 => NatRepr h
                 -> (forall n. (n <= h) => NatRepr n -> b -> m (a, b))
                 -> b
                 -> m (Vector (h + 1) a)
unfoldrWithIndexM :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector (h + 1) a)
unfoldrWithIndexM NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b)
gen b
start = Vector' a h -> Vector (h + 1) a
forall a (n :: Nat). Vector' a n -> Vector (n + 1) a
unVector' (Vector' a h -> Vector (h + 1) a)
-> m (Vector' a h) -> m (Vector (h + 1) a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector' a h)
forall (m :: Type -> Type) (h :: Nat) a b.
Monad m =>
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector' a h)
unfoldrWithIndexM' NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b)
gen b
start

-- | Unfold a vector, with access to the current index.
--
-- c.f. @Data.Vector.unfoldrExactN@
unfoldrWithIndex :: forall h a b
                . NatRepr h
                -> (forall n. (n <= h) => NatRepr n -> b -> (a, b))
                -> b
                -> Vector (h + 1) a
unfoldrWithIndex :: NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> (a, b))
-> b
-> Vector (h + 1) a
unfoldrWithIndex NatRepr h
h forall (n :: Nat). (n <= h) => NatRepr n -> b -> (a, b)
gen b
start =
  Vector' a h -> Vector (h + 1) a
forall a (n :: Nat). Vector' a n -> Vector (n + 1) a
unVector' (Vector' a h -> Vector (h + 1) a)
-> Vector' a h -> Vector (h + 1) a
forall a b. (a -> b) -> a -> b
$ Identity (Vector' a h) -> Vector' a h
forall a. Identity a -> a
runIdentity (Identity (Vector' a h) -> Vector' a h)
-> Identity (Vector' a h) -> Vector' a h
forall a b. (a -> b) -> a -> b
$ NatRepr h
-> (forall (n :: Nat).
    (n <= h) =>
    NatRepr n -> b -> Identity (a, b))
-> b
-> Identity (Vector' a h)
forall (m :: Type -> Type) (h :: Nat) a b.
Monad m =>
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector' a h)
unfoldrWithIndexM' NatRepr h
h (\NatRepr n
n b
v -> (a, b) -> Identity (a, b)
forall a. a -> Identity a
Identity (NatRepr n -> b -> (a, b)
forall (n :: Nat). (n <= h) => NatRepr n -> b -> (a, b)
gen NatRepr n
n b
v)) b
start

-- | Monadically construct a vector with exactly @h + 1@ elements by repeatedly
-- applying a generator function to a seed value.
--
-- c.f. @Data.Vector.unfoldrExactNM@
unfoldrM :: forall m h a b. (Monad m)
        => NatRepr h
        -> (b -> m (a, b))
        -> b
        -> m (Vector (h + 1) a)
unfoldrM :: NatRepr h -> (b -> m (a, b)) -> b -> m (Vector (h + 1) a)
unfoldrM NatRepr h
h b -> m (a, b)
gen b
start = NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector (h + 1) a)
forall (m :: Type -> Type) (h :: Nat) a b.
Monad m =>
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> m (a, b))
-> b
-> m (Vector (h + 1) a)
unfoldrWithIndexM NatRepr h
h (\NatRepr n
_ b
v -> b -> m (a, b)
gen b
v) b
start

-- | Construct a vector with exactly @h + 1@ elements by repeatedly applying a
-- generator function to a seed value.
--
-- c.f. @Data.Vector.unfoldrExactN@
unfoldr :: forall h a b
        . NatRepr h
       -> (b -> (a, b))
       -> b
       -> Vector (h + 1) a
unfoldr :: NatRepr h -> (b -> (a, b)) -> b -> Vector (h + 1) a
unfoldr NatRepr h
h b -> (a, b)
gen b
start = NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> (a, b))
-> b
-> Vector (h + 1) a
forall (h :: Nat) a b.
NatRepr h
-> (forall (n :: Nat). (n <= h) => NatRepr n -> b -> (a, b))
-> b
-> Vector (h + 1) a
unfoldrWithIndex NatRepr h
h (\NatRepr n
_ b
v -> b -> (a, b)
gen b
v) b
start

-- | Build a vector by repeatedly applying a monadic function to a seed value.
--
-- Compare to 'Vector.iterateNM'.
iterateNM :: Monad m => NatRepr n -> (a -> m a) -> a -> m (Vector (n + 1) a)
iterateNM :: NatRepr n -> (a -> m a) -> a -> m (Vector (n + 1) a)
iterateNM NatRepr n
h a -> m a
f a
start =
  case NatRepr n -> IsZeroNat n
forall (n :: Nat). NatRepr n -> IsZeroNat n
isZeroNat NatRepr n
h of
    IsZeroNat n
ZeroNat -> Vector 1 a -> m (Vector 1 a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a -> Vector 1 a
forall a. a -> Vector 1 a
singleton a
start)
    IsZeroNat n
NonZeroNat -> a -> Vector n a -> Vector (n + 1) a
forall (n :: Nat) a. a -> Vector n a -> Vector (n + 1) a
cons a
start (Vector n a -> Vector (n + 1) a)
-> m (Vector n a) -> m (Vector (n + 1) a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> NatRepr n -> (a -> m (a, a)) -> a -> m (Vector (n + 1) a)
forall (m :: Type -> Type) (h :: Nat) a b.
Monad m =>
NatRepr h -> (b -> m (a, b)) -> b -> m (Vector (h + 1) a)
unfoldrM (NatRepr (n + 1) -> NatRepr n
forall (n :: Nat). NatRepr (n + 1) -> NatRepr n
predNat NatRepr n
NatRepr (n + 1)
h) ((a -> (a, a)) -> m a -> m (a, a)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> (a, a)
forall b. b -> (b, b)
dup (m a -> m (a, a)) -> (a -> m a) -> a -> m (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m a
f) a
start
  where dup :: b -> (b, b)
dup b
x = (b
x, b
x)

-- | Build a vector by repeatedly applying a function to a seed value.
--
-- Compare to 'Vector.iterateN'
iterateN :: NatRepr n -> (a -> a) -> a -> Vector (n + 1) a
iterateN :: NatRepr n -> (a -> a) -> a -> Vector (n + 1) a
iterateN NatRepr n
h a -> a
f a
start = Identity (Vector (n + 1) a) -> Vector (n + 1) a
forall a. Identity a -> a
runIdentity (NatRepr n -> (a -> Identity a) -> a -> Identity (Vector (n + 1) a)
forall (m :: Type -> Type) (n :: Nat) a.
Monad m =>
NatRepr n -> (a -> m a) -> a -> m (Vector (n + 1) a)
iterateNM NatRepr n
h (a -> Identity a
forall a. a -> Identity a
Identity (a -> Identity a) -> (a -> a) -> a -> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) a
start)

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

coerceVec :: Coercible a b => Vector n a -> Vector n b
coerceVec :: Vector n a -> Vector n b
coerceVec = Vector n a -> Vector n b
coerce

-- | Monadically join a vector of values, using the given function.
-- This functionality can sometimes be reproduced by creating a newtype
-- wrapper and using @joinWith@, this implementation is provided for
-- convenience.
joinWithM ::
  forall m f n w.
  (1 <= w, Monad m) =>
  (forall l. (1 <= l) => NatRepr l -> f w -> f l -> m (f (w + l)))
  {- ^ A function for joining contained elements.  The first argument is
       the size of the accumulated third term, and the second argument
       is the element to join to the accumulated term.  The function
       can use any join strategy desired (prepending/"BigEndian",
       appending/"LittleEndian", etc.). -}
  -> NatRepr w
  -> Vector n (f w)
  -> m (f (n * w))

joinWithM :: (forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> f w -> f l -> m (f (w + l)))
-> NatRepr w -> Vector n (f w) -> m (f (n * w))
joinWithM forall (l :: Nat).
(1 <= l) =>
NatRepr l -> f w -> f l -> m (f (w + l))
jn NatRepr w
w = ((f (n * w), NatRepr (n * w)) -> f (n * w))
-> m (f (n * w), NatRepr (n * w)) -> m (f (n * w))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (f (n * w), NatRepr (n * w)) -> f (n * w)
forall a b. (a, b) -> a
fst (m (f (n * w), NatRepr (n * w)) -> m (f (n * w)))
-> (Vector n (f w) -> m (f (n * w), NatRepr (n * w)))
-> Vector n (f w)
-> m (f (n * w))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector n (f w) -> m (f (n * w), NatRepr (n * w))
forall (l :: Nat). Vector l (f w) -> m (f (l * w), NatRepr (l * w))
go
  where
  go :: forall l. Vector l (f w) -> m (f (l * w), NatRepr (l * w))
  go :: Vector l (f w) -> m (f (l * w), NatRepr (l * w))
go Vector l (f w)
exprs =
    case Vector l (f w) -> (f w, Either (l :~: 1) (Vector (l - 1) (f w)))
forall (n :: Nat) a.
Vector n a -> (a, Either (n :~: 1) (Vector (n - 1) a))
uncons Vector l (f w)
exprs of
      (f w
a, Left l :~: 1
Refl) -> (f w, NatRepr w) -> m (f w, NatRepr w)
forall (m :: Type -> Type) a. Monad m => a -> m a
return (f w
a, NatRepr w
w)
      (f w
a, Right Vector (l - 1) (f w)
rest) ->
        case Vector (l - 1) (f w) -> LeqProof 1 (l - 1)
forall (n :: Nat) a. Vector n a -> LeqProof 1 n
nonEmpty Vector (l - 1) (f w)
rest                of { LeqProof 1 (l - 1)
LeqProof ->
        case NatRepr (l - 1) -> NatRepr w -> LeqProof 1 ((l - 1) * w)
forall (p :: Nat -> Type) (q :: Nat -> Type) (x :: Nat) (y :: Nat).
(1 <= x, 1 <= y) =>
p x -> q y -> LeqProof 1 (x * y)
leqMulPos (Vector (l - 1) (f w) -> NatRepr (l - 1)
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector (l - 1) (f w)
rest) NatRepr w
w    of { LeqProof 1 ((l - 1) * w)
LeqProof ->
        case Vector l (f w) -> LeqProof 1 l
forall (n :: Nat) a. Vector n a -> LeqProof 1 n
nonEmpty Vector l (f w)
exprs               of { LeqProof 1 l
LeqProof ->
        case NatRepr w -> NatRepr l -> (w + ((l - 1) * w)) :~: (l * w)
forall (n :: Nat) (p :: Nat -> Type) (w :: Nat) (q :: Nat -> Type).
(1 <= n) =>
p w -> q n -> (w + ((n - 1) * w)) :~: (n * w)
lemmaMul NatRepr w
w (Vector l (f w) -> NatRepr l
forall (n :: Nat) a. Vector n a -> NatRepr n
length Vector l (f w)
exprs)    of { (w + ((l - 1) * w)) :~: (l * w)
Refl -> do
          -- @siddharthist: This could probably be written applicatively?
          (f ((l - 1) * w)
res, NatRepr ((l - 1) * w)
sz) <- Vector (l - 1) (f w) -> m (f ((l - 1) * w), NatRepr ((l - 1) * w))
forall (l :: Nat). Vector l (f w) -> m (f (l * w), NatRepr (l * w))
go Vector (l - 1) (f w)
rest
          f (w + ((l - 1) * w))
joined <- NatRepr ((l - 1) * w)
-> f w -> f ((l - 1) * w) -> m (f (w + ((l - 1) * w)))
forall (l :: Nat).
(1 <= l) =>
NatRepr l -> f w -> f l -> m (f (w + l))
jn NatRepr ((l - 1) * w)
sz f w
a f ((l - 1) * w)
res
          (f (w + ((l - 1) * w)), NatRepr (w + ((l - 1) * w)))
-> m (f (w + ((l - 1) * w)), NatRepr (w + ((l - 1) * w)))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (f (w + ((l - 1) * w))
joined, NatRepr w -> NatRepr ((l - 1) * w) -> NatRepr (w + ((l - 1) * w))
forall (m :: Nat) (n :: Nat).
NatRepr m -> NatRepr n -> NatRepr (m + n)
addNat NatRepr w
w NatRepr ((l - 1) * w)
sz)
        }}}}

-- | Join a vector of vectors, using the given function to combine the
-- sub-vectors.
joinWith ::
  forall f n w.
  (1 <= w) =>
  (forall l. (1 <= l) => NatRepr l -> f w -> f l -> f (w + l))
  {- ^ A function for joining contained elements.  The first argument is
       the size of the accumulated third term, and the second argument
       is the element to join to the accumulated term.  The function
       can use any join strategy desired (prepending/"BigEndian",
       appending/"LittleEndian", etc.). -}
  -> NatRepr w
  -> Vector n (f w)
  -> f (n * w)
joinWith :: (forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> f w -> f l -> f (w + l))
-> NatRepr w -> Vector n (f w) -> f (n * w)
joinWith forall (l :: Nat). (1 <= l) => NatRepr l -> f w -> f l -> f (w + l)
jn NatRepr w
w Vector n (f w)
v = Identity (f (n * w)) -> f (n * w)
forall a. Identity a -> a
runIdentity (Identity (f (n * w)) -> f (n * w))
-> Identity (f (n * w)) -> f (n * w)
forall a b. (a -> b) -> a -> b
$ (forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> f w -> f l -> Identity (f (w + l)))
-> NatRepr w -> Vector n (f w) -> Identity (f (n * w))
forall (m :: Type -> Type) (f :: Nat -> Type) (n :: Nat)
       (w :: Nat).
(1 <= w, Monad m) =>
(forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> f w -> f l -> m (f (w + l)))
-> NatRepr w -> Vector n (f w) -> m (f (n * w))
joinWithM (\NatRepr l
n f w
x -> f (w + l) -> Identity (f (w + l))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (f (w + l) -> Identity (f (w + l)))
-> (f l -> f (w + l)) -> f l -> Identity (f (w + l))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NatRepr l -> f w -> f l -> f (w + l)
forall (l :: Nat). (1 <= l) => NatRepr l -> f w -> f l -> f (w + l)
jn NatRepr l
n f w
x)) NatRepr w
w Vector n (f w)
v
{-# Inline joinWith #-}

-- | Split a vector into a vector of vectors.
--
-- The "Endian" parameter determines the ordering of the inner
-- vectors.  If "LittleEndian", then less significant bits go into
-- smaller indexes.  If "BigEndian", then less significant bits go
-- into larger indexes.  See the documentation for 'split' for more
-- details.
splitWith :: forall f w n.
  (1 <= w, 1 <= n) =>
  Endian ->
  (forall i. (i + w <= n * w) =>
             NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w)
  {- ^ A function for slicing out a chunk of length @w@, starting at @i@ -} ->
  NatRepr n -> NatRepr w -> f (n * w) -> Vector n (f w)
splitWith :: Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w)
-> NatRepr n
-> NatRepr w
-> f (n * w)
-> Vector n (f w)
splitWith Endian
endian forall (i :: Nat).
((i + w) <= (n * w)) =>
NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w
select NatRepr n
n NatRepr w
w f (n * w)
val = Vector (f w) -> Vector n (f w)
forall (n :: Nat) a. (1 <= n) => Vector a -> Vector n a
Vector ((forall s. ST s (MVector s (f w))) -> Vector (f w)
forall a. (forall s. ST s (MVector s a)) -> Vector a
Vector.create forall s. ST s (MVector s (f w))
initializer)
  where
  len :: Int
len          = NatRepr n -> Int
forall (n :: Nat). NatRepr n -> Int
widthVal NatRepr n
n
  start :: Int
  next :: Int -> Int
  (Int
start,Int -> Int
next) = case Endian
endian of
                   Endian
LittleEndian -> (Int
0, Int -> Int
forall a. Enum a => a -> a
succ)
                   Endian
BigEndian    -> (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int -> Int
forall a. Enum a => a -> a
pred)

  initializer :: forall s. ST s (MVector s (f w))
  initializer :: ST s (MVector s (f w))
initializer =
    do LeqProof 1 (n * w)
LeqProof <- LeqProof 1 (n * w) -> ST s (LeqProof 1 (n * w))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (NatRepr n -> NatRepr w -> LeqProof 1 (n * w)
forall (p :: Nat -> Type) (q :: Nat -> Type) (x :: Nat) (y :: Nat).
(1 <= x, 1 <= y) =>
p x -> q y -> LeqProof 1 (x * y)
leqMulPos NatRepr n
n NatRepr w
w)
       LeqProof w (n * w)
LeqProof <- LeqProof w (n * w) -> ST s (LeqProof w (n * w))
forall (m :: Type -> Type) a. Monad m => a -> m a
return (NatRepr n -> NatRepr w -> LeqProof w (n * w)
forall (x :: Nat) (p :: Nat -> Type) (q :: Nat -> Type) (y :: Nat).
(1 <= x) =>
p x -> q y -> LeqProof y (x * y)
leqMulMono NatRepr n
n NatRepr w
w)

       MVector s (f w)
v <- Int -> ST s (MVector (PrimState (ST s)) (f w))
forall (m :: Type -> Type) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MVector.new Int
len
       let fill :: Int -> NatRepr i -> ST s ()
           fill :: Int -> NatRepr i -> ST s ()
fill Int
loc NatRepr i
i =
             let end :: NatRepr (i + w)
end = NatRepr i -> NatRepr w -> NatRepr (i + w)
forall (m :: Nat) (n :: Nat).
NatRepr m -> NatRepr n -> NatRepr (m + n)
addNat NatRepr i
i NatRepr w
w in
             case NatRepr (i + w)
-> NatRepr (n * w) -> Maybe (LeqProof (i + w) (n * w))
forall (m :: Nat) (n :: Nat).
NatRepr m -> NatRepr n -> Maybe (LeqProof m n)
testLeq NatRepr (i + w)
end NatRepr (n * w)
inLen of
               Just LeqProof (i + w) (n * w)
LeqProof ->
                 do MVector (PrimState (ST s)) (f w) -> Int -> f w -> ST s ()
forall (m :: Type -> Type) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MVector.write MVector s (f w)
MVector (PrimState (ST s)) (f w)
v Int
loc (NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w
forall (i :: Nat).
((i + w) <= (n * w)) =>
NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w
select NatRepr (n * w)
inLen NatRepr i
i f (n * w)
val)
                    Int -> NatRepr (i + w) -> ST s ()
forall (i :: Nat). Int -> NatRepr i -> ST s ()
fill (Int -> Int
next Int
loc) NatRepr (i + w)
end
               Maybe (LeqProof (i + w) (n * w))
Nothing -> () -> ST s ()
forall (m :: Type -> Type) a. Monad m => a -> m a
return ()


       Int -> NatRepr 0 -> ST s ()
forall (i :: Nat). Int -> NatRepr i -> ST s ()
fill Int
start (KnownNat 0 => NatRepr 0
forall (n :: Nat). KnownNat n => NatRepr n
knownNat @0)
       MVector s (f w) -> ST s (MVector s (f w))
forall (m :: Type -> Type) a. Monad m => a -> m a
return MVector s (f w)
v

  inLen :: NatRepr (n * w)
  inLen :: NatRepr (n * w)
inLen = NatRepr n -> NatRepr w -> NatRepr (n * w)
forall (n :: Nat) (m :: Nat).
NatRepr n -> NatRepr m -> NatRepr (n * m)
natMultiply NatRepr n
n NatRepr w
w
{-# Inline splitWith #-}

-- We can sneakily put our functor in the parameter "f" of @splitWith@ using the
-- @Compose@ newtype.
-- | An applicative version of @splitWith@.
splitWithA :: forall f g w n. (Applicative f, 1 <= w, 1 <= n) =>
  Endian ->
  (forall i. (i + w <= n * w) =>
             NatRepr (n * w) -> NatRepr i -> g (n * w) -> f (g w))
  {- ^ f function for slicing out f chunk of length @w@, starting at @i@ -} ->
  NatRepr n -> NatRepr w -> g (n * w) -> f (Vector n (g w))
splitWithA :: Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w) -> NatRepr i -> g (n * w) -> f (g w))
-> NatRepr n
-> NatRepr w
-> g (n * w)
-> f (Vector n (g w))
splitWithA Endian
e forall (i :: Nat).
((i + w) <= (n * w)) =>
NatRepr (n * w) -> NatRepr i -> g (n * w) -> f (g w)
select NatRepr n
n NatRepr w
w g (n * w)
val = (Compose f g w -> f (g w))
-> Vector n (Compose f g w) -> f (Vector n (g w))
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Compose f g w -> f (g w)
forall k1 (f :: k1 -> Type) k2 (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (Vector n (Compose f g w) -> f (Vector n (g w)))
-> Vector n (Compose f g w) -> f (Vector n (g w))
forall a b. (a -> b) -> a -> b
$
  Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w)
    -> NatRepr i -> Compose f g (n * w) -> Compose f g w)
-> NatRepr n
-> NatRepr w
-> Compose f g (n * w)
-> Vector n (Compose f g w)
forall (f :: Nat -> Type) (w :: Nat) (n :: Nat).
(1 <= w, 1 <= n) =>
Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w)
-> NatRepr n
-> NatRepr w
-> f (n * w)
-> Vector n (f w)
splitWith @(Compose f g) Endian
e forall (i :: Nat).
((i + w) <= (n * w)) =>
NatRepr (n * w)
-> NatRepr i -> Compose f g (n * w) -> Compose f g w
select' NatRepr n
n NatRepr w
w (Compose f g (n * w) -> Vector n (Compose f g w))
-> Compose f g (n * w) -> Vector n (Compose f g w)
forall a b. (a -> b) -> a -> b
$ f (g (n * w)) -> Compose f g (n * w)
forall k k1 (f :: k -> Type) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (g (n * w) -> f (g (n * w))
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure g (n * w)
val)
  where -- Wrap everything in Compose
        select' :: (forall i. (i + w <= n * w)
                => NatRepr (n * w) -> NatRepr i -> Compose f g (n * w) -> Compose f g w)
        -- Whatever we pass in as "val" is what's passed to select anyway,
        -- so there's no need to examine the argument. Just use "val" directly here.
        select' :: NatRepr (n * w)
-> NatRepr i -> Compose f g (n * w) -> Compose f g w
select' NatRepr (n * w)
nw NatRepr i
i Compose f g (n * w)
_ = f (g w) -> Compose f g w
forall k k1 (f :: k -> Type) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (f (g w) -> Compose f g w) -> f (g w) -> Compose f g w
forall a b. (a -> b) -> a -> b
$ NatRepr (n * w) -> NatRepr i -> g (n * w) -> f (g w)
forall (i :: Nat).
((i + w) <= (n * w)) =>
NatRepr (n * w) -> NatRepr i -> g (n * w) -> f (g w)
select NatRepr (n * w)
nw NatRepr i
i g (n * w)
val

newtype Vec a n = Vec (Vector n a)

vSlice :: (i + w <= l, 1 <= w) =>
  NatRepr w -> NatRepr l -> NatRepr i -> Vec a l -> Vec a w
vSlice :: NatRepr w -> NatRepr l -> NatRepr i -> Vec a l -> Vec a w
vSlice NatRepr w
w NatRepr l
_ NatRepr i
i (Vec Vector l a
xs) = Vector w a -> Vec a w
forall a (n :: Nat). Vector n a -> Vec a n
Vec (NatRepr i -> NatRepr w -> Vector l a -> Vector w a
forall (i :: Nat) (w :: Nat) (n :: Nat) a.
((i + w) <= n, 1 <= w) =>
NatRepr i -> NatRepr w -> Vector n a -> Vector w a
slice NatRepr i
i NatRepr w
w Vector l a
xs)
{-# Inline vSlice #-}

-- | Append the two bit vectors.  The first argument is
-- at the lower indexes of the resulting vector.
vAppend :: NatRepr n -> Vec a m -> Vec a n -> Vec a (m + n)
vAppend :: NatRepr n -> Vec a m -> Vec a n -> Vec a (m + n)
vAppend NatRepr n
_ (Vec Vector m a
xs) (Vec Vector n a
ys) = Vector (m + n) a -> Vec a (m + n)
forall a (n :: Nat). Vector n a -> Vec a n
Vec (Vector m a -> Vector n a -> Vector (m + n) a
forall (m :: Nat) a (n :: Nat).
Vector m a -> Vector n a -> Vector (m + n) a
append Vector m a
xs Vector n a
ys)
{-# Inline vAppend #-}

-- | Split a vector into a vector of vectors.  The default ordering of
-- the outer result vector is "LittleEndian".
--
-- For example:
-- @
--   let wordsize = knownNat :: NatRepr 3
--       vecsize = knownNat :: NatRepr 12
--       numwords = knownNat :: NatRepr 4  (12 / 3)
--       Just inpvec = fromList vecsize [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 ]
--   in show (split numwords wordsize inpvec) == "[ [1,2,3], [4,5,6], [7,8,9], [10,11,12] ]"
-- @
-- whereas a BigEndian result would have been
-- @
--      [ [10,11,12], [7,8,9], [4,5,6], [1,2,3] ]
-- @
split :: (1 <= w, 1 <= n) =>
         NatRepr n -- ^ Inner vector size
      -> NatRepr w -- ^ Outer vector size
      -> Vector (n * w) a -- ^ Input vector
      -> Vector n (Vector w a)
split :: NatRepr n -> NatRepr w -> Vector (n * w) a -> Vector n (Vector w a)
split NatRepr n
n NatRepr w
w Vector (n * w) a
xs = Vector n (Vec a w) -> Vector n (Vector w a)
forall a b (n :: Nat). Coercible a b => Vector n a -> Vector n b
coerceVec (Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w) -> NatRepr i -> Vec a (n * w) -> Vec a w)
-> NatRepr n
-> NatRepr w
-> Vec a (n * w)
-> Vector n (Vec a w)
forall (f :: Nat -> Type) (w :: Nat) (n :: Nat).
(1 <= w, 1 <= n) =>
Endian
-> (forall (i :: Nat).
    ((i + w) <= (n * w)) =>
    NatRepr (n * w) -> NatRepr i -> f (n * w) -> f w)
-> NatRepr n
-> NatRepr w
-> f (n * w)
-> Vector n (f w)
splitWith Endian
LittleEndian (NatRepr w
-> NatRepr (n * w) -> NatRepr i -> Vec a (n * w) -> Vec a w
forall (i :: Nat) (w :: Nat) (l :: Nat) a.
((i + w) <= l, 1 <= w) =>
NatRepr w -> NatRepr l -> NatRepr i -> Vec a l -> Vec a w
vSlice NatRepr w
w) NatRepr n
n NatRepr w
w (Vector (n * w) a -> Vec a (n * w)
forall a (n :: Nat). Vector n a -> Vec a n
Vec Vector (n * w) a
xs))
{-# Inline split #-}

-- | Join a vector of vectors into a single vector.  Assumes an
-- append/"LittleEndian" join strategy: the order of the inner vectors
-- is preserved in the result vector.
--
-- @
--   let innersize = knownNat :: NatRepr 4
--       Just inner1 = fromList innersize [ 1, 2, 3, 4 ]
--       Just inner2 = fromList innersize [ 5, 6, 7, 8 ]
--       Just inner3 = fromList innersize [ 9, 10, 11, 12 ]
--       outersize = knownNat :: NatRepr 3
--       Just outer = fromList outersize [ inner1, inner2, inner3 ]
--   in show (join innersize outer) = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 ]
-- @
-- a prepend/"BigEndian" join strategy would have the result:
-- @
--   [ 9, 10, 11, 12, 5, 6, 7, 8, 1, 2, 3, 4 ]
-- @
join :: (1 <= w) => NatRepr w -> Vector n (Vector w a) -> Vector (n * w) a
join :: NatRepr w -> Vector n (Vector w a) -> Vector (n * w) a
join NatRepr w
w Vector n (Vector w a)
xs = Vector (n * w) a
ys
  where Vec Vector (n * w) a
ys = (forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> Vec a w -> Vec a l -> Vec a (w + l))
-> NatRepr w -> Vector n (Vec a w) -> Vec a (n * w)
forall (f :: Nat -> Type) (n :: Nat) (w :: Nat).
(1 <= w) =>
(forall (l :: Nat).
 (1 <= l) =>
 NatRepr l -> f w -> f l -> f (w + l))
-> NatRepr w -> Vector n (f w) -> f (n * w)
joinWith forall (l :: Nat).
(1 <= l) =>
NatRepr l -> Vec a w -> Vec a l -> Vec a (w + l)
forall (n :: Nat) a (m :: Nat).
NatRepr n -> Vec a m -> Vec a n -> Vec a (m + n)
vAppend NatRepr w
w (Vector n (Vector w a) -> Vector n (Vec a w)
forall a b (n :: Nat). Coercible a b => Vector n a -> Vector n b
coerceVec Vector n (Vector w a)
xs)
{-# Inline join #-}