{- | Fortran array representation primitives.

Fortran arrays are homogeneous: every element has the same type. That means a
@REAL(8)@ array must store only @REAL(8)@s, no @REAL(4)@s. We use some type
algebra to obtain correct-by-construction types.

Also, Fortran arrays are multi-dimensional. Rather than storing a single natural
number for a single dimension, or doing some recursive work for arrays of arrays
(which don't exist in Fortran), we instead store a list of naturals, defining
each dimension's extent.
-}

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE UndecidableInstances #-}

-- unwrapping somes gets you kind and length at the same time - I figure because
-- kind is just a convenience.

module Language.Fortran.Repr.Value.Array.Machine where

import Language.Fortran.Repr.Type.Array
import Language.Fortran.Repr.Type.Scalar
import Language.Fortran.Repr.Type.Scalar.Int
import Language.Fortran.Repr.Value.Scalar.Int.Machine
import Language.Fortran.Repr.Type.Scalar.Real
import Language.Fortran.Repr.Value.Scalar.Real
import Language.Fortran.Repr.Value.Scalar.Complex
import Language.Fortran.Repr.Value.Scalar.String
import Language.Fortran.Repr.Util ( natVal'' )

import GHC.TypeNats
import Language.Fortran.Repr.Compat.Natural

import qualified Data.Vector.Sized as V
import Data.Vector.Sized ( Vector )
import Data.Kind
import Data.Singletons

type Size :: [NaturalK] -> NaturalK
type family Size dims where
    Size (dim ': dims) = dim + Size dims
    Size '[]           = 0

-- can conveniently define kinded array types like so
data FVA (ft :: k -> Type) (fk :: k) (dims :: [NaturalK])
  = FVA { forall k (ft :: k -> *) (fk :: k) (dims :: [NaturalK]).
FVA ft fk dims -> Vector (Size dims) (ft fk)
unFVA :: Vector (Size dims) (ft fk) }
deriving stock instance Show (ft fk) => Show (FVA ft fk dims)

-- makes rank 1 array
mkFVA1 :: forall l ft fk. Vector l (ft fk) -> FVA ft fk '[l]
mkFVA1 :: forall {k} (l :: NaturalK) (ft :: k -> *) (fk :: k).
Vector l (ft fk) -> FVA ft fk '[l]
mkFVA1 = forall k (ft :: k -> *) (fk :: k) (dims :: [NaturalK]).
Vector (Size dims) (ft fk) -> FVA ft fk dims
FVA

-- reifies type info
fvaShape :: forall dims ft fk. KnownNats dims => FVA ft fk dims -> Shape
fvaShape :: forall {k} (dims :: [NaturalK]) (ft :: k -> *) (fk :: k).
KnownNats dims =>
FVA ft fk dims -> Shape
fvaShape FVA ft fk dims
_ = [NaturalK] -> Shape
Shape forall a b. (a -> b) -> a -> b
$ forall (ns :: [NaturalK]). KnownNats ns => [NaturalK]
natVals @dims

-- TODO
mkSomeFVA :: (forall l. KnownNat l => Vector l a -> r) -> [a] -> r
mkSomeFVA :: forall a r.
(forall (l :: NaturalK). KnownNat l => Vector l a -> r) -> [a] -> r
mkSomeFVA forall (l :: NaturalK). KnownNat l => Vector l a -> r
f [a]
as = forall a r.
[a] -> (forall (n :: NaturalK). KnownNat n => Vector n a -> r) -> r
V.withSizedList [a]
as forall (l :: NaturalK). KnownNat l => Vector l a -> r
f

-- | Reify a list of type-level 'Natural's.
class KnownNats (ns :: [NaturalK]) where natVals :: [Natural]
instance (KnownNat n, KnownNats ns) => KnownNats (n ': ns) where
    natVals :: [NaturalK]
natVals = forall (a :: NaturalK). KnownNat a => NaturalK
natVal'' @n forall a. a -> [a] -> [a]
: forall (ns :: [NaturalK]). KnownNats ns => [NaturalK]
natVals @ns
instance KnownNats '[] where natVals :: [NaturalK]
natVals = []

-- | Wrapper for defining an array of a kind-tagged Fortran type.
data SomeFVA k ft =
    forall (fk :: k) (dims :: [NaturalK]). (KnownNats dims, SingKind k, SingI fk)
        => SomeFVA { ()
unSomeFVA :: FVA ft fk dims }
deriving stock instance Show (SomeFVA FTInt    FInt)
deriving stock instance Show (SomeFVA FTReal   FReal)
deriving stock instance Show (SomeFVA FTReal   FComplex)
deriving stock instance Show (SomeFVA NaturalK FString)

someFVAKind :: SomeFVA k ft -> Demote k
someFVAKind :: forall k (ft :: k -> *). SomeFVA k ft -> Demote k
someFVAKind (SomeFVA (FVA ft fk dims
_ :: FVA ft fk dims)) = forall {k} (a :: k). (SingKind k, SingI a) => Demote k
demote @fk

someFVAShape :: SomeFVA k ft -> Shape
someFVAShape :: forall k (ft :: k -> *). SomeFVA k ft -> Shape
someFVAShape (SomeFVA FVA ft fk dims
a) = forall {k} (dims :: [NaturalK]) (ft :: k -> *) (fk :: k).
KnownNats dims =>
FVA ft fk dims -> Shape
fvaShape FVA ft fk dims
a

-- makes rank 1 array
mkSomeFVA1
    :: forall k ft (fk :: k). (SingKind k, SingI fk)
    => [ft fk] -> SomeFVA k ft
mkSomeFVA1 :: forall k (ft :: k -> *) (fk :: k).
(SingKind k, SingI fk) =>
[ft fk] -> SomeFVA k ft
mkSomeFVA1 = forall a r.
(forall (l :: NaturalK). KnownNat l => Vector l a -> r) -> [a] -> r
mkSomeFVA forall a b. (a -> b) -> a -> b
$ forall k (ft :: k -> *) (fk :: k) (dims :: [NaturalK]).
(KnownNats dims, SingKind k, SingI fk) =>
FVA ft fk dims -> SomeFVA k ft
SomeFVA forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (l :: NaturalK) (ft :: k -> *) (fk :: k).
Vector l (ft fk) -> FVA ft fk '[l]
mkFVA1

data FArrayValue
  = FAVInt     (SomeFVA FTInt    FInt)
  | FAVReal    (SomeFVA FTReal   FReal)
  | FAVComplex (SomeFVA FTReal   FComplex)
  | FAVLogical (SomeFVA FTInt    FInt)
  | FAVString  (SomeFVA NaturalK FString)
deriving stock instance Show FArrayValue

fArrayValueType :: FArrayValue -> FArrayType
fArrayValueType :: FArrayValue -> FArrayType
fArrayValueType = \case
  FAVInt     SomeFVA FTInt FInt
a -> forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go FTInt -> FScalarType
FSTInt     SomeFVA FTInt FInt
a
  FAVReal    SomeFVA FTReal FReal
a -> forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go FTReal -> FScalarType
FSTReal    SomeFVA FTReal FReal
a
  FAVComplex SomeFVA FTReal FComplex
a -> forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go FTReal -> FScalarType
FSTComplex SomeFVA FTReal FComplex
a
  FAVLogical SomeFVA FTInt FInt
a -> forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go FTInt -> FScalarType
FSTLogical SomeFVA FTInt FInt
a
  FAVString  SomeFVA NaturalK FString
a -> forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go NaturalK -> FScalarType
FSTString  SomeFVA NaturalK FString
a
  where
    go :: (Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
    go :: forall k (ft :: k -> *).
(Demote k -> FScalarType) -> SomeFVA k ft -> FArrayType
go Demote k -> FScalarType
f SomeFVA k ft
a = FScalarType -> Shape -> FArrayType
FArrayType (Demote k -> FScalarType
f (forall k (ft :: k -> *). SomeFVA k ft -> Demote k
someFVAKind SomeFVA k ft
a)) (forall k (ft :: k -> *). SomeFVA k ft -> Shape
someFVAShape SomeFVA k ft
a)