{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Arrays with type-level shape information, backed by unboxed vectors.
--
-- This module is a facade over 'Harpie.Fixed.Generic' with the backing
-- vector family fixed to 'Data.Vector.Unboxed.Vector'.
module Harpie.Fixed.Unboxed
  ( -- * Re-exports from the generic core
    module Harpie.Fixed.Generic,

    -- * Unboxed facade
    Array,
    FromVector (..),
    array,
    safeArray,
    unsafeArray,
    validate,
    unsafeModifyVector,
    vector,
    vector',
  )
where

import Data.Bool (bool)
import Data.Maybe (fromMaybe)
import Data.Vector.Unboxed (Unbox)
import Data.Vector.Unboxed qualified as VU
import GHC.TypeNats (KnownNat, SNat, withKnownNat)
import Harpie.Fixed.Generic hiding
  ( Array,
    FromVector (..),
    array,
    safeArray,
    unsafeArray,
    unsafeModifyVector,
    validate,
    vector,
    vector',
  )
import Harpie.Fixed.Generic qualified as G
import Harpie.Shape (KnownNats)
import Prelude hiding (cycle, drop, length, repeat, take, zipWith)

-- | A fixed-shape array backed by an unboxed vector.
type Array s a = G.Array VU.Vector s a

-- | Conversion to and from an unboxed 'VU.Vector'.
class (Unbox a) => FromVector t a | t -> a where
  asVector :: t -> VU.Vector a
  vectorAs :: VU.Vector a -> t

instance (Unbox a) => FromVector (VU.Vector a) a where
  asVector :: Vector a -> Vector a
asVector = Vector a -> Vector a
forall a. a -> a
id
  vectorAs :: Vector a -> Vector a
vectorAs = Vector a -> Vector a
forall a. a -> a
id

instance (Unbox a) => FromVector [a] a where
  asVector :: [a] -> Vector a
asVector = [a] -> Vector a
forall a. Unbox a => [a] -> Vector a
VU.fromList
  vectorAs :: Vector a -> [a]
vectorAs = Vector a -> [a]
forall a. Unbox a => Vector a -> [a]
VU.toList

instance (Unbox a, KnownNats s) => FromVector (Array s a) a where
  asVector :: Array s a -> Vector a
asVector (G.Array Vector a
v) = Vector a
v
  vectorAs :: Vector a -> Array s a
vectorAs Vector a
v = Vector a -> Array s a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
G.Array Vector a
v

-- | Construct an array without shape validation.
unsafeArray :: (KnownNats s, FromVector t a) => t -> Array s a
unsafeArray :: forall (s :: [Nat]) t a.
(KnownNats s, FromVector t a) =>
t -> Array s a
unsafeArray (t -> Vector a
forall t a. FromVector t a => t -> Vector a
asVector -> Vector a
v) = Vector a -> Array Vector s a
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
G.Array Vector a
v

-- | Validate the size and shape of an array.
validate :: (KnownNats s, Unbox a) => Array s a -> Bool
validate :: forall (s :: [Nat]) a. (KnownNats s, Unbox a) => Array s a -> Bool
validate Array s a
a = Array s a -> Int
forall (v :: * -> *) a (s :: [Nat]).
(KnownNats s, Vector v a) =>
Array v s a -> Int
G.size Array s a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length (Array s a -> Vector a
forall t a. FromVector t a => t -> Vector a
asVector Array s a
a)

-- | Construct an Array, checking shape.
safeArray :: (KnownNats s, FromVector t a) => t -> Maybe (Array s a)
safeArray :: forall (s :: [Nat]) t a.
(KnownNats s, FromVector t a) =>
t -> Maybe (Array s a)
safeArray t
v =
  Maybe (Array s a) -> Maybe (Array s a) -> Bool -> Maybe (Array s a)
forall a. a -> a -> Bool -> a
bool Maybe (Array s a)
forall a. Maybe a
Nothing (Array s a -> Maybe (Array s a)
forall a. a -> Maybe a
Just Array s a
a) (Array s a -> Bool
forall (s :: [Nat]) a. (KnownNats s, Unbox a) => Array s a -> Bool
validate Array s a
a)
  where
    a :: Array s a
a = t -> Array s a
forall (s :: [Nat]) t a.
(KnownNats s, FromVector t a) =>
t -> Array s a
unsafeArray t
v

-- | Construct an Array, throwing an exception on a bad shape.
array :: forall s a t. (KnownNats s, FromVector t a) => t -> Array s a
array :: forall (s :: [Nat]) a t.
(KnownNats s, FromVector t a) =>
t -> Array s a
array t
v =
  Array s a -> Maybe (Array s a) -> Array s a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Array s a
forall a. HasCallStack => [Char] -> a
error [Char]
"Shape Mismatch") (t -> Maybe (Array s a)
forall (s :: [Nat]) t a.
(KnownNats s, FromVector t a) =>
t -> Maybe (Array s a)
safeArray t
v)

-- | Unsafely modify an array vector.
unsafeModifyVector ::
  (KnownNats s, KnownNats s', Unbox a, Unbox b) =>
  (VU.Vector a -> VU.Vector b) ->
  Array s a ->
  Array s' b
unsafeModifyVector :: forall (s :: [Nat]) (s' :: [Nat]) a b.
(KnownNats s, KnownNats s', Unbox a, Unbox b) =>
(Vector a -> Vector b) -> Array s a -> Array s' b
unsafeModifyVector Vector a -> Vector b
f (G.Array Vector a
v) = Vector b -> Array Vector s' b
forall {k} (v :: k -> *) (s :: [Nat]) (a :: k). v a -> Array v s a
G.Array (Vector a -> Vector b
f Vector a
v)

-- | Create a one-dimensional array.
vector :: forall n a t. (FromVector t a, KnownNat n) => t -> Array '[n] a
vector :: forall (n :: Nat) a t.
(FromVector t a, KnownNat n) =>
t -> Array '[n] a
vector t
xs = t -> Array '[n] a
forall (s :: [Nat]) a t.
(KnownNats s, FromVector t a) =>
t -> Array s a
array t
xs

-- | Create a one-dimensional array with an explicit 'SNat'.
vector' :: forall a n t. (FromVector t a) => SNat n -> t -> Array '[n] a
vector' :: forall a (n :: Nat) t.
FromVector t a =>
SNat n -> t -> Array '[n] a
vector' SNat n
n t
xs = SNat n -> (KnownNat n => Array '[n] a) -> Array '[n] a
forall (n :: Nat) r. SNat n -> (KnownNat n => r) -> r
withKnownNat SNat n
n (t -> Array '[n] a
forall (n :: Nat) a t.
(FromVector t a, KnownNat n) =>
t -> Array '[n] a
vector t
xs)