{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.Array.Nested.Mixed.Shape where

import Control.DeepSeq (NFData(..))
import Data.Bifunctor (first)
import Data.Coerce
import Data.Foldable qualified as Foldable
import Data.Functor.Const
import Data.Functor.Product
import Data.Kind (Constraint, Type)
import Data.Monoid (Sum(..))
import Data.Type.Equality
import GHC.Exts (withDict)
import GHC.Generics (Generic)
import GHC.IsList (IsList)
import GHC.IsList qualified as IsList
import GHC.TypeLits

import Data.Array.Nested.Types


-- | The length of a type-level list. If the argument is a shape, then the
-- result is the rank of that shape.
type family Rank sh where
  Rank '[] = 0
  Rank (_ : sh) = Rank sh + 1


-- * Mixed lists

type role ListX nominal representational
type ListX :: [Maybe Nat] -> (Maybe Nat -> Type) -> Type
data ListX sh f where
  ZX :: ListX '[] f
  (::%) :: f n -> ListX sh f -> ListX (n : sh) f
deriving instance (forall n. Eq (f n)) => Eq (ListX sh f)
deriving instance (forall n. Ord (f n)) => Ord (ListX sh f)
infixr 3 ::%

#ifdef OXAR_DEFAULT_SHOW_INSTANCES
deriving instance (forall n. Show (f n)) => Show (ListX sh f)
#else
instance (forall n. Show (f n)) => Show (ListX sh f) where
  showsPrec :: Int -> ListX sh f -> ShowS
showsPrec Int
_ = (forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
(forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
listxShow f n -> ShowS
forall (n :: Maybe Nat). f n -> ShowS
forall a. Show a => a -> ShowS
shows
#endif

instance (forall n. NFData (f n)) => NFData (ListX sh f) where
  rnf :: ListX sh f -> ()
rnf ListX sh f
ZX = ()
  rnf (f n
x ::% ListX sh f
l) = f n -> ()
forall a. NFData a => a -> ()
rnf f n
x () -> () -> ()
forall a b. a -> b -> b
`seq` ListX sh f -> ()
forall a. NFData a => a -> ()
rnf ListX sh f
l

data UnconsListXRes f sh1 =
  forall n sh. (n : sh ~ sh1) => UnconsListXRes (ListX sh f) (f n)
listxUncons :: ListX sh1 f -> Maybe (UnconsListXRes f sh1)
listxUncons :: forall (sh1 :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh1 f -> Maybe (UnconsListXRes f sh1)
listxUncons (f n
i ::% ListX sh f
shl') = UnconsListXRes f sh1 -> Maybe (UnconsListXRes f sh1)
forall a. a -> Maybe a
Just (ListX sh f -> f n -> UnconsListXRes f sh1
forall (f :: Maybe Nat -> Type) (sh1 :: [Maybe Nat])
       (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
ListX sh f -> f n -> UnconsListXRes f sh1
UnconsListXRes ListX sh f
shl' f n
i)
listxUncons ListX sh1 f
ZX = Maybe (UnconsListXRes f sh1)
forall a. Maybe a
Nothing

-- | This checks only whether the types are equal; if the elements of the list
-- are not singletons, their values may still differ. This corresponds to
-- 'testEquality', except on the penultimate type parameter.
listxEqType :: TestEquality f => ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqType :: forall (f :: Maybe Nat -> Type) (sh :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
TestEquality f =>
ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqType ListX sh f
ZX ListX sh' f
ZX = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
listxEqType (f n
n ::% ListX sh f
sh) (f n
m ::% ListX sh f
sh')
  | Just n :~: n
Refl <- f n -> f n -> Maybe (n :~: n)
forall (a :: Maybe Nat) (b :: Maybe Nat).
f a -> f b -> Maybe (a :~: b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality f n
n f n
m
  , Just sh :~: sh
Refl <- ListX sh f -> ListX sh f -> Maybe (sh :~: sh)
forall (f :: Maybe Nat -> Type) (sh :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
TestEquality f =>
ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqType ListX sh f
sh ListX sh f
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
listxEqType ListX sh f
_ ListX sh' f
_ = Maybe (sh :~: sh')
forall a. Maybe a
Nothing

-- | This checks whether the two lists actually contain equal values. This is
-- more than 'testEquality', and corresponds to @geq@ from @Data.GADT.Compare@
-- in the @some@ package (except on the penultimate type parameter).
listxEqual :: (TestEquality f, forall n. Eq (f n)) => ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqual :: forall (f :: Maybe Nat -> Type) (sh :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
(TestEquality f, forall (n :: Maybe Nat). Eq (f n)) =>
ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqual ListX sh f
ZX ListX sh' f
ZX = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
listxEqual (f n
n ::% ListX sh f
sh) (f n
m ::% ListX sh f
sh')
  | Just n :~: n
Refl <- f n -> f n -> Maybe (n :~: n)
forall (a :: Maybe Nat) (b :: Maybe Nat).
f a -> f b -> Maybe (a :~: b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality f n
n f n
m
  , f n
n f n -> f n -> Bool
forall a. Eq a => a -> a -> Bool
== f n
f n
m
  , Just sh :~: sh
Refl <- ListX sh f -> ListX sh f -> Maybe (sh :~: sh)
forall (f :: Maybe Nat -> Type) (sh :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
(TestEquality f, forall (n :: Maybe Nat). Eq (f n)) =>
ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqual ListX sh f
sh ListX sh f
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
listxEqual ListX sh f
_ ListX sh' f
_ = Maybe (sh :~: sh')
forall a. Maybe a
Nothing

listxFmap :: (forall n. f n -> g n) -> ListX sh f -> ListX sh g
listxFmap :: forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat). f n -> g n) -> ListX sh f -> ListX sh g
listxFmap forall (n :: Maybe Nat). f n -> g n
_ ListX sh f
ZX = ListX sh g
ListX '[] g
forall (f :: Maybe Nat -> Type). ListX '[] f
ZX
listxFmap forall (n :: Maybe Nat). f n -> g n
f (f n
x ::% ListX sh f
xs) = f n -> g n
forall (n :: Maybe Nat). f n -> g n
f f n
x g n -> ListX sh g -> ListX (n : sh) g
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% (forall (n :: Maybe Nat). f n -> g n) -> ListX sh f -> ListX sh g
forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat). f n -> g n) -> ListX sh f -> ListX sh g
listxFmap f n -> g n
forall (n :: Maybe Nat). f n -> g n
f ListX sh f
xs

listxFold :: Monoid m => (forall n. f n -> m) -> ListX sh f -> m
listxFold :: forall m (f :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
Monoid m =>
(forall (n :: Maybe Nat). f n -> m) -> ListX sh f -> m
listxFold forall (n :: Maybe Nat). f n -> m
_ ListX sh f
ZX = m
forall a. Monoid a => a
mempty
listxFold forall (n :: Maybe Nat). f n -> m
f (f n
x ::% ListX sh f
xs) = f n -> m
forall (n :: Maybe Nat). f n -> m
f f n
x m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (forall (n :: Maybe Nat). f n -> m) -> ListX sh f -> m
forall m (f :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
Monoid m =>
(forall (n :: Maybe Nat). f n -> m) -> ListX sh f -> m
listxFold f n -> m
forall (n :: Maybe Nat). f n -> m
f ListX sh f
xs

listxLength :: ListX sh f -> Int
listxLength :: forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> Int
listxLength = Sum Int -> Int
forall a. Sum a -> a
getSum (Sum Int -> Int) -> (ListX sh f -> Sum Int) -> ListX sh f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (n :: Maybe Nat). f n -> Sum Int) -> ListX sh f -> Sum Int
forall m (f :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
Monoid m =>
(forall (n :: Maybe Nat). f n -> m) -> ListX sh f -> m
listxFold (\f n
_ -> Int -> Sum Int
forall a. a -> Sum a
Sum Int
1)

listxRank :: ListX sh f -> SNat (Rank sh)
listxRank :: forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank ListX sh f
ZX = SNat 0
SNat (Rank sh)
forall (n :: Nat). KnownNat n => SNat n
SNat
listxRank (f n
_ ::% ListX sh f
l) | SNat (Rank sh)
SNat <- ListX sh f -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank ListX sh f
l = SNat (Rank sh + 1)
SNat (Rank sh)
forall (n :: Nat). KnownNat n => SNat n
SNat

listxShow :: forall sh f. (forall n. f n -> ShowS) -> ListX sh f -> ShowS
listxShow :: forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
(forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
listxShow forall (n :: Maybe Nat). f n -> ShowS
f ListX sh f
l = String -> ShowS
showString String
"[" ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ListX sh f -> ShowS
forall (sh' :: [Maybe Nat]). String -> ListX sh' f -> ShowS
go String
"" ListX sh f
l ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
"]"
  where
    go :: String -> ListX sh' f -> ShowS
    go :: forall (sh' :: [Maybe Nat]). String -> ListX sh' f -> ShowS
go String
_ ListX sh' f
ZX = ShowS
forall a. a -> a
id
    go String
prefix (f n
x ::% ListX sh f
xs) = String -> ShowS
showString String
prefix ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f n -> ShowS
forall (n :: Maybe Nat). f n -> ShowS
f f n
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ListX sh f -> ShowS
forall (sh' :: [Maybe Nat]). String -> ListX sh' f -> ShowS
go String
"," ListX sh f
xs

listxFromList :: StaticShX sh -> [i] -> ListX sh (Const i)
listxFromList :: forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
listxFromList StaticShX sh
topssh [i]
topl = StaticShX sh -> [i] -> ListX sh (Const i)
forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
go StaticShX sh
topssh [i]
topl
  where
    go :: StaticShX sh' -> [i] -> ListX sh' (Const i)
    go :: forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
go StaticShX sh'
ZKX [] = ListX sh' (Const i)
ListX '[] (Const i)
forall (f :: Maybe Nat -> Type). ListX '[] f
ZX
    go (SMayNat () SNat n
_ :!% StaticShX sh
sh) (i
i : [i]
is) = i -> Const i n
forall {k} a (b :: k). a -> Const a b
Const i
i Const i n -> ListX sh (Const i) -> ListX (n : sh) (Const i)
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% StaticShX sh -> [i] -> ListX sh (Const i)
forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
go StaticShX sh
sh [i]
is
    go StaticShX sh'
_ [i]
_ = String -> ListX sh' (Const i)
forall a. HasCallStack => String -> a
error (String -> ListX sh' (Const i)) -> String -> ListX sh' (Const i)
forall a b. (a -> b) -> a -> b
$ String
"listxFromList: Mismatched list length (type says "
                       String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (StaticShX sh -> Int
forall (sh :: [Maybe Nat]). StaticShX sh -> Int
ssxLength StaticShX sh
topssh) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", list has length "
                       String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([i] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [i]
topl) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

listxToList :: ListX sh' (Const i) -> [i]
listxToList :: forall (sh' :: [Maybe Nat]) i. ListX sh' (Const i) -> [i]
listxToList ListX sh' (Const i)
ZX = []
listxToList (Const i
i ::% ListX sh (Const i)
is) = i
i i -> [i] -> [i]
forall a. a -> [a] -> [a]
: ListX sh (Const i) -> [i]
forall (sh' :: [Maybe Nat]) i. ListX sh' (Const i) -> [i]
listxToList ListX sh (Const i)
is

listxHead :: ListX (mn ': sh) f -> f mn
listxHead :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (f :: Maybe Nat -> Type).
ListX (mn : sh) f -> f mn
listxHead (f n
i ::% ListX sh f
_) = f mn
f n
i

listxTail :: ListX (n : sh) i -> ListX sh i
listxTail :: forall (n :: Maybe Nat) (sh :: [Maybe Nat])
       (i :: Maybe Nat -> Type).
ListX (n : sh) i -> ListX sh i
listxTail (i n
_ ::% ListX sh i
sh) = ListX sh i
ListX sh i
sh

listxAppend :: ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
listxAppend :: forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (sh' :: [Maybe Nat]).
ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
listxAppend ListX sh f
ZX ListX sh' f
idx' = ListX sh' f
ListX (sh ++ sh') f
idx'
listxAppend (f n
i ::% ListX sh f
idx) ListX sh' f
idx' = f n
i f n -> ListX (sh ++ sh') f -> ListX (n : (sh ++ sh')) f
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (sh' :: [Maybe Nat]).
ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
listxAppend ListX sh f
idx ListX sh' f
idx'

listxDrop :: forall f g sh sh'. ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop :: forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop ListX sh g
ZX ListX (sh ++ sh') f
long = ListX sh' f
ListX (sh ++ sh') f
long
listxDrop (g n
_ ::% ListX sh g
short) ListX (sh ++ sh') f
long = case ListX (sh ++ sh') f
long of f n
_ ::% ListX sh f
long' -> ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop ListX sh g
short ListX sh f
ListX (sh ++ sh') f
long'

listxInit :: forall f n sh. ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit :: forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit (f n
i ::% sh :: ListX sh f
sh@(f n
_ ::% ListX sh f
_)) = f n
i f n -> ListX (Init (n : sh)) f -> ListX (n : Init (n : sh)) f
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX (n : sh) f -> ListX (Init (n : sh)) f
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit ListX sh f
ListX (n : sh) f
sh
listxInit (f n
_ ::% ListX sh f
ZX) = ListX '[] f
ListX (Init (n : sh)) f
forall (f :: Maybe Nat -> Type). ListX '[] f
ZX

listxLast :: forall f n sh. ListX (n : sh) f -> f (Last (n : sh))
listxLast :: forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> f (Last (n : sh))
listxLast (f n
_ ::% sh :: ListX sh f
sh@(f n
_ ::% ListX sh f
_)) = ListX (n : sh) f -> f (Last (n : sh))
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> f (Last (n : sh))
listxLast ListX sh f
ListX (n : sh) f
sh
listxLast (f n
x ::% ListX sh f
ZX) = f n
f (Last (n : sh))
x

listxZip :: ListX sh f -> ListX sh g -> ListX sh (Product f g)
listxZip :: forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (g :: Maybe Nat -> Type).
ListX sh f -> ListX sh g -> ListX sh (Product f g)
listxZip ListX sh f
ZX ListX sh g
ZX = ListX sh (Product f g)
ListX '[] (Product f g)
forall (f :: Maybe Nat -> Type). ListX '[] f
ZX
listxZip (f n
i ::% ListX sh f
irest) (g n
j ::% ListX sh g
jrest) =
  f n -> g n -> Product f g n
forall {k} (f :: k -> Type) (g :: k -> Type) (a :: k).
f a -> g a -> Product f g a
Pair f n
i g n
g n
j Product f g n
-> ListX sh (Product f g) -> ListX (n : sh) (Product f g)
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX sh f -> ListX sh g -> ListX sh (Product f g)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (g :: Maybe Nat -> Type).
ListX sh f -> ListX sh g -> ListX sh (Product f g)
listxZip ListX sh f
irest ListX sh g
ListX sh g
jrest

listxZipWith :: (forall a. f a -> g a -> h a) -> ListX sh f -> ListX sh g
             -> ListX sh h
listxZipWith :: forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (h :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
(forall (a :: Maybe Nat). f a -> g a -> h a)
-> ListX sh f -> ListX sh g -> ListX sh h
listxZipWith forall (a :: Maybe Nat). f a -> g a -> h a
_ ListX sh f
ZX ListX sh g
ZX = ListX sh h
ListX '[] h
forall (f :: Maybe Nat -> Type). ListX '[] f
ZX
listxZipWith forall (a :: Maybe Nat). f a -> g a -> h a
f (f n
i ::% ListX sh f
is) (g n
j ::% ListX sh g
js) =
  f n -> g n -> h n
forall (a :: Maybe Nat). f a -> g a -> h a
f f n
i g n
g n
j h n -> ListX sh h -> ListX (n : sh) h
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% (forall (a :: Maybe Nat). f a -> g a -> h a)
-> ListX sh f -> ListX sh g -> ListX sh h
forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (h :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
(forall (a :: Maybe Nat). f a -> g a -> h a)
-> ListX sh f -> ListX sh g -> ListX sh h
listxZipWith f a -> g a -> h a
forall (a :: Maybe Nat). f a -> g a -> h a
f ListX sh f
is ListX sh g
ListX sh g
js


-- * Mixed indices

-- | An index into a mixed-typed array.
type role IxX nominal representational
type IxX :: [Maybe Nat] -> Type -> Type
newtype IxX sh i = IxX (ListX sh (Const i))
  deriving (IxX sh i -> IxX sh i -> Bool
(IxX sh i -> IxX sh i -> Bool)
-> (IxX sh i -> IxX sh i -> Bool) -> Eq (IxX sh i)
forall (sh :: [Maybe Nat]) i. Eq i => IxX sh i -> IxX sh i -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (sh :: [Maybe Nat]) i. Eq i => IxX sh i -> IxX sh i -> Bool
== :: IxX sh i -> IxX sh i -> Bool
$c/= :: forall (sh :: [Maybe Nat]) i. Eq i => IxX sh i -> IxX sh i -> Bool
/= :: IxX sh i -> IxX sh i -> Bool
Eq, Eq (IxX sh i)
Eq (IxX sh i) =>
(IxX sh i -> IxX sh i -> Ordering)
-> (IxX sh i -> IxX sh i -> Bool)
-> (IxX sh i -> IxX sh i -> Bool)
-> (IxX sh i -> IxX sh i -> Bool)
-> (IxX sh i -> IxX sh i -> Bool)
-> (IxX sh i -> IxX sh i -> IxX sh i)
-> (IxX sh i -> IxX sh i -> IxX sh i)
-> Ord (IxX sh i)
IxX sh i -> IxX sh i -> Bool
IxX sh i -> IxX sh i -> Ordering
IxX sh i -> IxX sh i -> IxX sh i
forall (sh :: [Maybe Nat]) i. Ord i => Eq (IxX sh i)
forall (sh :: [Maybe Nat]) i. Ord i => IxX sh i -> IxX sh i -> Bool
forall (sh :: [Maybe Nat]) i.
Ord i =>
IxX sh i -> IxX sh i -> Ordering
forall (sh :: [Maybe Nat]) i.
Ord i =>
IxX sh i -> IxX sh i -> IxX sh i
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (sh :: [Maybe Nat]) i.
Ord i =>
IxX sh i -> IxX sh i -> Ordering
compare :: IxX sh i -> IxX sh i -> Ordering
$c< :: forall (sh :: [Maybe Nat]) i. Ord i => IxX sh i -> IxX sh i -> Bool
< :: IxX sh i -> IxX sh i -> Bool
$c<= :: forall (sh :: [Maybe Nat]) i. Ord i => IxX sh i -> IxX sh i -> Bool
<= :: IxX sh i -> IxX sh i -> Bool
$c> :: forall (sh :: [Maybe Nat]) i. Ord i => IxX sh i -> IxX sh i -> Bool
> :: IxX sh i -> IxX sh i -> Bool
$c>= :: forall (sh :: [Maybe Nat]) i. Ord i => IxX sh i -> IxX sh i -> Bool
>= :: IxX sh i -> IxX sh i -> Bool
$cmax :: forall (sh :: [Maybe Nat]) i.
Ord i =>
IxX sh i -> IxX sh i -> IxX sh i
max :: IxX sh i -> IxX sh i -> IxX sh i
$cmin :: forall (sh :: [Maybe Nat]) i.
Ord i =>
IxX sh i -> IxX sh i -> IxX sh i
min :: IxX sh i -> IxX sh i -> IxX sh i
Ord, (forall x. IxX sh i -> Rep (IxX sh i) x)
-> (forall x. Rep (IxX sh i) x -> IxX sh i) -> Generic (IxX sh i)
forall (sh :: [Maybe Nat]) i x. Rep (IxX sh i) x -> IxX sh i
forall (sh :: [Maybe Nat]) i x. IxX sh i -> Rep (IxX sh i) x
forall x. Rep (IxX sh i) x -> IxX sh i
forall x. IxX sh i -> Rep (IxX sh i) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall (sh :: [Maybe Nat]) i x. IxX sh i -> Rep (IxX sh i) x
from :: forall x. IxX sh i -> Rep (IxX sh i) x
$cto :: forall (sh :: [Maybe Nat]) i x. Rep (IxX sh i) x -> IxX sh i
to :: forall x. Rep (IxX sh i) x -> IxX sh i
Generic)

pattern ZIX :: forall sh i. () => sh ~ '[] => IxX sh i
pattern $mZIX :: forall {r} {sh :: [Maybe Nat]} {i}.
IxX sh i -> ((sh ~ '[]) => r) -> ((# #) -> r) -> r
$bZIX :: forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX = IxX ZX

pattern (:.%)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => i -> IxX sh i -> IxX sh1 i
pattern i $m:.% :: forall {r} {sh1 :: [Maybe Nat]} {i}.
IxX sh1 i
-> (forall {n :: Maybe Nat} {sh :: [Maybe Nat]}.
    ((n : sh) ~ sh1) =>
    i -> IxX sh i -> r)
-> ((# #) -> r)
-> r
$b:.% :: forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% shl <- IxX (listxUncons -> Just (UnconsListXRes (IxX -> shl) (getConst -> i)))
  where i
i :.% IxX ListX sh (Const i)
shl = ListX sh1 (Const i) -> IxX sh1 i
forall (sh :: [Maybe Nat]) i. ListX sh (Const i) -> IxX sh i
IxX (i -> Const i n
forall {k} a (b :: k). a -> Const a b
Const i
i Const i n -> ListX sh (Const i) -> ListX (n : sh) (Const i)
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX sh (Const i)
shl)
infixr 3 :.%

{-# COMPLETE ZIX, (:.%) #-}

-- For convenience, this contains regular 'Int's instead of bounded integers
-- (traditionally called \"@Fin@\").
type IIxX sh = IxX sh Int

#ifdef OXAR_DEFAULT_SHOW_INSTANCES
deriving instance Show i => Show (IxX sh i)
#else
instance Show i => Show (IxX sh i) where
  showsPrec :: Int -> IxX sh i -> ShowS
showsPrec Int
_ (IxX ListX sh (Const i)
l) = (forall (n :: Maybe Nat). Const i n -> ShowS)
-> ListX sh (Const i) -> ShowS
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
(forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
listxShow (i -> ShowS
forall a. Show a => a -> ShowS
shows (i -> ShowS) -> (Const i n -> i) -> Const i n -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Const i n -> i
forall {k} a (b :: k). Const a b -> a
getConst) ListX sh (Const i)
l
#endif

instance Functor (IxX sh) where
  fmap :: forall a b. (a -> b) -> IxX sh a -> IxX sh b
fmap a -> b
f (IxX ListX sh (Const a)
l) = ListX sh (Const b) -> IxX sh b
forall (sh :: [Maybe Nat]) i. ListX sh (Const i) -> IxX sh i
IxX ((forall (n :: Maybe Nat). Const a n -> Const b n)
-> ListX sh (Const a) -> ListX sh (Const b)
forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat). f n -> g n) -> ListX sh f -> ListX sh g
listxFmap (b -> Const b n
forall {k} a (b :: k). a -> Const a b
Const (b -> Const b n) -> (Const a n -> b) -> Const a n -> Const b n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f (a -> b) -> (Const a n -> a) -> Const a n -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Const a n -> a
forall {k} a (b :: k). Const a b -> a
getConst) ListX sh (Const a)
l)

instance Foldable (IxX sh) where
  foldMap :: forall m a. Monoid m => (a -> m) -> IxX sh a -> m
foldMap a -> m
f (IxX ListX sh (Const a)
l) = (forall (n :: Maybe Nat). Const a n -> m)
-> ListX sh (Const a) -> m
forall m (f :: Maybe Nat -> Type) (sh :: [Maybe Nat]).
Monoid m =>
(forall (n :: Maybe Nat). f n -> m) -> ListX sh f -> m
listxFold (a -> m
f (a -> m) -> (Const a n -> a) -> Const a n -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Const a n -> a
forall {k} a (b :: k). Const a b -> a
getConst) ListX sh (Const a)
l

instance NFData i => NFData (IxX sh i)

ixxLength :: IxX sh i -> Int
ixxLength :: forall (sh :: [Maybe Nat]) a. IxX sh a -> Int
ixxLength (IxX ListX sh (Const i)
l) = ListX sh (Const i) -> Int
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> Int
listxLength ListX sh (Const i)
l

ixxRank :: IxX sh i -> SNat (Rank sh)
ixxRank :: forall (sh :: [Maybe Nat]) i. IxX sh i -> SNat (Rank sh)
ixxRank (IxX ListX sh (Const i)
l) = ListX sh (Const i) -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank ListX sh (Const i)
l

ixxZero :: StaticShX sh -> IIxX sh
ixxZero :: forall (sh :: [Maybe Nat]). StaticShX sh -> IIxX sh
ixxZero StaticShX sh
ZKX = IxX sh Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxZero (SMayNat () SNat n
_ :!% StaticShX sh
ssh) = Int
0 Int -> IxX sh Int -> IxX sh Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% StaticShX sh -> IxX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> IIxX sh
ixxZero StaticShX sh
ssh

ixxZero' :: IShX sh -> IIxX sh
ixxZero' :: forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh
ixxZero' ShX sh Int
ZSX = IxX sh Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxZero' (SMayNat Int SNat n
_ :$% ShX sh Int
sh) = Int
0 Int -> IxX sh Int -> IxX sh Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% ShX sh Int -> IxX sh Int
forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh
ixxZero' ShX sh Int
sh

ixxFromList :: forall sh i. StaticShX sh -> [i] -> IxX sh i
ixxFromList :: forall (sh :: [Maybe Nat]) i. StaticShX sh -> [i] -> IxX sh i
ixxFromList = (StaticShX sh -> [i] -> ListX sh (Const i))
-> StaticShX sh -> [i] -> IxX sh i
forall a b. Coercible a b => a -> b
coerce (forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
listxFromList @_ @i)

ixxHead :: IxX (n : sh) i -> i
ixxHead :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i. IxX (n : sh) i -> i
ixxHead (IxX ListX (n : sh) (Const i)
list) = Const i n -> i
forall {k} a (b :: k). Const a b -> a
getConst (ListX (n : sh) (Const i) -> Const i n
forall (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (f :: Maybe Nat -> Type).
ListX (mn : sh) f -> f mn
listxHead ListX (n : sh) (Const i)
list)

ixxTail :: IxX (n : sh) i -> IxX sh i
ixxTail :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
IxX (n : sh) i -> IxX sh i
ixxTail (IxX ListX (n : sh) (Const i)
list) = ListX sh (Const i) -> IxX sh i
forall (sh :: [Maybe Nat]) i. ListX sh (Const i) -> IxX sh i
IxX (ListX (n : sh) (Const i) -> ListX sh (Const i)
forall (n :: Maybe Nat) (sh :: [Maybe Nat])
       (i :: Maybe Nat -> Type).
ListX (n : sh) i -> ListX sh i
listxTail ListX (n : sh) (Const i)
list)

ixxAppend :: forall sh sh' i. IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
ixxAppend :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
ixxAppend = (ListX sh (Const i)
 -> ListX sh' (Const i) -> ListX (sh ++ sh') (Const i))
-> IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
forall a b. Coercible a b => a -> b
coerce (forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (sh' :: [Maybe Nat]).
ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
listxAppend @_ @(Const i))

ixxDrop :: forall sh sh' i. IxX sh i -> IxX (sh ++ sh') i -> IxX sh' i
ixxDrop :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
IxX sh i -> IxX (sh ++ sh') i -> IxX sh' i
ixxDrop = (ListX sh (Const i)
 -> ListX (sh ++ sh') (Const i) -> ListX sh' (Const i))
-> IxX sh i -> IxX (sh ++ sh') i -> IxX sh' i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(Const i) @(Const i))

ixxInit :: forall n sh i. IxX (n : sh) i -> IxX (Init (n : sh)) i
ixxInit :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
IxX (n : sh) i -> IxX (Init (n : sh)) i
ixxInit = (ListX (n : sh) (Const i) -> ListX (Init (n : sh)) (Const i))
-> IxX (n : sh) i -> IxX (Init (n : sh)) i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit @(Const i))

ixxLast :: forall n sh i. IxX (n : sh) i -> i
ixxLast :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i. IxX (n : sh) i -> i
ixxLast = (ListX (n : sh) (Const i) -> Const i (Last (n : sh)))
-> IxX (n : sh) i -> i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> f (Last (n : sh))
listxLast @(Const i))

ixxCast :: StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast :: forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX sh'
ZKX IxX sh i
ZIX = IxX sh' i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxCast (SMayNat () SNat n
_ :!% StaticShX sh
sh) (i
i :.% IxX sh i
idx) = i
i i -> IxX sh i -> IxX sh' i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% StaticShX sh -> IxX sh i -> IxX sh i
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX sh
sh IxX sh i
idx
ixxCast StaticShX sh'
_ IxX sh i
_ = String -> IxX sh' i
forall a. HasCallStack => String -> a
error String
"ixxCast: ranks don't match"

ixxZip :: IxX sh i -> IxX sh j -> IxX sh (i, j)
ixxZip :: forall (sh :: [Maybe Nat]) i j.
IxX sh i -> IxX sh j -> IxX sh (i, j)
ixxZip IxX sh i
ZIX IxX sh j
ZIX = IxX sh (i, j)
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxZip (i
i :.% IxX sh i
is) (j
j :.% IxX sh j
js) = (i
i, j
j) (i, j) -> IxX sh (i, j) -> IxX sh (i, j)
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% IxX sh i -> IxX sh j -> IxX sh (i, j)
forall (sh :: [Maybe Nat]) i j.
IxX sh i -> IxX sh j -> IxX sh (i, j)
ixxZip IxX sh i
is IxX sh j
IxX sh j
js

ixxZipWith :: (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
ixxZipWith :: forall i j k (sh :: [Maybe Nat]).
(i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
ixxZipWith i -> j -> k
_ IxX sh i
ZIX IxX sh j
ZIX = IxX sh k
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxZipWith i -> j -> k
f (i
i :.% IxX sh i
is) (j
j :.% IxX sh j
js) = i -> j -> k
f i
i j
j k -> IxX sh k -> IxX sh k
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
forall i j k (sh :: [Maybe Nat]).
(i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
ixxZipWith i -> j -> k
f IxX sh i
is IxX sh j
IxX sh j
js

ixxFromLinear :: IShX sh -> Int -> IIxX sh
ixxFromLinear :: forall (sh :: [Maybe Nat]). IShX sh -> Int -> IIxX sh
ixxFromLinear = \IShX sh
sh Int
i -> case IShX sh -> Int -> (IIxX sh, Int)
forall (sh :: [Maybe Nat]). IShX sh -> Int -> (IIxX sh, Int)
go IShX sh
sh Int
i of
  (IIxX sh
idx, Int
0) -> IIxX sh
idx
  (IIxX sh, Int)
_ -> String -> IIxX sh
forall a. HasCallStack => String -> a
error (String -> IIxX sh) -> String -> IIxX sh
forall a b. (a -> b) -> a -> b
$ String
"ixxFromLinear: out of range (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++
               String
" in array of shape " String -> ShowS
forall a. [a] -> [a] -> [a]
++ IShX sh -> String
forall a. Show a => a -> String
show IShX sh
sh String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
  where
    -- returns (index in subarray, remaining index in enclosing array)
    go :: IShX sh -> Int -> (IIxX sh, Int)
    go :: forall (sh :: [Maybe Nat]). IShX sh -> Int -> (IIxX sh, Int)
go ShX sh Int
ZSX Int
i = (IxX sh Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX, Int
i)
    go (SMayNat Int SNat n
n :$% ShX sh Int
sh) Int
i =
      let (IIxX sh
idx, Int
i') = ShX sh Int -> Int -> (IIxX sh, Int)
forall (sh :: [Maybe Nat]). IShX sh -> Int -> (IIxX sh, Int)
go ShX sh Int
sh Int
i
          (Int
upi, Int
locali) = Int
i' Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` SMayNat Int SNat n -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat n
n
      in (Int
locali Int -> IIxX sh -> IxX sh Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% IIxX sh
idx, Int
upi)

ixxToLinear :: IShX sh -> IIxX sh -> Int
ixxToLinear :: forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh -> Int
ixxToLinear = \IShX sh
sh IIxX sh
i -> (Int, Int) -> Int
forall a b. (a, b) -> a
fst (IShX sh -> IIxX sh -> (Int, Int)
forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh -> (Int, Int)
go IShX sh
sh IIxX sh
i)
  where
    -- returns (index in subarray, size of subarray)
    go :: IShX sh -> IIxX sh -> (Int, Int)
    go :: forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh -> (Int, Int)
go ShX sh Int
ZSX IxX sh Int
ZIX = (Int
0, Int
1)
    go (SMayNat Int SNat n
n :$% ShX sh Int
sh) (Int
i :.% IxX sh Int
ix) =
      let (Int
lidx, Int
sz) = ShX sh Int -> IIxX sh -> (Int, Int)
forall (sh :: [Maybe Nat]). IShX sh -> IIxX sh -> (Int, Int)
go ShX sh Int
sh IIxX sh
IxX sh Int
ix
      in (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lidx, SMayNat Int SNat n -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
sz)


-- * Mixed shapes

data SMayNat i f n where
  SUnknown :: i -> SMayNat i f Nothing
  SKnown :: f n -> SMayNat i f (Just n)
deriving instance (Show i, forall m. Show (f m)) => Show (SMayNat i f n)
deriving instance (Eq i, forall m. Eq (f m)) => Eq (SMayNat i f n)
deriving instance (Ord i, forall m. Ord (f m)) => Ord (SMayNat i f n)

instance (NFData i, forall m. NFData (f m)) => NFData (SMayNat i f n) where
  rnf :: SMayNat i f n -> ()
rnf (SUnknown i
i) = i -> ()
forall a. NFData a => a -> ()
rnf i
i
  rnf (SKnown f n
x) = f n -> ()
forall a. NFData a => a -> ()
rnf f n
x

instance TestEquality f => TestEquality (SMayNat i f) where
  testEquality :: forall (a :: Maybe k) (b :: Maybe k).
SMayNat i f a -> SMayNat i f b -> Maybe (a :~: b)
testEquality SUnknown{} SUnknown{} = (a :~: b) -> Maybe (a :~: b)
forall a. a -> Maybe a
Just a :~: a
a :~: b
forall {k} (a :: k). a :~: a
Refl
  testEquality (SKnown f n
n) (SKnown f n
m) | Just n :~: n
Refl <- f n -> f n -> Maybe (n :~: n)
forall (a :: k) (b :: k). f a -> f b -> Maybe (a :~: b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality f n
n f n
m = (a :~: b) -> Maybe (a :~: b)
forall a. a -> Maybe a
Just a :~: a
a :~: b
forall {k} (a :: k). a :~: a
Refl
  testEquality SMayNat i f a
_ SMayNat i f b
_ = Maybe (a :~: b)
forall a. Maybe a
Nothing

fromSMayNat :: (n ~ Nothing => i -> r)
            -> (forall m. n ~ Just m => f m -> r)
            -> SMayNat i f n -> r
fromSMayNat :: forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat (n ~ 'Nothing) => i -> r
f forall (m :: k). (n ~ 'Just m) => f m -> r
_ (SUnknown i
i) = i -> r
(n ~ 'Nothing) => i -> r
f i
i
fromSMayNat (n ~ 'Nothing) => i -> r
_ forall (m :: k). (n ~ 'Just m) => f m -> r
g (SKnown f n
s) = f n -> r
forall (m :: k). (n ~ 'Just m) => f m -> r
g f n
s

fromSMayNat' :: SMayNat Int SNat n -> Int
fromSMayNat' :: forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' = ((n ~ 'Nothing) => Int -> Int)
-> (forall (m :: Nat). (n ~ 'Just m) => SNat m -> Int)
-> SMayNat Int SNat n
-> Int
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat Int -> Int
(n ~ 'Nothing) => Int -> Int
forall a. a -> a
id SNat m -> Int
forall (m :: Nat). (n ~ 'Just m) => SNat m -> Int
forall (n :: Nat). SNat n -> Int
fromSNat'

type family AddMaybe n m where
  AddMaybe Nothing _ = Nothing
  AddMaybe (Just _) Nothing = Nothing
  AddMaybe (Just n) (Just m) = Just (n + m)

smnAddMaybe :: SMayNat Int SNat n -> SMayNat Int SNat m -> SMayNat Int SNat (AddMaybe n m)
smnAddMaybe :: forall (n :: Maybe Nat) (m :: Maybe Nat).
SMayNat Int SNat n
-> SMayNat Int SNat m -> SMayNat Int SNat (AddMaybe n m)
smnAddMaybe (SUnknown Int
n) SMayNat Int SNat m
m = Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ SMayNat Int SNat m -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat m
m)
smnAddMaybe (SKnown SNat n
n) (SUnknown Int
m) = Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown (SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m)
smnAddMaybe (SKnown SNat n
n) (SKnown SNat n
m) = SNat (n + n) -> SMayNat Int SNat ('Just (n + n))
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown (SNat n -> SNat n -> SNat (n + n)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n
n SNat n
m)


-- | This is a newtype over 'ListX'.
type role ShX nominal representational
type ShX :: [Maybe Nat] -> Type -> Type
newtype ShX sh i = ShX (ListX sh (SMayNat i SNat))
  deriving (ShX sh i -> ShX sh i -> Bool
(ShX sh i -> ShX sh i -> Bool)
-> (ShX sh i -> ShX sh i -> Bool) -> Eq (ShX sh i)
forall (sh :: [Maybe Nat]) i. Eq i => ShX sh i -> ShX sh i -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (sh :: [Maybe Nat]) i. Eq i => ShX sh i -> ShX sh i -> Bool
== :: ShX sh i -> ShX sh i -> Bool
$c/= :: forall (sh :: [Maybe Nat]) i. Eq i => ShX sh i -> ShX sh i -> Bool
/= :: ShX sh i -> ShX sh i -> Bool
Eq, Eq (ShX sh i)
Eq (ShX sh i) =>
(ShX sh i -> ShX sh i -> Ordering)
-> (ShX sh i -> ShX sh i -> Bool)
-> (ShX sh i -> ShX sh i -> Bool)
-> (ShX sh i -> ShX sh i -> Bool)
-> (ShX sh i -> ShX sh i -> Bool)
-> (ShX sh i -> ShX sh i -> ShX sh i)
-> (ShX sh i -> ShX sh i -> ShX sh i)
-> Ord (ShX sh i)
ShX sh i -> ShX sh i -> Bool
ShX sh i -> ShX sh i -> Ordering
ShX sh i -> ShX sh i -> ShX sh i
forall (sh :: [Maybe Nat]) i. Ord i => Eq (ShX sh i)
forall (sh :: [Maybe Nat]) i. Ord i => ShX sh i -> ShX sh i -> Bool
forall (sh :: [Maybe Nat]) i.
Ord i =>
ShX sh i -> ShX sh i -> Ordering
forall (sh :: [Maybe Nat]) i.
Ord i =>
ShX sh i -> ShX sh i -> ShX sh i
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (sh :: [Maybe Nat]) i.
Ord i =>
ShX sh i -> ShX sh i -> Ordering
compare :: ShX sh i -> ShX sh i -> Ordering
$c< :: forall (sh :: [Maybe Nat]) i. Ord i => ShX sh i -> ShX sh i -> Bool
< :: ShX sh i -> ShX sh i -> Bool
$c<= :: forall (sh :: [Maybe Nat]) i. Ord i => ShX sh i -> ShX sh i -> Bool
<= :: ShX sh i -> ShX sh i -> Bool
$c> :: forall (sh :: [Maybe Nat]) i. Ord i => ShX sh i -> ShX sh i -> Bool
> :: ShX sh i -> ShX sh i -> Bool
$c>= :: forall (sh :: [Maybe Nat]) i. Ord i => ShX sh i -> ShX sh i -> Bool
>= :: ShX sh i -> ShX sh i -> Bool
$cmax :: forall (sh :: [Maybe Nat]) i.
Ord i =>
ShX sh i -> ShX sh i -> ShX sh i
max :: ShX sh i -> ShX sh i -> ShX sh i
$cmin :: forall (sh :: [Maybe Nat]) i.
Ord i =>
ShX sh i -> ShX sh i -> ShX sh i
min :: ShX sh i -> ShX sh i -> ShX sh i
Ord, (forall x. ShX sh i -> Rep (ShX sh i) x)
-> (forall x. Rep (ShX sh i) x -> ShX sh i) -> Generic (ShX sh i)
forall (sh :: [Maybe Nat]) i x. Rep (ShX sh i) x -> ShX sh i
forall (sh :: [Maybe Nat]) i x. ShX sh i -> Rep (ShX sh i) x
forall x. Rep (ShX sh i) x -> ShX sh i
forall x. ShX sh i -> Rep (ShX sh i) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall (sh :: [Maybe Nat]) i x. ShX sh i -> Rep (ShX sh i) x
from :: forall x. ShX sh i -> Rep (ShX sh i) x
$cto :: forall (sh :: [Maybe Nat]) i x. Rep (ShX sh i) x -> ShX sh i
to :: forall x. Rep (ShX sh i) x -> ShX sh i
Generic)

pattern ZSX :: forall sh i. () => sh ~ '[] => ShX sh i
pattern $mZSX :: forall {r} {sh :: [Maybe Nat]} {i}.
ShX sh i -> ((sh ~ '[]) => r) -> ((# #) -> r) -> r
$bZSX :: forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX = ShX ZX

pattern (:$%)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => SMayNat i SNat n -> ShX sh i -> ShX sh1 i
pattern i $m:$% :: forall {r} {sh1 :: [Maybe Nat]} {i}.
ShX sh1 i
-> (forall {n :: Maybe Nat} {sh :: [Maybe Nat]}.
    ((n : sh) ~ sh1) =>
    SMayNat i SNat n -> ShX sh i -> r)
-> ((# #) -> r)
-> r
$b:$% :: forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% shl <- ShX (listxUncons -> Just (UnconsListXRes (ShX -> shl) i))
  where SMayNat i SNat n
i :$% ShX ListX sh (SMayNat i SNat)
shl = ListX sh1 (SMayNat i SNat) -> ShX sh1 i
forall (sh :: [Maybe Nat]) i. ListX sh (SMayNat i SNat) -> ShX sh i
ShX (SMayNat i SNat n
i SMayNat i SNat n
-> ListX sh (SMayNat i SNat) -> ListX (n : sh) (SMayNat i SNat)
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX sh (SMayNat i SNat)
shl)
infixr 3 :$%

{-# COMPLETE ZSX, (:$%) #-}

type IShX sh = ShX sh Int

#ifdef OXAR_DEFAULT_SHOW_INSTANCES
deriving instance Show i => Show (ShX sh i)
#else
instance Show i => Show (ShX sh i) where
  showsPrec :: Int -> ShX sh i -> ShowS
showsPrec Int
_ (ShX ListX sh (SMayNat i SNat)
l) = (forall (n :: Maybe Nat). SMayNat i SNat n -> ShowS)
-> ListX sh (SMayNat i SNat) -> ShowS
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
(forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
listxShow (((n ~ 'Nothing) => i -> ShowS)
-> (forall (m :: Nat). (n ~ 'Just m) => SNat m -> ShowS)
-> SMayNat i SNat n
-> ShowS
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat i -> ShowS
(n ~ 'Nothing) => i -> ShowS
forall a. Show a => a -> ShowS
shows (Integer -> ShowS
forall a. Show a => a -> ShowS
shows (Integer -> ShowS) -> (SNat m -> Integer) -> SNat m -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat m -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat)) ListX sh (SMayNat i SNat)
l
#endif

instance Functor (ShX sh) where
  fmap :: forall a b. (a -> b) -> ShX sh a -> ShX sh b
fmap a -> b
f (ShX ListX sh (SMayNat a SNat)
l) = ListX sh (SMayNat b SNat) -> ShX sh b
forall (sh :: [Maybe Nat]) i. ListX sh (SMayNat i SNat) -> ShX sh i
ShX ((forall (n :: Maybe Nat). SMayNat a SNat n -> SMayNat b SNat n)
-> ListX sh (SMayNat a SNat) -> ListX sh (SMayNat b SNat)
forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat). f n -> g n) -> ListX sh f -> ListX sh g
listxFmap (((n ~ 'Nothing) => a -> SMayNat b SNat n)
-> (forall (m :: Nat). (n ~ 'Just m) => SNat m -> SMayNat b SNat n)
-> SMayNat a SNat n
-> SMayNat b SNat n
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat (b -> SMayNat b SNat n
b -> SMayNat b SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown (b -> SMayNat b SNat n) -> (a -> b) -> a -> SMayNat b SNat n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) SNat m -> SMayNat b SNat n
SNat m -> SMayNat b SNat ('Just m)
forall (m :: Nat). (n ~ 'Just m) => SNat m -> SMayNat b SNat n
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown) ListX sh (SMayNat a SNat)
l)

instance NFData i => NFData (ShX sh i) where
  rnf :: ShX sh i -> ()
rnf (ShX ListX sh (SMayNat i SNat)
ZX) = ()
  rnf (ShX (SUnknown i
i ::% ListX sh (SMayNat i SNat)
l)) = i -> ()
forall a. NFData a => a -> ()
rnf i
i () -> () -> ()
forall a b. a -> b -> b
`seq` ShX sh i -> ()
forall a. NFData a => a -> ()
rnf (ListX sh (SMayNat i SNat) -> ShX sh i
forall (sh :: [Maybe Nat]) i. ListX sh (SMayNat i SNat) -> ShX sh i
ShX ListX sh (SMayNat i SNat)
l)
  rnf (ShX (SKnown SNat n
SNat ::% ListX sh (SMayNat i SNat)
l)) = ShX sh i -> ()
forall a. NFData a => a -> ()
rnf (ListX sh (SMayNat i SNat) -> ShX sh i
forall (sh :: [Maybe Nat]) i. ListX sh (SMayNat i SNat) -> ShX sh i
ShX ListX sh (SMayNat i SNat)
l)

-- | This checks only whether the types are equal; unknown dimensions might
-- still differ. This corresponds to 'testEquality', except on the penultimate
-- type parameter.
shxEqType :: ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqType :: forall (sh :: [Maybe Nat]) i (sh' :: [Maybe Nat]).
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqType ShX sh i
ZSX ShX sh' i
ZSX = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqType (SKnown n :: SNat n
n@SNat n
SNat :$% ShX sh i
sh) (SKnown m :: SNat n
m@SNat n
SNat :$% ShX sh i
sh')
  | Just n :~: n
Refl <- SNat n -> SNat n -> Maybe (n :~: n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe (a :~: b)
sameNat SNat n
n SNat n
m
  , Just sh :~: sh
Refl <- ShX sh i -> ShX sh i -> Maybe (sh :~: sh)
forall (sh :: [Maybe Nat]) i (sh' :: [Maybe Nat]).
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqType ShX sh i
sh ShX sh i
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqType (SUnknown i
_ :$% ShX sh i
sh) (SUnknown i
_ :$% ShX sh i
sh')
  | Just sh :~: sh
Refl <- ShX sh i -> ShX sh i -> Maybe (sh :~: sh)
forall (sh :: [Maybe Nat]) i (sh' :: [Maybe Nat]).
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqType ShX sh i
sh ShX sh i
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqType ShX sh i
_ ShX sh' i
_ = Maybe (sh :~: sh')
forall a. Maybe a
Nothing

-- | This checks whether all dimensions have the same value. This is more than
-- 'testEquality', and corresponds to @geq@ from @Data.GADT.Compare@ in the
-- @some@ package (except on the penultimate type parameter).
shxEqual :: Eq i => ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqual :: forall i (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Eq i =>
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqual ShX sh i
ZSX ShX sh' i
ZSX = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqual (SKnown n :: SNat n
n@SNat n
SNat :$% ShX sh i
sh) (SKnown m :: SNat n
m@SNat n
SNat :$% ShX sh i
sh')
  | Just n :~: n
Refl <- SNat n -> SNat n -> Maybe (n :~: n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe (a :~: b)
sameNat SNat n
n SNat n
m
  , Just sh :~: sh
Refl <- ShX sh i -> ShX sh i -> Maybe (sh :~: sh)
forall i (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Eq i =>
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqual ShX sh i
sh ShX sh i
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqual (SUnknown i
i :$% ShX sh i
sh) (SUnknown i
j :$% ShX sh i
sh')
  | i
i i -> i -> Bool
forall a. Eq a => a -> a -> Bool
== i
j
  , Just sh :~: sh
Refl <- ShX sh i -> ShX sh i -> Maybe (sh :~: sh)
forall i (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Eq i =>
ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
shxEqual ShX sh i
sh ShX sh i
sh'
  = (sh :~: sh') -> Maybe (sh :~: sh')
forall a. a -> Maybe a
Just sh :~: sh
sh :~: sh'
forall {k} (a :: k). a :~: a
Refl
shxEqual ShX sh i
_ ShX sh' i
_ = Maybe (sh :~: sh')
forall a. Maybe a
Nothing

shxLength :: ShX sh i -> Int
shxLength :: forall (sh :: [Maybe Nat]) i. ShX sh i -> Int
shxLength (ShX ListX sh (SMayNat i SNat)
l) = ListX sh (SMayNat i SNat) -> Int
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> Int
listxLength ListX sh (SMayNat i SNat)
l

shxRank :: ShX sh i -> SNat (Rank sh)
shxRank :: forall (sh :: [Maybe Nat]) i. ShX sh i -> SNat (Rank sh)
shxRank (ShX ListX sh (SMayNat i SNat)
l) = ListX sh (SMayNat i SNat) -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank ListX sh (SMayNat i SNat)
l

-- | The number of elements in an array described by this shape.
shxSize :: IShX sh -> Int
shxSize :: forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize ShX sh Int
ZSX = Int
1
shxSize (SMayNat Int SNat n
n :$% ShX sh Int
sh) = SMayNat Int SNat n -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* ShX sh Int -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize ShX sh Int
sh

shxFromList :: StaticShX sh -> [Int] -> IShX sh
shxFromList :: forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
shxFromList StaticShX sh
topssh [Int]
topl = StaticShX sh -> [Int] -> IShX sh
forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
go StaticShX sh
topssh [Int]
topl
  where
    go :: StaticShX sh' -> [Int] -> IShX sh'
    go :: forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
go StaticShX sh'
ZKX [] = ShX sh' Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
    go (SKnown SNat n
sn :!% StaticShX sh
sh) (Int
i : [Int]
is)
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn = SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
sn SMayNat Int SNat ('Just n) -> ShX sh Int -> ShX sh' Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX sh -> [Int] -> ShX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
go StaticShX sh
sh [Int]
is
      | Bool
otherwise = String -> ShX sh' Int
forall a. HasCallStack => String -> a
error (String -> ShX sh' Int) -> String -> ShX sh' Int
forall a b. (a -> b) -> a -> b
$ String
"shxFromList: Value does not match typing (type says "
                              String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", list contains " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
    go (SUnknown () :!% StaticShX sh
sh) (Int
i : [Int]
is) = Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown Int
i SMayNat Int SNat 'Nothing -> ShX sh Int -> ShX sh' Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX sh -> [Int] -> ShX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
go StaticShX sh
sh [Int]
is
    go StaticShX sh'
_ [Int]
_ = String -> ShX sh' Int
forall a. HasCallStack => String -> a
error (String -> ShX sh' Int) -> String -> ShX sh' Int
forall a b. (a -> b) -> a -> b
$ String
"shxFromList: Mismatched list length (type says "
                       String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (StaticShX sh -> Int
forall (sh :: [Maybe Nat]). StaticShX sh -> Int
ssxLength StaticShX sh
topssh) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", list has length "
                       String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([Int] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Int]
topl) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

shxToList :: IShX sh -> [Int]
shxToList :: forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList ShX sh Int
ZSX = []
shxToList (SMayNat Int SNat n
smn :$% ShX sh Int
sh) = SMayNat Int SNat n -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat n
smn Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: ShX sh Int -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList ShX sh Int
sh

shxFromSSX :: StaticShX (MapJust sh) -> ShX (MapJust sh) i
shxFromSSX :: forall (sh :: [Nat]) i.
StaticShX (MapJust sh) -> ShX (MapJust sh) i
shxFromSSX StaticShX (MapJust sh)
ZKX = ShX '[] i
ShX (MapJust sh) i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxFromSSX (SKnown SNat n
n :!% StaticShX sh
sh :: StaticShX (MapJust sh))
  | sh :~: (n : Tail sh)
Refl <- forall (sh :: [Nat]) (n :: Nat) (sh' :: [Maybe Nat]).
(MapJust sh :~: ('Just n : sh')) -> sh :~: (n : Tail sh)
forall {a} (sh :: [a]) (n :: a) (sh' :: [Maybe a]).
(MapJust sh :~: ('Just n : sh')) -> sh :~: (n : Tail sh)
lemMapJustCons @sh ('Just n : sh) :~: ('Just n : sh)
MapJust sh :~: ('Just n : sh)
forall {k} (a :: k). a :~: a
Refl
  = SNat n -> SMayNat i SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
n SMayNat i SNat ('Just n) -> ShX sh i -> ShX ('Just n : sh) i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX (MapJust (Tail sh)) -> ShX (MapJust (Tail sh)) i
forall (sh :: [Nat]) i.
StaticShX (MapJust sh) -> ShX (MapJust sh) i
shxFromSSX StaticShX sh
StaticShX (MapJust (Tail sh))
sh
shxFromSSX (SUnknown ()
_ :!% StaticShX sh
_) = String -> ShX ('Nothing : sh) i
forall a. HasCallStack => String -> a
error String
"unreachable"

-- | This may fail if @sh@ has @Nothing@s in it.
shxFromSSX2 :: StaticShX sh -> Maybe (ShX sh i)
shxFromSSX2 :: forall (sh :: [Maybe Nat]) i. StaticShX sh -> Maybe (ShX sh i)
shxFromSSX2 StaticShX sh
ZKX = ShX sh i -> Maybe (ShX sh i)
forall a. a -> Maybe a
Just ShX sh i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxFromSSX2 (SKnown SNat n
n :!% StaticShX sh
sh) = (SNat n -> SMayNat i SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
n SMayNat i SNat ('Just n) -> ShX sh i -> ShX sh i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (ShX sh i -> ShX sh i) -> Maybe (ShX sh i) -> Maybe (ShX sh i)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticShX sh -> Maybe (ShX sh i)
forall (sh :: [Maybe Nat]) i. StaticShX sh -> Maybe (ShX sh i)
shxFromSSX2 StaticShX sh
sh
shxFromSSX2 (SUnknown ()
_ :!% StaticShX sh
_) = Maybe (ShX sh i)
forall a. Maybe a
Nothing

shxAppend :: forall sh sh' i. ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend = (ListX sh (SMayNat i SNat)
 -> ListX sh' (SMayNat i SNat)
 -> ListX (sh ++ sh') (SMayNat i SNat))
-> ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
forall a b. Coercible a b => a -> b
coerce (forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type)
       (sh' :: [Maybe Nat]).
ListX sh f -> ListX sh' f -> ListX (sh ++ sh') f
listxAppend @_ @(SMayNat i SNat))

shxHead :: ShX (n : sh) i -> SMayNat i SNat n
shxHead :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX (n : sh) i -> SMayNat i SNat n
shxHead (ShX ListX (n : sh) (SMayNat i SNat)
list) = ListX (n : sh) (SMayNat i SNat) -> SMayNat i SNat n
forall (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (f :: Maybe Nat -> Type).
ListX (mn : sh) f -> f mn
listxHead ListX (n : sh) (SMayNat i SNat)
list

shxTail :: ShX (n : sh) i -> ShX sh i
shxTail :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX (n : sh) i -> ShX sh i
shxTail (ShX ListX (n : sh) (SMayNat i SNat)
list) = ListX sh (SMayNat i SNat) -> ShX sh i
forall (sh :: [Maybe Nat]) i. ListX sh (SMayNat i SNat) -> ShX sh i
ShX (ListX (n : sh) (SMayNat i SNat) -> ListX sh (SMayNat i SNat)
forall (n :: Maybe Nat) (sh :: [Maybe Nat])
       (i :: Maybe Nat -> Type).
ListX (n : sh) i -> ListX sh i
listxTail ListX (n : sh) (SMayNat i SNat)
list)

shxDropSSX :: forall sh sh' i. StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
shxDropSSX :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
shxDropSSX = (ListX sh (SMayNat () SNat)
 -> ListX (sh ++ sh') (SMayNat i SNat)
 -> ListX sh' (SMayNat i SNat))
-> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat i SNat) @(SMayNat () SNat))

shxDropIx :: forall sh sh' i j. IxX sh j -> ShX (sh ++ sh') i -> ShX sh' i
shxDropIx :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i j.
IxX sh j -> ShX (sh ++ sh') i -> ShX sh' i
shxDropIx = (ListX sh (Const j)
 -> ListX (sh ++ sh') (SMayNat i SNat)
 -> ListX sh' (SMayNat i SNat))
-> IxX sh j -> ShX (sh ++ sh') i -> ShX sh' i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat i SNat) @(Const j))

shxDropSh :: forall sh sh' i. ShX sh i -> ShX (sh ++ sh') i -> ShX sh' i
shxDropSh :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> ShX (sh ++ sh') i -> ShX sh' i
shxDropSh = (ListX sh (SMayNat i SNat)
 -> ListX (sh ++ sh') (SMayNat i SNat)
 -> ListX sh' (SMayNat i SNat))
-> ShX sh i -> ShX (sh ++ sh') i -> ShX sh' i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat i SNat) @(SMayNat i SNat))

shxInit :: forall n sh i. ShX (n : sh) i -> ShX (Init (n : sh)) i
shxInit :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX (n : sh) i -> ShX (Init (n : sh)) i
shxInit = (ListX (n : sh) (SMayNat i SNat)
 -> ListX (Init (n : sh)) (SMayNat i SNat))
-> ShX (n : sh) i -> ShX (Init (n : sh)) i
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit @(SMayNat i SNat))

shxLast :: forall n sh i. ShX (n : sh) i -> SMayNat i SNat (Last (n : sh))
shxLast :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX (n : sh) i -> SMayNat i SNat (Last (n : sh))
shxLast = (ListX (n : sh) (SMayNat i SNat) -> SMayNat i SNat (Last (n : sh)))
-> ShX (n : sh) i -> SMayNat i SNat (Last (n : sh))
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> f (Last (n : sh))
listxLast @(SMayNat i SNat))

shxTakeSSX :: forall sh sh' i proxy. proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
shxTakeSSX :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> Type).
proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
shxTakeSSX proxy sh'
_ StaticShX sh
ZKX ShX (sh ++ sh') i
_ = ShX sh i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxTakeSSX proxy sh'
p (SMayNat () SNat n
_ :!% StaticShX sh
ssh1) (SMayNat i SNat n
n :$% ShX sh i
sh) = SMayNat i SNat n
n SMayNat i SNat n -> ShX sh i -> ShX sh i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> Type).
proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
shxTakeSSX proxy sh'
p StaticShX sh
ssh1 ShX sh i
ShX (sh ++ sh') i
sh

shxZipWith :: (forall n. SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
           -> ShX sh i -> ShX sh j -> ShX sh k
shxZipWith :: forall i j k (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat).
 SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
-> ShX sh i -> ShX sh j -> ShX sh k
shxZipWith forall (n :: Maybe Nat).
SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
_ ShX sh i
ZSX ShX sh j
ZSX = ShX sh k
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxZipWith forall (n :: Maybe Nat).
SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
f (SMayNat i SNat n
i :$% ShX sh i
is) (SMayNat j SNat n
j :$% ShX sh j
js) = SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
forall (n :: Maybe Nat).
SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
f SMayNat i SNat n
i SMayNat j SNat n
SMayNat j SNat n
j SMayNat k SNat n -> ShX sh k -> ShX sh k
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% (forall (n :: Maybe Nat).
 SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
-> ShX sh i -> ShX sh j -> ShX sh k
forall i j k (sh :: [Maybe Nat]).
(forall (n :: Maybe Nat).
 SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n)
-> ShX sh i -> ShX sh j -> ShX sh k
shxZipWith SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
forall (n :: Maybe Nat).
SMayNat i SNat n -> SMayNat j SNat n -> SMayNat k SNat n
f ShX sh i
is ShX sh j
ShX sh j
js

-- This is a weird operation, so it has a long name
shxCompleteZeros :: StaticShX sh -> IShX sh
shxCompleteZeros :: forall (sh :: [Maybe Nat]). StaticShX sh -> IShX sh
shxCompleteZeros StaticShX sh
ZKX = ShX sh Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxCompleteZeros (SUnknown () :!% StaticShX sh
ssh) = Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown Int
0 SMayNat Int SNat 'Nothing -> ShX sh Int -> ShX sh Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX sh -> ShX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> IShX sh
shxCompleteZeros StaticShX sh
ssh
shxCompleteZeros (SKnown SNat n
n :!% StaticShX sh
ssh) = SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
n SMayNat Int SNat ('Just n) -> ShX sh Int -> ShX sh Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX sh -> ShX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> IShX sh
shxCompleteZeros StaticShX sh
ssh

shxSplitApp :: proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp :: forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp proxy sh'
_ StaticShX sh
ZKX ShX (sh ++ sh') i
idx = (ShX sh i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX, ShX sh' i
ShX (sh ++ sh') i
idx)
shxSplitApp proxy sh'
p (SMayNat () SNat n
_ :!% StaticShX sh
ssh) (SMayNat i SNat n
i :$% ShX sh i
idx) = (ShX sh i -> ShX sh i)
-> (ShX sh i, ShX sh' i) -> (ShX sh i, ShX sh' i)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: Type -> Type -> Type) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (SMayNat i SNat n
i SMayNat i SNat n -> ShX sh i -> ShX sh i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp proxy sh'
p StaticShX sh
ssh ShX sh i
ShX (sh ++ sh') i
idx)

shxEnum :: IShX sh -> [IIxX sh]
shxEnum :: forall (sh :: [Maybe Nat]). IShX sh -> [IIxX sh]
shxEnum = \IShX sh
sh -> IShX sh -> (IIxX sh -> IIxX sh) -> [IIxX sh] -> [IIxX sh]
forall (sh :: [Maybe Nat]) a.
IShX sh -> (IIxX sh -> a) -> [a] -> [a]
go IShX sh
sh IIxX sh -> IIxX sh
forall a. a -> a
id []
  where
    go :: IShX sh -> (IIxX sh -> a) -> [a] -> [a]
    go :: forall (sh :: [Maybe Nat]) a.
IShX sh -> (IIxX sh -> a) -> [a] -> [a]
go ShX sh Int
ZSX IIxX sh -> a
f = (IIxX sh -> a
f IIxX sh
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX a -> [a] -> [a]
forall a. a -> [a] -> [a]
:)
    go (SMayNat Int SNat n
n :$% ShX sh Int
sh) IIxX sh -> a
f = (([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a])
-> ([a] -> [a]) -> [[a] -> [a]] -> [a] -> [a]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) [a] -> [a]
forall a. a -> a
id [ShX sh Int -> (IIxX sh -> a) -> [a] -> [a]
forall (sh :: [Maybe Nat]) a.
IShX sh -> (IIxX sh -> a) -> [a] -> [a]
go ShX sh Int
sh (IIxX sh -> a
f (IIxX sh -> a) -> (IIxX sh -> IIxX sh) -> IIxX sh -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
i Int -> IIxX sh -> IIxX sh
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.%)) | Int
i <- [Int
0 .. SMayNat Int SNat n -> Int
forall (n :: Maybe Nat). SMayNat Int SNat n -> Int
fromSMayNat' SMayNat Int SNat n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast :: forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh'
ZKX ShX sh Int
ZSX = IShX sh' -> Maybe (IShX sh')
forall a. a -> Maybe a
Just IShX sh'
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxCast (SKnown SNat n
m    :!% StaticShX sh
ssh) (SKnown SNat n
n   :$% ShX sh Int
sh) | Just n :~: n
Refl <- SNat n -> SNat n -> Maybe (n :~: n)
forall (a :: Nat) (b :: Nat). SNat a -> SNat b -> Maybe (a :~: b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality SNat n
n SNat n
m = (SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
n SMayNat Int SNat ('Just n) -> ShX sh Int -> IShX sh'
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (ShX sh Int -> IShX sh') -> Maybe (ShX sh Int) -> Maybe (IShX sh')
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticShX sh -> ShX sh Int -> Maybe (ShX sh Int)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh
ssh ShX sh Int
sh
shxCast (SKnown SNat n
m    :!% StaticShX sh
ssh) (SUnknown Int
n :$% ShX sh Int
sh) | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
m              = (SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
m SMayNat Int SNat ('Just n) -> ShX sh Int -> IShX sh'
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (ShX sh Int -> IShX sh') -> Maybe (ShX sh Int) -> Maybe (IShX sh')
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticShX sh -> ShX sh Int -> Maybe (ShX sh Int)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh
ssh ShX sh Int
sh
shxCast (SUnknown () :!% StaticShX sh
ssh) (SKnown SNat n
n   :$% ShX sh Int
sh)                                 = (Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown (SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
n) SMayNat Int SNat 'Nothing -> ShX sh Int -> IShX sh'
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (ShX sh Int -> IShX sh') -> Maybe (ShX sh Int) -> Maybe (IShX sh')
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticShX sh -> ShX sh Int -> Maybe (ShX sh Int)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh
ssh ShX sh Int
sh
shxCast (SUnknown () :!% StaticShX sh
ssh) (SUnknown Int
n :$% ShX sh Int
sh)                                 = (Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown Int
n SMayNat Int SNat 'Nothing -> ShX sh Int -> IShX sh'
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$%) (ShX sh Int -> IShX sh') -> Maybe (ShX sh Int) -> Maybe (IShX sh')
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticShX sh -> ShX sh Int -> Maybe (ShX sh Int)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh
ssh ShX sh Int
sh
shxCast StaticShX sh'
_ ShX sh Int
_ = Maybe (IShX sh')
forall a. Maybe a
Nothing

-- | Partial version of 'shxCast'.
shxCast' :: StaticShX sh' -> IShX sh -> IShX sh'
shxCast' :: forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> IShX sh'
shxCast' StaticShX sh'
ssh IShX sh
sh = case StaticShX sh' -> IShX sh -> Maybe (IShX sh')
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]).
StaticShX sh' -> IShX sh -> Maybe (IShX sh')
shxCast StaticShX sh'
ssh IShX sh
sh of
  Just IShX sh'
sh' -> IShX sh'
sh'
  Maybe (IShX sh')
Nothing -> String -> IShX sh'
forall a. HasCallStack => String -> a
error (String -> IShX sh') -> String -> IShX sh'
forall a b. (a -> b) -> a -> b
$ String
"shxCast': Mismatch: (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ IShX sh -> String
forall a. Show a => a -> String
show IShX sh
sh String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") does not match (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticShX sh' -> String
forall a. Show a => a -> String
show StaticShX sh'
ssh String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"


-- * Static mixed shapes

-- | The part of a shape that is statically known. (A newtype over 'ListX'.)
type StaticShX :: [Maybe Nat] -> Type
newtype StaticShX sh = StaticShX (ListX sh (SMayNat () SNat))
  deriving (StaticShX sh -> StaticShX sh -> Bool
(StaticShX sh -> StaticShX sh -> Bool)
-> (StaticShX sh -> StaticShX sh -> Bool) -> Eq (StaticShX sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
== :: StaticShX sh -> StaticShX sh -> Bool
$c/= :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
/= :: StaticShX sh -> StaticShX sh -> Bool
Eq, Eq (StaticShX sh)
Eq (StaticShX sh) =>
(StaticShX sh -> StaticShX sh -> Ordering)
-> (StaticShX sh -> StaticShX sh -> Bool)
-> (StaticShX sh -> StaticShX sh -> Bool)
-> (StaticShX sh -> StaticShX sh -> Bool)
-> (StaticShX sh -> StaticShX sh -> Bool)
-> (StaticShX sh -> StaticShX sh -> StaticShX sh)
-> (StaticShX sh -> StaticShX sh -> StaticShX sh)
-> Ord (StaticShX sh)
StaticShX sh -> StaticShX sh -> Bool
StaticShX sh -> StaticShX sh -> Ordering
StaticShX sh -> StaticShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]). Eq (StaticShX sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
forall (sh :: [Maybe Nat]).
StaticShX sh -> StaticShX sh -> Ordering
forall (sh :: [Maybe Nat]).
StaticShX sh -> StaticShX sh -> StaticShX sh
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (sh :: [Maybe Nat]).
StaticShX sh -> StaticShX sh -> Ordering
compare :: StaticShX sh -> StaticShX sh -> Ordering
$c< :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
< :: StaticShX sh -> StaticShX sh -> Bool
$c<= :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
<= :: StaticShX sh -> StaticShX sh -> Bool
$c> :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
> :: StaticShX sh -> StaticShX sh -> Bool
$c>= :: forall (sh :: [Maybe Nat]). StaticShX sh -> StaticShX sh -> Bool
>= :: StaticShX sh -> StaticShX sh -> Bool
$cmax :: forall (sh :: [Maybe Nat]).
StaticShX sh -> StaticShX sh -> StaticShX sh
max :: StaticShX sh -> StaticShX sh -> StaticShX sh
$cmin :: forall (sh :: [Maybe Nat]).
StaticShX sh -> StaticShX sh -> StaticShX sh
min :: StaticShX sh -> StaticShX sh -> StaticShX sh
Ord)

pattern ZKX :: forall sh. () => sh ~ '[] => StaticShX sh
pattern $mZKX :: forall {r} {sh :: [Maybe Nat]}.
StaticShX sh -> ((sh ~ '[]) => r) -> ((# #) -> r) -> r
$bZKX :: forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX = StaticShX ZX

pattern (:!%)
  :: forall {sh1}.
     forall n sh. (n : sh ~ sh1)
  => SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
pattern i $m:!% :: forall {r} {sh1 :: [Maybe Nat]}.
StaticShX sh1
-> (forall {n :: Maybe Nat} {sh :: [Maybe Nat]}.
    ((n : sh) ~ sh1) =>
    SMayNat () SNat n -> StaticShX sh -> r)
-> ((# #) -> r)
-> r
$b:!% :: forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% shl <- StaticShX (listxUncons -> Just (UnconsListXRes (StaticShX -> shl) i))
  where SMayNat () SNat n
i :!% StaticShX ListX sh (SMayNat () SNat)
shl = ListX sh1 (SMayNat () SNat) -> StaticShX sh1
forall (sh :: [Maybe Nat]).
ListX sh (SMayNat () SNat) -> StaticShX sh
StaticShX (SMayNat () SNat n
i SMayNat () SNat n
-> ListX sh (SMayNat () SNat) -> ListX (n : sh) (SMayNat () SNat)
forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
f n -> ListX sh f -> ListX (n : sh) f
::% ListX sh (SMayNat () SNat)
shl)
infixr 3 :!%

{-# COMPLETE ZKX, (:!%) #-}

#ifdef OXAR_DEFAULT_SHOW_INSTANCES
deriving instance Show (StaticShX sh)
#else
instance Show (StaticShX sh) where
  showsPrec :: Int -> StaticShX sh -> ShowS
showsPrec Int
_ (StaticShX ListX sh (SMayNat () SNat)
l) = (forall (n :: Maybe Nat). SMayNat () SNat n -> ShowS)
-> ListX sh (SMayNat () SNat) -> ShowS
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
(forall (n :: Maybe Nat). f n -> ShowS) -> ListX sh f -> ShowS
listxShow (((n ~ 'Nothing) => () -> ShowS)
-> (forall (m :: Nat). (n ~ 'Just m) => SNat m -> ShowS)
-> SMayNat () SNat n
-> ShowS
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat (n ~ 'Nothing) => () -> ShowS
() -> ShowS
forall a. Show a => a -> ShowS
shows (Integer -> ShowS
forall a. Show a => a -> ShowS
shows (Integer -> ShowS) -> (SNat m -> Integer) -> SNat m -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat m -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat)) ListX sh (SMayNat () SNat)
l
#endif

instance NFData (StaticShX sh) where
  rnf :: StaticShX sh -> ()
rnf (StaticShX ListX sh (SMayNat () SNat)
ZX) = ()
  rnf (StaticShX (SUnknown () ::% ListX sh (SMayNat () SNat)
l)) = StaticShX sh -> ()
forall a. NFData a => a -> ()
rnf (ListX sh (SMayNat () SNat) -> StaticShX sh
forall (sh :: [Maybe Nat]).
ListX sh (SMayNat () SNat) -> StaticShX sh
StaticShX ListX sh (SMayNat () SNat)
l)
  rnf (StaticShX (SKnown SNat n
SNat ::% ListX sh (SMayNat () SNat)
l)) = StaticShX sh -> ()
forall a. NFData a => a -> ()
rnf (ListX sh (SMayNat () SNat) -> StaticShX sh
forall (sh :: [Maybe Nat]).
ListX sh (SMayNat () SNat) -> StaticShX sh
StaticShX ListX sh (SMayNat () SNat)
l)

instance TestEquality StaticShX where
  testEquality :: forall (a :: [Maybe Nat]) (b :: [Maybe Nat]).
StaticShX a -> StaticShX b -> Maybe (a :~: b)
testEquality (StaticShX ListX a (SMayNat () SNat)
l1) (StaticShX ListX b (SMayNat () SNat)
l2) = ListX a (SMayNat () SNat)
-> ListX b (SMayNat () SNat) -> Maybe (a :~: b)
forall (f :: Maybe Nat -> Type) (sh :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
TestEquality f =>
ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
listxEqType ListX a (SMayNat () SNat)
l1 ListX b (SMayNat () SNat)
l2

ssxLength :: StaticShX sh -> Int
ssxLength :: forall (sh :: [Maybe Nat]). StaticShX sh -> Int
ssxLength (StaticShX ListX sh (SMayNat () SNat)
l) = ListX sh (SMayNat () SNat) -> Int
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> Int
listxLength ListX sh (SMayNat () SNat)
l

ssxRank :: StaticShX sh -> SNat (Rank sh)
ssxRank :: forall (sh :: [Maybe Nat]). StaticShX sh -> SNat (Rank sh)
ssxRank (StaticShX ListX sh (SMayNat () SNat)
l) = ListX sh (SMayNat () SNat) -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank ListX sh (SMayNat () SNat)
l

-- | @ssxEqType = 'testEquality'@. Provided for consistency.
ssxEqType :: StaticShX sh -> StaticShX sh' -> Maybe (sh :~: sh')
ssxEqType :: forall (a :: [Maybe Nat]) (b :: [Maybe Nat]).
StaticShX a -> StaticShX b -> Maybe (a :~: b)
ssxEqType = StaticShX sh -> StaticShX sh' -> Maybe (sh :~: sh')
forall (a :: [Maybe Nat]) (b :: [Maybe Nat]).
StaticShX a -> StaticShX b -> Maybe (a :~: b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality

ssxAppend :: StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ZKX StaticShX sh'
sh' = StaticShX sh'
StaticShX (sh ++ sh')
sh'
ssxAppend (SMayNat () SNat n
n :!% StaticShX sh
sh) StaticShX sh'
sh' = SMayNat () SNat n
n SMayNat () SNat n
-> StaticShX (sh ++ sh') -> StaticShX (n : (sh ++ sh'))
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
sh StaticShX sh'
sh'

ssxHead :: StaticShX (n : sh) -> SMayNat () SNat n
ssxHead :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
StaticShX (n : sh) -> SMayNat () SNat n
ssxHead (StaticShX ListX (n : sh) (SMayNat () SNat)
list) = ListX (n : sh) (SMayNat () SNat) -> SMayNat () SNat n
forall (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (f :: Maybe Nat -> Type).
ListX (mn : sh) f -> f mn
listxHead ListX (n : sh) (SMayNat () SNat)
list

ssxTail :: StaticShX (n : sh) -> StaticShX sh
ssxTail :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
StaticShX (n : sh) -> StaticShX sh
ssxTail (SMayNat () SNat n
_ :!% StaticShX sh
ssh) = StaticShX sh
StaticShX sh
ssh

ssxDropSSX :: forall sh sh'. StaticShX sh -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropSSX :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropSSX = (ListX sh (SMayNat () SNat)
 -> ListX (sh ++ sh') (SMayNat () SNat)
 -> ListX sh' (SMayNat () SNat))
-> StaticShX sh -> StaticShX (sh ++ sh') -> StaticShX sh'
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat () SNat) @(SMayNat () SNat))

ssxDropIx :: forall sh sh' i. IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropIx :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropIx = (ListX sh (Const i)
 -> ListX (sh ++ sh') (SMayNat () SNat)
 -> ListX sh' (SMayNat () SNat))
-> IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat () SNat) @(Const i))

ssxDropSh :: forall sh sh' i. ShX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropSh :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
ssxDropSh = (ListX sh (SMayNat i SNat)
 -> ListX (sh ++ sh') (SMayNat () SNat)
 -> ListX sh' (SMayNat () SNat))
-> ShX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (g :: Maybe Nat -> Type)
       (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
ListX sh g -> ListX (sh ++ sh') f -> ListX sh' f
listxDrop @(SMayNat () SNat) @(SMayNat i SNat))

ssxInit :: forall n sh. StaticShX (n : sh) -> StaticShX (Init (n : sh))
ssxInit :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
StaticShX (n : sh) -> StaticShX (Init (n : sh))
ssxInit = (ListX (n : sh) (SMayNat () SNat)
 -> ListX (Init (n : sh)) (SMayNat () SNat))
-> StaticShX (n : sh) -> StaticShX (Init (n : sh))
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> ListX (Init (n : sh)) f
listxInit @(SMayNat () SNat))

ssxLast :: forall n sh. StaticShX (n : sh) -> SMayNat () SNat (Last (n : sh))
ssxLast :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
StaticShX (n : sh) -> SMayNat () SNat (Last (n : sh))
ssxLast = (ListX (n : sh) (SMayNat () SNat)
 -> SMayNat () SNat (Last (n : sh)))
-> StaticShX (n : sh) -> SMayNat () SNat (Last (n : sh))
forall a b. Coercible a b => a -> b
coerce (forall (f :: Maybe Nat -> Type) (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
ListX (n : sh) f -> f (Last (n : sh))
listxLast @(SMayNat () SNat))

ssxReplicate :: SNat n -> StaticShX (Replicate n Nothing)
ssxReplicate :: forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxReplicate SNat n
SZ = StaticShX '[]
StaticShX (Replicate n 'Nothing)
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX
ssxReplicate (SS (SNat n
n :: SNat n'))
  | ('Nothing : Replicate n 'Nothing) :~: Replicate (n + 1) 'Nothing
Refl <- forall (a2 :: Maybe Nat) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
forall {a1} (a2 :: a1) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
lemReplicateSucc @(Nothing @Nat) @n'
  = () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing
-> StaticShX (Replicate n 'Nothing)
-> StaticShX ('Nothing : Replicate n 'Nothing)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% SNat n -> StaticShX (Replicate n 'Nothing)
forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxReplicate SNat n
n

ssxIotaFrom :: StaticShX sh -> Int -> [Int]
ssxIotaFrom :: forall (sh :: [Maybe Nat]). StaticShX sh -> Int -> [Int]
ssxIotaFrom StaticShX sh
ZKX Int
_ = []
ssxIotaFrom (SMayNat () SNat n
_ :!% StaticShX sh
ssh) Int
i = Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: StaticShX sh -> Int -> [Int]
forall (sh :: [Maybe Nat]). StaticShX sh -> Int -> [Int]
ssxIotaFrom StaticShX sh
ssh (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

ssxFromShX :: ShX sh i -> StaticShX sh
ssxFromShX :: forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX ShX sh i
ZSX = StaticShX sh
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX
ssxFromShX (SMayNat i SNat n
n :$% ShX sh i
sh) = ((n ~ 'Nothing) => i -> SMayNat () SNat n)
-> (forall (m :: Nat).
    (n ~ 'Just m) =>
    SNat m -> SMayNat () SNat n)
-> SMayNat i SNat n
-> SMayNat () SNat n
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat (\i
_ -> () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown ()) SNat m -> SMayNat () SNat n
SNat m -> SMayNat () SNat ('Just m)
forall (m :: Nat). (n ~ 'Just m) => SNat m -> SMayNat () SNat n
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SMayNat i SNat n
n SMayNat () SNat n -> StaticShX sh -> StaticShX sh
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% ShX sh i -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX ShX sh i
sh

ssxFromSNat :: SNat n -> StaticShX (Replicate n Nothing)
ssxFromSNat :: forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxFromSNat SNat n
SZ = StaticShX '[]
StaticShX (Replicate n 'Nothing)
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX
ssxFromSNat (SS (SNat n
n :: SNat nm1)) | ('Nothing : Replicate n 'Nothing) :~: Replicate (n + 1) 'Nothing
Refl <- forall (a2 :: Maybe Nat) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
forall {a1} (a2 :: a1) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
lemReplicateSucc @(Nothing @Nat) @nm1 = () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing
-> StaticShX (Replicate n 'Nothing)
-> StaticShX ('Nothing : Replicate n 'Nothing)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% SNat n -> StaticShX (Replicate n 'Nothing)
forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxFromSNat SNat n
n


-- | Evidence for the static part of a shape. This pops up only when you are
-- polymorphic in the element type of an array.
type KnownShX :: [Maybe Nat] -> Constraint
class KnownShX sh where knownShX :: StaticShX sh
instance KnownShX '[] where knownShX :: StaticShX '[]
knownShX = StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX
instance (KnownNat n, KnownShX sh) => KnownShX (Just n : sh) where knownShX :: StaticShX ('Just n : sh)
knownShX = SNat n -> SMayNat () SNat ('Just n)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat n
forall (n :: Nat). KnownNat n => SNat n
natSing SMayNat () SNat ('Just n)
-> StaticShX sh -> StaticShX ('Just n : sh)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX sh
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX
instance KnownShX sh => KnownShX (Nothing : sh) where knownShX :: StaticShX ('Nothing : sh)
knownShX = () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing
-> StaticShX sh -> StaticShX ('Nothing : sh)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX sh
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX

withKnownShX :: forall sh r. StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX :: forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX = forall (cls :: Constraint) meth r.
WithDict cls meth =>
meth -> (cls => r) -> r
withDict @(KnownShX sh)


-- * Flattening

type Flatten sh = Flatten' 1 sh

type family Flatten' acc sh where
  Flatten' acc '[] = Just acc
  Flatten' acc (Nothing : sh) = Nothing
  Flatten' acc (Just n : sh) = Flatten' (acc * n) sh

-- This function is currently unused
ssxFlatten :: StaticShX sh -> SMayNat () SNat (Flatten sh)
ssxFlatten :: forall (sh :: [Maybe Nat]).
StaticShX sh -> SMayNat () SNat (Flatten sh)
ssxFlatten = SNat 1 -> StaticShX sh -> SMayNat () SNat (Flatten' 1 sh)
forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> StaticShX sh -> SMayNat () SNat (Flatten' acc sh)
go (forall (n :: Nat). KnownNat n => SNat n
SNat @1)
  where
    go :: SNat acc -> StaticShX sh -> SMayNat () SNat (Flatten' acc sh)
    go :: forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> StaticShX sh -> SMayNat () SNat (Flatten' acc sh)
go SNat acc
acc StaticShX sh
ZKX = SNat acc -> SMayNat () SNat ('Just acc)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat acc
acc
    go SNat acc
_ (SUnknown () :!% StaticShX sh
_) = () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown ()
    go SNat acc
acc (SKnown SNat n
sn :!% StaticShX sh
sh) = SNat (acc * n)
-> StaticShX sh -> SMayNat () SNat (Flatten' (acc * n) sh)
forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> StaticShX sh -> SMayNat () SNat (Flatten' acc sh)
go (SNat acc -> SNat n -> SNat (acc * n)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n * m)
snatMul SNat acc
acc SNat n
sn) StaticShX sh
sh

shxFlatten :: IShX sh -> SMayNat Int SNat (Flatten sh)
shxFlatten :: forall (sh :: [Maybe Nat]).
IShX sh -> SMayNat Int SNat (Flatten sh)
shxFlatten = SNat 1 -> ShX sh Int -> SMayNat Int SNat (Flatten' 1 sh)
forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> IShX sh -> SMayNat Int SNat (Flatten' acc sh)
go (forall (n :: Nat). KnownNat n => SNat n
SNat @1)
  where
    go :: SNat acc -> IShX sh -> SMayNat Int SNat (Flatten' acc sh)
    go :: forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> IShX sh -> SMayNat Int SNat (Flatten' acc sh)
go SNat acc
acc ShX sh Int
ZSX = SNat acc -> SMayNat Int SNat ('Just acc)
forall {k} (f :: k -> Type) (n :: k) i.
f n -> SMayNat i f ('Just n)
SKnown SNat acc
acc
    go SNat acc
acc (SUnknown Int
n :$% ShX sh Int
sh) = Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown (Int -> ShX sh Int -> Int
forall (sh :: [Maybe Nat]). Int -> IShX sh -> Int
goUnknown (SNat acc -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat acc
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) ShX sh Int
sh)
    go SNat acc
acc (SKnown SNat n
sn :$% ShX sh Int
sh) = SNat (acc * n)
-> ShX sh Int -> SMayNat Int SNat (Flatten' (acc * n) sh)
forall (acc :: Nat) (sh :: [Maybe Nat]).
SNat acc -> IShX sh -> SMayNat Int SNat (Flatten' acc sh)
go (SNat acc -> SNat n -> SNat (acc * n)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n * m)
snatMul SNat acc
acc SNat n
sn) ShX sh Int
sh

    goUnknown :: Int -> IShX sh -> Int
    goUnknown :: forall (sh :: [Maybe Nat]). Int -> IShX sh -> Int
goUnknown Int
acc ShX sh Int
ZSX = Int
acc
    goUnknown Int
acc (SUnknown Int
n :$% ShX sh Int
sh) = Int -> ShX sh Int -> Int
forall (sh :: [Maybe Nat]). Int -> IShX sh -> Int
goUnknown (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) ShX sh Int
sh
    goUnknown Int
acc (SKnown SNat n
sn :$% ShX sh Int
sh) = Int -> ShX sh Int -> Int
forall (sh :: [Maybe Nat]). Int -> IShX sh -> Int
goUnknown (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
* SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn) ShX sh Int
sh


-- | Very untyped: only length is checked (at runtime).
instance KnownShX sh => IsList (ListX sh (Const i)) where
  type Item (ListX sh (Const i)) = i
  fromList :: [Item (ListX sh (Const i))] -> ListX sh (Const i)
fromList = StaticShX sh -> [i] -> ListX sh (Const i)
forall (sh :: [Maybe Nat]) i.
StaticShX sh -> [i] -> ListX sh (Const i)
listxFromList (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh)
  toList :: ListX sh (Const i) -> [Item (ListX sh (Const i))]
toList = ListX sh (Const i) -> [i]
ListX sh (Const i) -> [Item (ListX sh (Const i))]
forall (sh' :: [Maybe Nat]) i. ListX sh' (Const i) -> [i]
listxToList

-- | Very untyped: only length is checked (at runtime), index bounds are __not checked__.
instance KnownShX sh => IsList (IxX sh i) where
  type Item (IxX sh i) = i
  fromList :: [Item (IxX sh i)] -> IxX sh i
fromList = ListX sh (Const i) -> IxX sh i
forall (sh :: [Maybe Nat]) i. ListX sh (Const i) -> IxX sh i
IxX (ListX sh (Const i) -> IxX sh i)
-> ([i] -> ListX sh (Const i)) -> [i] -> IxX sh i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [i] -> ListX sh (Const i)
[Item (ListX sh (Const i))] -> ListX sh (Const i)
forall l. IsList l => [Item l] -> l
IsList.fromList
  toList :: IxX sh i -> [Item (IxX sh i)]
toList = IxX sh i -> [i]
IxX sh i -> [Item (IxX sh i)]
forall a. IxX sh a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList

-- | Untyped: length and known dimensions are checked (at runtime).
instance KnownShX sh => IsList (ShX sh Int) where
  type Item (ShX sh Int) = Int
  fromList :: [Item (ShX sh Int)] -> ShX sh Int
fromList = StaticShX sh -> [Int] -> ShX sh Int
forall (sh :: [Maybe Nat]). StaticShX sh -> [Int] -> IShX sh
shxFromList (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh)
  toList :: ShX sh Int -> [Item (ShX sh Int)]
toList = ShX sh Int -> [Int]
ShX sh Int -> [Item (ShX sh Int)]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList