1 files changed, 628 insertions, 0 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
new file mode 100644
index 0000000..a01e0f3
--- /dev/null
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -0,0 +1,628 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE NoStarIsType #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE StrictData #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE ViewPatterns #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
+module Data.Array.Nested.Mixed.Shape (
+  module Data.Array.Nested.Mixed.Shape,
+  Rank,
+) where
+
+import Control.DeepSeq (NFData(..))
+import Control.Exception (assert)
+import Data.Bifunctor (first)
+import Data.Coerce
+import Data.Foldable qualified as Foldable
+import Data.Kind (Constraint, Type)
+import Data.Monoid (Sum(..))
+import Data.Proxy
+import Data.Type.Equality
+import GHC.Exts (Int(..), Int#, build, quotRemInt#, withDict)
+import GHC.IsList (IsList)
+import GHC.IsList qualified as IsList
+import GHC.TypeLits
+#if !MIN_VERSION_GLASGOW_HASKELL(9,8,0,0)
+import GHC.TypeLits.Orphans ()
+#endif
+
+import Data.Array.Nested.Mixed.ListX
+import Data.Array.Nested.Types
+
+
+-- * 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 i)
+  deriving (Eq, Ord, NFData, Functor, Foldable)
+
+pattern ZIX :: forall sh i. () => sh ~ '[] => IxX sh i
+pattern ZIX = IxX ZX
+
+pattern (:.%)
+  :: forall {sh1} {i}.
+     forall n sh. (n : sh ~ sh1)
+  => i -> IxX sh i -> IxX sh1 i
+pattern i :.% l <- IxX (i ::% (IxX -> l))
+  where i :.% IxX l = IxX (i ::% l)
+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 _ (IxX l) = listxShow shows l
+#endif
+
+{-# INLINE ixxFromList #-}
+ixxFromList :: StaticShX sh -> [i] -> IxX sh i
+ixxFromList sh l = assert (ssxLength sh == length l) $ IsList.fromList l
+
+ixxRank :: IxX sh i -> SNat (Rank sh)
+ixxRank ZIX = SNat
+ixxRank (_ :.% l) | SNat <- ixxRank l = SNat
+
+ixxZero :: StaticShX sh -> IIxX sh
+ixxZero ZKX = ZIX
+ixxZero (_ :!% ssh) = 0 :.% ixxZero ssh
+
+ixxZero' :: IShX sh -> IIxX sh
+ixxZero' ZSX = ZIX
+ixxZero' (_ :$% sh) = 0 :.% ixxZero' sh
+
+ixxHead :: IxX (mn ': sh) i -> i
+ixxHead (i :.% _) = i
+
+ixxTail :: IxX (n : sh) i -> IxX sh i
+ixxTail (_ :.% sh) = sh
+
+ixxAppend :: forall sh sh' i. IxX sh i -> IxX sh' i -> IxX (sh ++ sh') i
+ixxAppend (IxX l1) (IxX l2) = IxX $ lazilyConcat (++) l1 l2
+
+ixxDrop :: forall i j sh sh'. IxX sh j -> IxX (sh ++ sh') i -> IxX sh' i
+ixxDrop ZIX long = long
+ixxDrop (_ :.% short) long = case long of _ :.% long' -> ixxDrop short long'
+
+ixxInit :: forall i n sh. IxX (n : sh) i -> IxX (Init (n : sh)) i
+ixxInit (i :.% sh@(_ :.% _)) = i :.% ixxInit sh
+ixxInit (_ :.% ZIX) = ZIX
+
+ixxLast :: forall i n sh. IxX (n : sh) i -> i
+ixxLast (_ :.% sh@(_ :.% _)) = ixxLast sh
+ixxLast (x :.% ZIX) = x
+
+ixxZip :: IxX sh i -> IxX sh j -> IxX sh (i, j)
+ixxZip ZIX ZIX = ZIX
+ixxZip (i :.% is) (j :.% js) = (i, j) :.% ixxZip is js
+
+{-# INLINE ixxZipWith #-}
+ixxZipWith :: (i -> j -> k) -> IxX sh i -> IxX sh j -> IxX sh k
+ixxZipWith _ ZIX ZIX = ZIX
+ixxZipWith f (i :.% is) (j :.% js) = f i j :.% ixxZipWith f is js
+
+ixxCast :: StaticShX sh' -> IxX sh i -> IxX sh' i
+ixxCast ZKX ZIX = ZIX
+ixxCast (_ :!% sh) (i :.% idx) = i :.% ixxCast sh idx
+ixxCast _ _ = error "ixxCast: ranks don't match"
+
+-- | Given a multidimensional index, get the corresponding linear
+-- index into the buffer.
+{-# INLINEABLE ixxToLinear #-}
+ixxToLinear :: Num i => IShX sh -> IxX sh i -> i
+ixxToLinear = \sh i -> go sh i 0
+  where
+    go :: Num i => IShX sh -> IxX sh i -> i -> i
+    go ZSX ZIX !a = a
+    go (n :$% sh) (i :.% ix) a = go sh ix (fromIntegral (fromSMayNat' n) * a + i)
+
+{-# INLINEABLE ixxFromLinear #-}
+ixxFromLinear :: Num i => IShX sh -> Int -> IxX sh i
+ixxFromLinear = \sh ->  -- give this function arity 1 so that suffixes is shared when it's called many times
+    let suffixes = drop 1 (scanr (*) 1 (shxToList sh))
+    in fromLin0 sh suffixes
+  where
+    -- Unfold first iteration of fromLin to do the range check.
+    -- Don't inline this function at first to allow GHC to inline the outer
+    -- function and realise that 'suffixes' is shared. But then later inline it
+    -- anyway, to avoid the function call. Removing the pragma makes GHC
+    -- somehow unable to recognise that 'suffixes' can be shared in a loop.
+    {-# NOINLINE [0] fromLin0 #-}
+    fromLin0 :: Num i => IShX sh -> [Int] -> Int -> IxX sh i
+    fromLin0 sh suffixes i =
+        if i < 0 then outrange sh i else
+        case (sh, suffixes) of
+          (ZSX, _) | i > 0 -> outrange sh i
+                   | otherwise -> ZIX
+          ((fromSMayNat' -> n) :$% sh', suff : suffs) ->
+            let (q, r) = i `quotRem` suff
+            in if q >= n then outrange sh i else
+               fromIntegral q :.% fromLin sh' suffs r
+          _ -> error "impossible"
+
+    fromLin :: Num i => IShX sh -> [Int] -> Int -> IxX sh i
+    fromLin ZSX _ !_ = ZIX
+    fromLin (_ :$% sh') (suff : suffs) i =
+      let (q, r) = i `quotRem` suff  -- suff == shrSize sh'
+      in fromIntegral q :.% fromLin sh' suffs r
+    fromLin _ _ _ = error "impossible"
+
+    {-# NOINLINE outrange #-}
+    outrange :: IShX sh -> Int -> a
+    outrange sh i = error $ "ixxFromLinear: out of range (" ++ show i ++
+                            " in array of shape " ++ show sh ++ ")"
+
+shxEnum :: IShX sh -> [IIxX sh]
+shxEnum = shxEnum'
+
+{-# INLINABLE shxEnum' #-}  -- ensure this can be specialised at use site
+shxEnum' :: Num i => IShX sh -> [IxX sh i]
+shxEnum' sh = [fromLin sh suffixes li# | I# li# <- [0 .. shxSize sh - 1]]
+  where
+    suffixes = drop 1 (scanr (*) 1 (shxToList sh))
+
+    fromLin :: Num i => IShX sh -> [Int] -> Int# -> IxX sh i
+    fromLin ZSX _ _ = ZIX
+    fromLin (_ :$% sh') (I# suff# : suffs) i# =
+      let !(# q#, r# #) = i# `quotRemInt#` suff#  -- suff == shrSize sh'
+      in fromIntegral (I# q#) :.% fromLin sh' suffs r#
+    fromLin _ _ _ = error "impossible"
+
+
+-- * Mixed shapes
+
+data SMayNat i n where
+  SUnknown :: i -> SMayNat i Nothing
+  SKnown :: SNat n -> SMayNat i (Just n)
+deriving instance Show i => Show (SMayNat i n)
+deriving instance Eq i => Eq (SMayNat i n)
+deriving instance Ord i => Ord (SMayNat i n)
+
+instance NFData i => NFData (SMayNat i n) where
+  rnf (SUnknown i) = rnf i
+  rnf (SKnown SNat) = ()
+
+instance TestEquality (SMayNat i) where
+  testEquality SUnknown{} SUnknown{} = Just Refl
+  testEquality (SKnown n) (SKnown m) | Just Refl <- testEquality n m = Just Refl
+  testEquality _ _ = Nothing
+
+{-# INLINE fromSMayNat #-}
+fromSMayNat :: (n ~ Nothing => i -> r)
+            -> (forall m. n ~ Just m => SNat m -> r)
+            -> SMayNat i n -> r
+fromSMayNat f _ (SUnknown i) = f i
+fromSMayNat _ g (SKnown s) = g s
+
+{-# INLINE fromSMayNat' #-}
+fromSMayNat' :: SMayNat Int n -> Int
+fromSMayNat' = fromSMayNat id 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 n -> SMayNat Int m -> SMayNat Int (AddMaybe n m)
+smnAddMaybe (SUnknown n) m = SUnknown (n + fromSMayNat' m)
+smnAddMaybe (SKnown n) (SUnknown m) = SUnknown (fromSNat' n + m)
+smnAddMaybe (SKnown n) (SKnown m) = SKnown (snatPlus n m)
+
+
+type role ShX nominal representational
+type ShX :: [Maybe Nat] -> Type -> Type
+data ShX sh i where
+  ZSX :: ShX '[] i
+  ConsUnknown :: forall sh i. i -> ShX sh i -> ShX (Nothing : sh) i
+-- TODO: bring this UNPACK back when GHC no longer crashes:
+-- ConsKnown :: forall n sh i. {-# UNPACK #-} SNat n -> ShX sh i -> ShX (Just n : sh) i
+  ConsKnown :: forall n sh i. SNat n -> ShX sh i -> ShX (Just n : sh) i
+deriving instance Ord i => Ord (ShX sh i)
+
+-- A manually defined instance and this INLINEABLE is needed to specialize
+-- mdot1Inner (otherwise GHC warns specialization breaks down here).
+instance Eq i => Eq (ShX sh i) where
+  {-# INLINEABLE (==) #-}
+  ZSX == ZSX = True
+  ConsUnknown i1 sh1 == ConsUnknown i2 sh2 = i1 == i2 && sh1 == sh2
+  ConsKnown _ sh1 == ConsKnown _ sh2 = sh1 == sh2
+
+#ifdef OXAR_DEFAULT_SHOW_INSTANCES
+deriving instance Show i => Show (ShX sh i)
+#else
+instance Show i => Show (ShX sh i) where
+  showsPrec _ l = shxShow (fromSMayNat shows (shows . fromSNat)) l
+#endif
+
+instance NFData i => NFData (ShX sh i) where
+  rnf ZSX = ()
+  rnf (x `ConsUnknown` l) = rnf x `seq` rnf l
+  rnf (SNat `ConsKnown` l) = rnf l
+
+instance Functor (ShX sh) where
+  {-# INLINE fmap #-}
+  fmap f l = shxFmap (fromSMayNat (SUnknown . f) SKnown) l
+
+data UnconsShXRes i sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsShXRes (SMayNat i n) (ShX sh i)
+shxUncons :: ShX sh1 i -> Maybe (UnconsShXRes i sh1)
+shxUncons (i `ConsUnknown` shl') = Just (UnconsShXRes (SUnknown i) shl')
+shxUncons (i `ConsKnown` shl') = Just (UnconsShXRes (SKnown i) shl')
+shxUncons ZSX = 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.
+shxEqType :: ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
+shxEqType ZSX ZSX = Just Refl
+shxEqType (_ `ConsUnknown` sh) (_ `ConsUnknown` sh')
+  | Just Refl <- shxEqType sh sh'
+  = Just Refl
+shxEqType (n `ConsKnown` sh) (m `ConsKnown` sh')
+  | Just Refl <- testEquality n m
+  , Just Refl <- shxEqType sh sh'
+  = Just Refl
+shxEqType _ _ = 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).
+shxEqual :: Eq i => ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
+shxEqual ZSX ZSX = Just Refl
+shxEqual (n `ConsUnknown` sh) (m `ConsUnknown` sh')
+  | n == m
+  , Just Refl <- shxEqual sh sh'
+  = Just Refl
+shxEqual (n `ConsKnown` sh) (m `ConsKnown` sh')
+  | Just Refl <- testEquality n m
+  , Just Refl <- shxEqual sh sh'
+  = Just Refl
+shxEqual _ _ = Nothing
+
+{-# INLINE shxFmap #-}
+shxFmap :: (forall n. SMayNat i n -> SMayNat j n) -> ShX sh i -> ShX sh j
+shxFmap _ ZSX = ZSX
+shxFmap f (x `ConsUnknown` xs) = case f (SUnknown x) of
+  SUnknown y -> y `ConsUnknown` shxFmap f xs
+shxFmap f (x `ConsKnown` xs) = case f (SKnown x) of
+  SKnown y -> y `ConsKnown` shxFmap f xs
+
+{-# INLINE shxFoldMap #-}
+shxFoldMap :: Monoid m => (forall n. SMayNat i n -> m) -> ShX sh i -> m
+shxFoldMap _ ZSX = mempty
+shxFoldMap f (x `ConsUnknown` xs) = f (SUnknown x) <> shxFoldMap f xs
+shxFoldMap f (x `ConsKnown` xs) = f (SKnown x) <> shxFoldMap f xs
+
+shxLength :: ShX sh i -> Int
+shxLength = getSum . shxFoldMap (\_ -> Sum 1)
+
+shxRank :: ShX sh i -> SNat (Rank sh)
+shxRank ZSX = SNat
+shxRank (_ `ConsUnknown` l) | SNat <- shxRank l = SNat
+shxRank (_ `ConsKnown` l) | SNat <- shxRank l = SNat
+
+{-# INLINE shxShow #-}
+shxShow :: forall sh i. (forall n. SMayNat i n -> ShowS) -> ShX sh i -> ShowS
+shxShow f l = showString "[" . go "" l . showString "]"
+  where
+    go :: String -> ShX sh' i -> ShowS
+    go _ ZSX = id
+    go prefix (x `ConsUnknown` xs) = showString prefix . f (SUnknown x) . go "," xs
+    go prefix (x `ConsKnown` xs) = showString prefix . f (SKnown x) . go "," xs
+
+shxHead :: ShX (mn ': sh) i -> SMayNat i mn
+shxHead (i `ConsUnknown` _) = SUnknown i
+shxHead (i `ConsKnown` _) = SKnown i
+
+shxTail :: ShX (n : sh) i -> ShX sh i
+shxTail (_ `ConsUnknown` sh) = sh
+shxTail (_ `ConsKnown` sh) = sh
+
+shxAppend :: ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
+shxAppend ZSX idx' = idx'
+shxAppend (i `ConsUnknown` idx) idx' = i `ConsUnknown` shxAppend idx idx'
+shxAppend (i `ConsKnown` idx) idx' = i `ConsKnown` shxAppend idx idx'
+
+shxDropSh :: forall sh sh' i j. ShX sh j -> ShX (sh ++ sh') i -> ShX sh' i
+shxDropSh ZSX long = long
+shxDropSh (_ `ConsUnknown` short) long = case long of
+  _ `ConsUnknown` long' -> shxDropSh short long'
+shxDropSh (_ `ConsKnown` short) long = case long of
+  _ `ConsKnown` long' -> shxDropSh short long'
+
+shxDropSSX :: forall sh sh' i. StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
+shxDropSSX = coerce (shxDropSh @_ @_ @i @())
+
+shxInit :: forall i n sh. ShX (n : sh) i -> ShX (Init (n : sh)) i
+shxInit (i `ConsUnknown` sh@(_ `ConsUnknown` _)) = i `ConsUnknown` shxInit sh
+shxInit (i `ConsUnknown` sh@(_ `ConsKnown` _)) = i `ConsUnknown` shxInit sh
+shxInit (_ `ConsUnknown` ZSX) = ZSX
+shxInit (i `ConsKnown` sh@(_ `ConsUnknown` _)) = i `ConsKnown` shxInit sh
+shxInit (i `ConsKnown` sh@(_ `ConsKnown` _)) = i `ConsKnown` shxInit sh
+shxInit (_ `ConsKnown` ZSX) = ZSX
+
+shxLast :: forall i n sh. ShX (n : sh) i -> SMayNat i (Last (n : sh))
+shxLast (_ `ConsUnknown` sh@(_ `ConsUnknown` _)) = shxLast sh
+shxLast (_ `ConsUnknown` sh@(_ `ConsKnown` _)) = shxLast sh
+shxLast (x `ConsUnknown` ZSX) = SUnknown x
+shxLast (_ `ConsKnown` sh@(_ `ConsUnknown` _)) = shxLast sh
+shxLast (_ `ConsKnown` sh@(_ `ConsKnown` _)) = shxLast sh
+shxLast (x `ConsKnown` ZSX) = SKnown x
+
+pattern (:$%)
+  :: forall {sh1} {i}.
+     forall n sh. (n : sh ~ sh1)
+  => SMayNat i n -> ShX sh i -> ShX sh1 i
+pattern i :$% shl <- (shxUncons -> Just (UnconsShXRes i shl))
+  where i :$% shl = case i of; SUnknown x -> x `ConsUnknown` shl; SKnown x -> x `ConsKnown` shl
+infixr 3 :$%
+
+{-# COMPLETE ZSX, (:$%) #-}
+
+type IShX sh = ShX sh Int
+
+-- | The number of elements in an array described by this shape.
+shxSize :: IShX sh -> Int
+shxSize ZSX = 1
+shxSize (n :$% sh) = fromSMayNat' n * shxSize sh
+
+-- We don't report the size of the list in case of errors in order not to retain the list.
+{-# INLINEABLE shxFromList #-}
+shxFromList :: StaticShX sh -> [Int] -> IShX sh
+shxFromList (StaticShX topssh) topl = go topssh topl
+  where
+    go :: ShX sh' () -> [Int] -> ShX sh' Int
+    go ZSX [] = ZSX
+    go ZSX _ = error $ "shxFromList: List too long (type says " ++ show (shxLength topssh) ++ ")"
+    go (ConsKnown sn sh) (i : is)
+      | i == fromSNat' sn = ConsKnown sn (go sh is)
+      | otherwise = error "shxFromList: Value does not match typing"
+    go (ConsUnknown () sh) (i : is) = ConsUnknown i (go sh is)
+    go _ _ = error $ "shxFromList: List too short (type says " ++ show (shxLength topssh) ++ ")"
+
+{-# INLINEABLE shxToList #-}
+shxToList :: IShX sh -> [Int]
+shxToList l = build (\(cons :: i -> is -> is) (nil :: is) ->
+  let go :: ShX sh Int -> is
+      go ZSX = nil
+      go (ConsUnknown i rest) = i `cons` go rest
+      go (ConsKnown sn rest) = fromSNat' sn `cons` go rest
+  in go l)
+
+-- If it ever matters for performance, this is unsafeCoercible.
+shxFromSSX :: StaticShX (MapJust sh) -> ShX (MapJust sh) i
+shxFromSSX ZKX = ZSX
+shxFromSSX (SKnown n :!% sh :: StaticShX (MapJust sh))
+  | Refl <- lemMapJustCons @sh Refl
+  = SKnown n :$% shxFromSSX sh
+shxFromSSX (SUnknown _ :!% _) = error "unreachable"
+
+-- | This may fail if @sh@ has @Nothing@s in it.
+shxFromSSX2 :: StaticShX sh -> Maybe (ShX sh i)
+shxFromSSX2 ZKX = Just ZSX
+shxFromSSX2 (SKnown n :!% sh) = (SKnown n :$%) <$> shxFromSSX2 sh
+shxFromSSX2 (SUnknown _ :!% _) = Nothing
+
+shxTakeSSX :: forall sh sh' i proxy. proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
+shxTakeSSX _ ZKX _ = ZSX
+shxTakeSSX p (_ :!% ssh1) (n :$% sh) = n :$% shxTakeSSX p ssh1 sh
+
+shxTakeSh :: forall sh sh' i proxy. proxy sh' -> ShX sh i -> ShX (sh ++ sh') i -> ShX sh i
+shxTakeSh _ ZSX _ = ZSX
+shxTakeSh p (_ :$% ssh1) (n :$% sh) = n :$% shxTakeSh p ssh1 sh
+
+{-# INLINEABLE shxTakeIx #-}
+shxTakeIx :: forall sh sh' i j. Proxy sh' -> IxX sh j -> ShX (sh ++ sh') i -> ShX sh i
+shxTakeIx _ (IxX ZX) _ = ZSX
+shxTakeIx proxy (IxX (_ ::% long)) short = case short of i :$% short' -> i :$% shxTakeIx proxy (IxX long) short'
+
+{-# INLINEABLE shxDropIx #-}
+shxDropIx :: forall sh sh' i j. IxX sh j -> ShX (sh ++ sh') i -> ShX sh' i
+shxDropIx ZIX long = long
+shxDropIx (_ :.% short) long = case long of _ :$% long' -> shxDropIx short long'
+
+{-# INLINE shxZipWith #-}
+shxZipWith :: (forall n. SMayNat i n -> SMayNat j n -> SMayNat k n)
+           -> ShX sh i -> ShX sh j -> ShX sh k
+shxZipWith _ ZSX ZSX = ZSX
+shxZipWith f (i :$% is) (j :$% js) = f i j :$% shxZipWith f is js
+
+-- This is a weird operation, so it has a long name
+shxCompleteZeros :: StaticShX sh -> IShX sh
+shxCompleteZeros ZKX = ZSX
+shxCompleteZeros (SUnknown () :!% ssh) = SUnknown 0 :$% shxCompleteZeros ssh
+shxCompleteZeros (SKnown n :!% ssh) = SKnown n :$% shxCompleteZeros ssh
+
+shxSplitApp :: proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
+shxSplitApp _ ZKX idx = (ZSX, idx)
+shxSplitApp p (_ :!% ssh) (i :$% idx) = first (i :$%) (shxSplitApp p ssh idx)
+
+shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh')
+shxCast ZKX ZSX = Just ZSX
+shxCast (SKnown m    :!% ssh) (SKnown n   :$% sh) | Just Refl <- testEquality n m = (SKnown n :$%) <$> shxCast ssh sh
+shxCast (SKnown m    :!% ssh) (SUnknown n :$% sh) | n == fromSNat' m              = (SKnown m :$%) <$> shxCast ssh sh
+shxCast (SUnknown () :!% ssh) (SKnown n   :$% sh)                                 = (SUnknown (fromSNat' n) :$%) <$> shxCast ssh sh
+shxCast (SUnknown () :!% ssh) (SUnknown n :$% sh)                                 = (SUnknown n :$%) <$> shxCast ssh sh
+shxCast _ _ = Nothing
+
+-- | Partial version of 'shxCast'.
+shxCast' :: StaticShX sh' -> IShX sh -> IShX sh'
+shxCast' ssh sh = case shxCast ssh sh of
+  Just sh' -> sh'
+  Nothing -> error $ "shxCast': Mismatch: (" ++ show sh ++ ") does not match (" ++ show ssh ++ ")"
+
+
+-- * Static mixed shapes
+
+-- | The part of a shape that is statically known. (A newtype over 'ShX'.)
+type StaticShX :: [Maybe Nat] -> Type
+newtype StaticShX sh = StaticShX (ShX sh ())
+  deriving (NFData)
+
+instance Eq (StaticShX sh) where _ == _ = True
+instance Ord (StaticShX sh) where compare _ _ = EQ
+
+pattern ZKX :: forall sh. () => sh ~ '[] => StaticShX sh
+pattern ZKX = StaticShX ZSX
+
+pattern (:!%)
+  :: forall {sh1}.
+     forall n sh. (n : sh ~ sh1)
+  => SMayNat () n -> StaticShX sh -> StaticShX sh1
+pattern i :!% shl <- StaticShX (shxUncons -> Just (UnconsShXRes i (StaticShX -> shl)))
+  where i :!% StaticShX shl = case i of; SUnknown () -> StaticShX (() `ConsUnknown` shl); SKnown x -> StaticShX (x `ConsKnown` shl)
+
+infixr 3 :!%
+
+{-# COMPLETE ZKX, (:!%) #-}
+
+#ifdef OXAR_DEFAULT_SHOW_INSTANCES
+deriving instance Show (StaticShX sh)
+#else
+instance Show (StaticShX sh) where
+  showsPrec _ (StaticShX l) = shxShow (fromSMayNat shows (shows . fromSNat)) l
+#endif
+
+instance TestEquality StaticShX where
+  testEquality (StaticShX l1) (StaticShX l2) = shxEqType l1 l2
+
+ssxLength :: StaticShX sh -> Int
+ssxLength (StaticShX l) = shxLength l
+
+ssxRank :: StaticShX sh -> SNat (Rank sh)
+ssxRank (StaticShX l) = shxRank l
+
+-- | @ssxEqType = 'testEquality'@. Provided for consistency.
+ssxEqType :: StaticShX sh -> StaticShX sh' -> Maybe (sh :~: sh')
+ssxEqType = testEquality
+
+ssxAppend :: StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
+ssxAppend = coerce (shxAppend @_ @())
+
+ssxHead :: StaticShX (n : sh) -> SMayNat () n
+ssxHead (StaticShX list) = shxHead list
+
+ssxTail :: StaticShX (n : sh) -> StaticShX sh
+ssxTail (StaticShX list) = StaticShX (shxTail list)
+
+ssxTakeIx :: forall sh sh' i. Proxy sh' -> IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh
+ssxTakeIx _ (IxX ZX) _ = ZKX
+ssxTakeIx proxy (IxX (_ ::% long)) short = case short of i :!% short' -> i :!% ssxTakeIx proxy (IxX long) short'
+
+ssxDropIx :: forall sh sh' i. IxX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
+ssxDropIx (IxX ZX) long = long
+ssxDropIx (IxX (_ ::% short)) long = case long of _ :!% long' -> ssxDropIx (IxX short) long'
+
+ssxDropSh :: forall sh sh' i. ShX sh i -> StaticShX (sh ++ sh') -> StaticShX sh'
+ssxDropSh = coerce (shxDropSh @_ @_ @() @i)
+
+ssxDropSSX :: forall sh sh'. StaticShX sh -> StaticShX (sh ++ sh') -> StaticShX sh'
+ssxDropSSX = coerce (shxDropSh @_ @_ @() @())
+
+ssxInit :: forall n sh. StaticShX (n : sh) -> StaticShX (Init (n : sh))
+ssxInit = coerce (shxInit @())
+
+ssxLast :: forall n sh. StaticShX (n : sh) -> SMayNat () (Last (n : sh))
+ssxLast = coerce (shxLast @())
+
+ssxReplicate :: SNat n -> StaticShX (Replicate n Nothing)
+ssxReplicate SZ = ZKX
+ssxReplicate (SS (n :: SNat n'))
+  | Refl <- lemReplicateSucc @(Nothing @Nat) n
+  = SUnknown () :!% ssxReplicate n
+
+ssxIotaFrom :: StaticShX sh -> Int -> [Int]
+ssxIotaFrom ZKX _ = []
+ssxIotaFrom (_ :!% ssh) i = i : ssxIotaFrom ssh (i + 1)
+
+ssxFromShX :: ShX sh i -> StaticShX sh
+ssxFromShX ZSX = ZKX
+ssxFromShX (n :$% sh) = fromSMayNat (\_ -> SUnknown ()) SKnown n :!% ssxFromShX sh
+
+ssxFromSNat :: SNat n -> StaticShX (Replicate n Nothing)
+ssxFromSNat SZ = ZKX
+ssxFromSNat (SS (n :: SNat nm1)) | Refl <- lemReplicateSucc @(Nothing @Nat) n = SUnknown () :!% ssxFromSNat 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 = ZKX
+instance (KnownNat n, KnownShX sh) => KnownShX (Just n : sh) where knownShX = SKnown natSing :!% knownShX
+instance KnownShX sh => KnownShX (Nothing : sh) where knownShX = SUnknown () :!% knownShX
+
+withKnownShX :: forall sh r. StaticShX sh -> (KnownShX sh => r) -> r
+withKnownShX = 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 () (Flatten sh)
+ssxFlatten = go (SNat @1)
+  where
+    go :: SNat acc -> StaticShX sh -> SMayNat () (Flatten' acc sh)
+    go acc ZKX = SKnown acc
+    go _ (SUnknown () :!% _) = SUnknown ()
+    go acc (SKnown sn :!% sh) = go (snatMul acc sn) sh
+
+shxFlatten :: IShX sh -> SMayNat Int (Flatten sh)
+shxFlatten = go (SNat @1)
+  where
+    go :: SNat acc -> IShX sh -> SMayNat Int (Flatten' acc sh)
+    go acc ZSX = SKnown acc
+    go acc (SUnknown n :$% sh) = SUnknown (goUnknown (fromSNat' acc * n) sh)
+    go acc (SKnown sn :$% sh) = go (snatMul acc sn) sh
+
+    goUnknown :: Int -> IShX sh -> Int
+    goUnknown acc ZSX = acc
+    goUnknown acc (SUnknown n :$% sh) = goUnknown (acc * n) sh
+    goUnknown acc (SKnown sn :$% sh) = goUnknown (acc * fromSNat' sn) sh
+
+
+-- | Very untyped: only length is checked (at runtime), index bounds are __not checked__.
+instance IsList (IxX sh i) where
+  type Item (IxX sh i) = i
+  fromList = IxX . IsList.fromList
+  toList = Foldable.toList
+
+-- | Untyped: length and known dimensions are checked (at runtime).
+instance KnownShX sh => IsList (IShX sh) where
+  type Item (IShX sh) = Int
+  fromList = shxFromList (knownShX @sh)
+  toList = shxToList