{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}

-- | Arrays with shape information and computations at a value-level.
--
-- This module is a unboxed-vector facade over 'Harpie.Array.Generic'.
module Harpie.Array.Unboxed
  ( -- * Re-exports from the generic core
    module Harpie.Array.Generic,

    -- * Unboxed facade
    Array,
    array,
    (><),
    validate,
    safeArray,
    unsafeModifyVector,

    -- * Conversion
    FromVector (..),
    FromArray (..),
  )
where

import Data.Bool
import Data.Vector.Unboxed qualified as VU
import Harpie.Array.Generic hiding
  ( Array,
    array,
    arrayV,
    asVector,
    safeArray,
    unsafeModifyVector,
    validate,
    (><),
  )
import Harpie.Array.Generic qualified as G
import Harpie.Shape qualified as S
import Prelude hiding (cycle, drop, length, repeat, take, zip, zipWith)

type Array a = G.Array VU.Vector a

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

instance (VU.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 (VU.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 (VU.Unbox a) => FromVector (Array a) a where
  asVector :: Array a -> Vector a
asVector (G.UnsafeArray Vector Int
_ Vector Int
_ Vector a
v) = Vector a
v
  vectorAs :: Vector a -> Array a
vectorAs Vector a
v = [Int] -> Vector a -> Array a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
G.unsafeArrayL [Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v] Vector a
v

-- | Conversion to and from an 'Array'.
class (VU.Unbox a) => FromArray t a | t -> a where
  asArray :: t -> Array a
  arrayAs :: Array a -> t

instance (VU.Unbox a) => FromArray (Array a) a where
  asArray :: Array a -> Array a
asArray = Array a -> Array a
forall a. a -> a
id
  arrayAs :: Array a -> Array a
arrayAs = Array a -> Array a
forall a. a -> a
id

instance (VU.Unbox a) => FromArray [a] a where
  asArray :: [a] -> Array a
asArray [a]
l = [Int] -> Vector a -> Array a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
G.unsafeArrayL [[a] -> Int
forall a. [a] -> Int
S.rankL [a]
l] ([a] -> Vector a
forall a. Unbox a => [a] -> Vector a
VU.fromList [a]
l)
  arrayAs :: Array a -> [a]
arrayAs (G.UnsafeArray Vector Int
_ Vector Int
_ Vector a
v) = Vector a -> [a]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector a
v

instance (VU.Unbox a) => FromArray (VU.Vector a) a where
  asArray :: Vector a -> Array a
asArray Vector a
v = [Int] -> Vector a -> Array a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
G.unsafeArrayL [Vector a -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector a
v] Vector a
v
  arrayAs :: Array a -> Vector a
arrayAs (G.UnsafeArray Vector Int
_ Vector Int
_ Vector a
v) = Vector a
v

-- | Construct an array from a shape and a value without any shape validation.
array :: (FromVector t a) => [Int] -> t -> Array a
array :: forall t a. FromVector t a => [Int] -> t -> Array a
array [Int]
s (t -> Vector a
forall t a. FromVector t a => t -> Vector a
asVector -> Vector a
v) = [Int] -> Vector a -> Array Vector a
forall {k} (v :: k -> *) (a :: k). [Int] -> v a -> Array v a
G.unsafeArrayL [Int]
s Vector a
v

infixl 4 ><

-- | Construct an Array.
(><) :: (FromVector t a) => [Int] -> t -> Array a
>< :: forall t a. FromVector t a => [Int] -> t -> Array a
(><) = [Int] -> t -> Array a
forall t a. FromVector t a => [Int] -> t -> Array a
array

-- | Validate the size and shape of an array.
validate :: (VU.Unbox a) => Array a -> Bool
validate :: forall a. Unbox a => Array a -> Bool
validate Array a
a = Array a -> Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Int
G.size Array 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 a -> Vector a
forall t a. FromVector t a => t -> Vector a
asVector Array a
a)

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

-- | Unsafely modify an array vector.
unsafeModifyVector :: (FromVector u a, FromVector v b) => (u -> v) -> Array a -> Array b
unsafeModifyVector :: forall u a v b.
(FromVector u a, FromVector v b) =>
(u -> v) -> Array a -> Array b
unsafeModifyVector u -> v
f (G.UnsafeArray Vector Int
s Vector Int
_ Vector a
v) = Vector Int -> Vector b -> Array Vector b
forall {k} (v :: k -> *) (a :: k). Vector Int -> v a -> Array v a
G.unsafeArray Vector Int
s (v -> Vector b
forall t a. FromVector t a => t -> Vector a
asVector (u -> v
f (Vector a -> u
forall t a. FromVector t a => Vector a -> t
vectorAs Vector a
v)))