1 files changed, 323 insertions, 0 deletions
diff --git a/src/Data/Array/Nested/Ranked.hs b/src/Data/Array/Nested/Ranked.hs
new file mode 100644
index 0000000..9778c54
--- /dev/null
+++ b/src/Data/Array/Nested/Ranked.hs
@@ -0,0 +1,323 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ViewPatterns #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
+module Data.Array.Nested.Ranked (
+  Ranked(Ranked),
+  rquotArray, rremArray, ratan2Array,
+  rshape, rrank,
+  module Data.Array.Nested.Ranked,
+  liftRanked1, liftRanked2,
+) where
+
+import Prelude hiding (mappend, mconcat)
+
+import Data.Array.RankedS qualified as S
+import Data.Bifunctor (first)
+import Data.Coerce (coerce)
+import Data.List.NonEmpty (NonEmpty)
+import Data.Proxy
+import Data.Type.Equality
+import Data.Vector.Storable qualified as VS
+import Foreign.Storable (Storable)
+import GHC.TypeLits
+import GHC.TypeNats qualified as TN
+
+import Data.Array.Nested.Convert
+import Data.Array.Nested.Lemmas
+import Data.Array.Nested.Mixed
+import Data.Array.Nested.Mixed.Shape
+import Data.Array.Nested.Permutation
+import Data.Array.Nested.Ranked.Base
+import Data.Array.Nested.Ranked.Shape
+import Data.Array.Nested.Types
+import Data.Array.Strided.Arith
+import Data.Array.XArray (XArray(..))
+import Data.Array.XArray qualified as X
+
+
+remptyArray :: KnownElt a => Ranked 1 a
+remptyArray = mtoRanked (memptyArray ZSX)
+
+-- | The total number of elements in the array.
+rsize :: Elt a => Ranked n a -> Int
+rsize = shrSize . rshape
+
+rindex :: Elt a => Ranked n a -> IIxR n -> a
+rindex (Ranked arr) idx = mindex arr (ixxFromIxR idx)
+
+rindexPartial :: forall n m a. Elt a => Ranked (n + m) a -> IIxR n -> Ranked m a
+rindexPartial (Ranked arr) idx =
+  Ranked (mindexPartial @a @(Replicate n Nothing) @(Replicate m Nothing)
+            (castWith (subst2 (lemReplicatePlusApp (ixrRank idx) (Proxy @m) (Proxy @Nothing))) arr)
+            (ixxFromIxR idx))
+
+-- | __WARNING__: All values returned from the function must have equal shape.
+-- See the documentation of 'mgenerate' for more details.
+rgenerate :: forall n a. KnownElt a => IShR n -> (IIxR n -> a) -> Ranked n a
+rgenerate sh f
+  | sn@SNat <- shrRank sh
+  , Dict <- lemKnownReplicate sn
+  , Refl <- lemRankReplicate sn
+  = Ranked (mgenerate (shxFromShR sh) (f . ixrFromIxX))
+
+-- | See the documentation of 'mlift'.
+rlift :: forall n1 n2 a. Elt a
+      => SNat n2
+      -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b)
+      -> Ranked n1 a -> Ranked n2 a
+rlift sn2 f (Ranked arr) = Ranked (mlift (ssxFromSNat sn2) f arr)
+
+-- | See the documentation of 'mlift2'.
+rlift2 :: forall n1 n2 n3 a. Elt a
+       => SNat n3
+       -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b -> XArray (Replicate n3 Nothing ++ sh') b)
+       -> Ranked n1 a -> Ranked n2 a -> Ranked n3 a
+rlift2 sn3 f (Ranked arr1) (Ranked arr2) = Ranked (mlift2 (ssxFromSNat sn3) f arr1 arr2)
+
+rsumOuter1P :: forall n a.
+               (Storable a, NumElt a)
+            => Ranked (n + 1) (Primitive a) -> Ranked n (Primitive a)
+rsumOuter1P (Ranked arr)
+  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = Ranked (msumOuter1P arr)
+
+rsumOuter1 :: forall n a. (NumElt a, PrimElt a)
+           => Ranked (n + 1) a -> Ranked n a
+rsumOuter1 = rfromPrimitive . rsumOuter1P . rtoPrimitive
+
+rsumAllPrim :: (PrimElt a, NumElt a) => Ranked n a -> a
+rsumAllPrim (Ranked arr) = msumAllPrim arr
+
+rtranspose :: forall n a. Elt a => PermR -> Ranked n a -> Ranked n a
+rtranspose perm arr
+  | sn@SNat <- rrank arr
+  , Dict <- lemKnownReplicate sn
+  , length perm <= fromIntegral (natVal (Proxy @n))
+  = rlift sn
+          (\ssh' -> X.transposeUntyped (natSing @n) ssh' perm)
+          arr
+  | otherwise
+  = error "Data.Array.Nested.rtranspose: Permutation longer than rank of array"
+
+rconcat :: forall n a. Elt a => NonEmpty (Ranked (n + 1) a) -> Ranked (n + 1) a
+rconcat
+  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = coerce mconcat
+
+rappend :: forall n a. Elt a
+        => Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked (n + 1) a
+rappend arr1 arr2
+  | sn@SNat <- rrank arr1
+  , Dict <- lemKnownReplicate sn
+  , Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = coerce (mappend @Nothing @Nothing @(Replicate n Nothing))
+      arr1 arr2
+
+rscalar :: Elt a => a -> Ranked 0 a
+rscalar x = Ranked (mscalar x)
+
+rfromVectorP :: forall n a. Storable a => IShR n -> VS.Vector a -> Ranked n (Primitive a)
+rfromVectorP sh v
+  | Dict <- lemKnownReplicate (shrRank sh)
+  = Ranked (mfromVectorP (shxFromShR sh) v)
+
+rfromVector :: forall n a. PrimElt a => IShR n -> VS.Vector a -> Ranked n a
+rfromVector sh v = rfromPrimitive (rfromVectorP sh v)
+
+rtoVectorP :: Storable a => Ranked n (Primitive a) -> VS.Vector a
+rtoVectorP = coerce mtoVectorP
+
+rtoVector :: PrimElt a => Ranked n a -> VS.Vector a
+rtoVector = coerce mtoVector
+
+rfromList1 :: Elt a => NonEmpty a -> Ranked 1 a
+rfromList1 l = Ranked (mfromList1 l)
+
+rfromListOuter :: forall n a. Elt a => NonEmpty (Ranked n a) -> Ranked (n + 1) a
+rfromListOuter l
+  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = Ranked (mfromListOuter (coerce l :: NonEmpty (Mixed (Replicate n Nothing) a)))
+
+rfromListLinear :: forall n a. Elt a => IShR n -> NonEmpty a -> Ranked n a
+rfromListLinear sh l = rreshape sh (rfromList1 l)
+
+rfromListPrim :: PrimElt a => [a] -> Ranked 1 a
+rfromListPrim l = Ranked (mfromListPrim l)
+
+rfromListPrimLinear :: PrimElt a => IShR n -> [a] -> Ranked n a
+rfromListPrimLinear sh l =
+  let M_Primitive _ xarr = toPrimitive (mfromListPrim l)
+  in Ranked $ fromPrimitive $ M_Primitive (shxFromShR sh) (X.reshape (SUnknown () :!% ZKX) (shxFromShR sh) xarr)
+
+rtoList :: Elt a => Ranked 1 a -> [a]
+rtoList = map runScalar . rtoListOuter
+
+rtoListOuter :: forall n a. Elt a => Ranked (n + 1) a -> [Ranked n a]
+rtoListOuter (Ranked arr)
+  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = coerce (mtoListOuter @a @Nothing @(Replicate n Nothing) arr)
+
+rtoListLinear :: Elt a => Ranked n a -> [a]
+rtoListLinear (Ranked arr) = mtoListLinear arr
+
+rfromOrthotope :: PrimElt a => SNat n -> S.Array n a -> Ranked n a
+rfromOrthotope sn arr
+  | Refl <- lemRankReplicate sn
+  = let xarr = XArray arr
+    in Ranked (fromPrimitive (M_Primitive (X.shape (ssxFromSNat sn) xarr) xarr))
+
+rtoOrthotope :: PrimElt a => Ranked n a -> S.Array n a
+rtoOrthotope (rtoPrimitive -> Ranked (M_Primitive sh (XArray arr)))
+  | Refl <- lemRankReplicate (shrRank $ shrFromShX2 sh)
+  = arr
+
+runScalar :: Elt a => Ranked 0 a -> a
+runScalar arr = rindex arr ZIR
+
+rnest :: forall n m a. Elt a => SNat n -> Ranked (n + m) a -> Ranked n (Ranked m a)
+rnest n arr
+  | Refl <- lemReplicatePlusApp n (Proxy @m) (Proxy @(Nothing @Nat))
+  = coerce (mnest (ssxFromSNat n) (coerce arr))
+
+runNest :: forall n m a. Elt a => Ranked n (Ranked m a) -> Ranked (n + m) a
+runNest rarr@(Ranked (M_Ranked (M_Nest _ arr)))
+  | Refl <- lemReplicatePlusApp (rrank rarr) (Proxy @m) (Proxy @(Nothing @Nat))
+  = Ranked arr
+
+rzip :: (Elt a, Elt b) => Ranked n a -> Ranked n b -> Ranked n (a, b)
+rzip = coerce mzip
+
+runzip :: Ranked n (a, b) -> (Ranked n a, Ranked n b)
+runzip = coerce munzip
+
+rrerankP :: forall n1 n2 n a b. (Storable a, Storable b)
+         => SNat n -> IShR n2
+         -> (Ranked n1 (Primitive a) -> Ranked n2 (Primitive b))
+         -> Ranked (n + n1) (Primitive a) -> Ranked (n + n2) (Primitive b)
+rrerankP sn sh2 f (Ranked arr)
+  | Refl <- lemReplicatePlusApp sn (Proxy @n1) (Proxy @(Nothing @Nat))
+  , Refl <- lemReplicatePlusApp sn (Proxy @n2) (Proxy @(Nothing @Nat))
+  = Ranked (mrerankP (ssxFromSNat sn) (shxFromShR sh2)
+                     (\a -> let Ranked r = f (Ranked a) in r)
+                     arr)
+
+-- | If there is a zero-sized dimension in the @n@-prefix of the shape of the
+-- input array, then there is no way to deduce the full shape of the output
+-- array (more precisely, the @n2@ part): that could only come from calling
+-- @f@, and there are no subarrays to call @f@ on. @orthotope@ errors out in
+-- this case; we choose to fill the @n2@ part of the output shape with zeros.
+--
+-- For example, if:
+--
+-- @
+-- arr :: Ranked 5 Int   -- of shape [3, 0, 4, 2, 21]
+-- f :: Ranked 2 Int -> Ranked 3 Float
+-- @
+--
+-- then:
+--
+-- @
+-- rrerank _ _ _ f arr :: Ranked 6 Float
+-- @
+--
+-- and this result will have shape @[3, 0, 4, 0, 0, 0]@. Note that the
+-- "reranked" part (the last 3 entries) are zero; we don't know if @f@ intended
+-- to return an array with shape all-0 here (it probably didn't), but there is
+-- no better number to put here absent a subarray of the input to pass to @f@.
+rrerank :: forall n1 n2 n a b. (PrimElt a, PrimElt b)
+        => SNat n -> IShR n2
+        -> (Ranked n1 a -> Ranked n2 b)
+        -> Ranked (n + n1) a -> Ranked (n + n2) b
+rrerank sn sh2 f (rtoPrimitive -> arr) =
+  rfromPrimitive $ rrerankP sn sh2 (rtoPrimitive . f . rfromPrimitive) arr
+
+rreplicate :: forall n m a. Elt a
+           => IShR n -> Ranked m a -> Ranked (n + m) a
+rreplicate sh (Ranked arr)
+  | Refl <- lemReplicatePlusApp (shrRank sh) (Proxy @m) (Proxy @(Nothing @Nat))
+  = Ranked (mreplicate (shxFromShR sh) arr)
+
+rreplicateScalP :: forall n a. Storable a => IShR n -> a -> Ranked n (Primitive a)
+rreplicateScalP sh x
+  | Dict <- lemKnownReplicate (shrRank sh)
+  = Ranked (mreplicateScalP (shxFromShR sh) x)
+
+rreplicateScal :: forall n a. PrimElt a
+               => IShR n -> a -> Ranked n a
+rreplicateScal sh x = rfromPrimitive (rreplicateScalP sh x)
+
+rslice :: forall n a. Elt a => Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a
+rslice i n arr
+  | Refl <- lemReplicateSucc @(Nothing @Nat) @n
+  = rlift (rrank arr)
+          (\_ -> X.sliceU i n)
+          arr
+
+rrev1 :: forall n a. Elt a => Ranked (n + 1) a -> Ranked (n + 1) a
+rrev1 arr =
+  rlift (rrank arr)
+        (\(_ :: StaticShX sh') ->
+          case lemReplicateSucc @(Nothing @Nat) @n of
+            Refl -> X.rev1 @Nothing @(Replicate n Nothing ++ sh'))
+        arr
+
+rreshape :: forall n n' a. Elt a
+         => IShR n' -> Ranked n a -> Ranked n' a
+rreshape sh' rarr@(Ranked arr)
+  | Dict <- lemKnownReplicate (rrank rarr)
+  , Dict <- lemKnownReplicate (shrRank sh')
+  = Ranked (mreshape (shxFromShR sh') arr)
+
+rflatten :: Elt a => Ranked n a -> Ranked 1 a
+rflatten (Ranked arr) = mtoRanked (mflatten arr)
+
+riota :: (Enum a, PrimElt a) => Int -> Ranked 1 a
+riota n = TN.withSomeSNat (fromIntegral n) $ mtoRanked . miota
+
+-- | Throws if the array is empty.
+rminIndexPrim :: (PrimElt a, NumElt a) => Ranked n a -> IIxR n
+rminIndexPrim rarr@(Ranked arr)
+  | Refl <- lemRankReplicate (rrank (rtoPrimitive rarr))
+  = ixrFromIxX (mminIndexPrim arr)
+
+-- | Throws if the array is empty.
+rmaxIndexPrim :: (PrimElt a, NumElt a) => Ranked n a -> IIxR n
+rmaxIndexPrim rarr@(Ranked arr)
+  | Refl <- lemRankReplicate (rrank (rtoPrimitive rarr))
+  = ixrFromIxX (mmaxIndexPrim arr)
+
+rdot1Inner :: forall n a. (PrimElt a, NumElt a) => Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked n a
+rdot1Inner arr1 arr2
+  | SNat <- rrank arr1
+  , Refl <- lemReplicatePlusApp (SNat @n) (Proxy @1) (Proxy @(Nothing @Nat))
+  = coerce (mdot1Inner (Proxy @(Nothing @Nat))) arr1 arr2
+
+-- | This has a temporary, suboptimal implementation in terms of 'mflatten'.
+-- Prefer 'rdot1Inner' if applicable.
+rdot :: (PrimElt a, NumElt a) => Ranked n a -> Ranked n a -> a
+rdot = coerce mdot
+
+rtoXArrayPrimP :: Ranked n (Primitive a) -> (IShR n, XArray (Replicate n Nothing) a)
+rtoXArrayPrimP (Ranked arr) = first shrFromShX2 (mtoXArrayPrimP arr)
+
+rtoXArrayPrim :: PrimElt a => Ranked n a -> (IShR n, XArray (Replicate n Nothing) a)
+rtoXArrayPrim (Ranked arr) = first shrFromShX2 (mtoXArrayPrim arr)
+
+rfromXArrayPrimP :: SNat n -> XArray (Replicate n Nothing) a -> Ranked n (Primitive a)
+rfromXArrayPrimP sn arr = Ranked (mfromXArrayPrimP (ssxFromShX (X.shape (ssxFromSNat sn) arr)) arr)
+
+rfromXArrayPrim :: PrimElt a => SNat n -> XArray (Replicate n Nothing) a -> Ranked n a
+rfromXArrayPrim sn arr = Ranked (mfromXArrayPrim (ssxFromShX (X.shape (ssxFromSNat sn) arr)) arr)
+
+rfromPrimitive :: PrimElt a => Ranked n (Primitive a) -> Ranked n a
+rfromPrimitive (Ranked arr) = Ranked (fromPrimitive arr)
+
+rtoPrimitive :: PrimElt a => Ranked n a -> Ranked n (Primitive a)
+rtoPrimitive (Ranked arr) = Ranked (toPrimitive arr)