From fc2a1370d67f12b50e3e5750d17aefd33bc3d8a3 Mon Sep 17 00:00:00 2001 From: Mikolaj Konarski Date: Mon, 16 Mar 2026 14:36:36 +0100 Subject: Fill and clean up *TakeIx and *DropIx functions --- src/Data/Array/Nested/Shaped/Shape.hs | 11 +++++++++++ 1 file changed, 11 insertions(+) (limited to 'src/Data/Array/Nested/Shaped') diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs index f98c860..ef91c7b 100644 --- a/src/Data/Array/Nested/Shaped/Shape.hs +++ b/src/Data/Array/Nested/Shaped/Shape.hs @@ -30,6 +30,7 @@ import Data.Array.Shape qualified as O import Data.Coerce (coerce) import Data.Foldable qualified as Foldable import Data.Kind (Constraint, Type) +import Data.Proxy import Data.Type.Equality import GHC.Exts (build, withDict) import GHC.IsList (IsList) @@ -394,6 +395,16 @@ shsHead (ShS shx) = case shxHead shx of shsTail :: forall n sh. ShS (n : sh) -> ShS sh shsTail = coerce (shxTail @_ @_ @Int) +{-# INLINEABLE shsTakeIx #-} +shsTakeIx :: forall sh sh' j. Proxy sh' -> IxS sh j -> ShS (sh ++ sh') -> ShS sh +shsTakeIx _ ZIS _ = ZSS +shsTakeIx p (_ :.$ idx) sh = case sh of n :$$ sh' -> n :$$ shsTakeIx p idx sh' + +{-# INLINEABLE shsDropIx #-} +shsDropIx :: forall sh sh' j. IxS sh j -> ShS (sh ++ sh') -> ShS sh' +shsDropIx ZIS long = long +shsDropIx (_ :.$ short) long = case long of _ :$$ long' -> shsDropIx short long' + shsInit :: forall n sh. ShS (n : sh) -> ShS (Init (n : sh)) shsInit = gcastWith (unsafeCoerceRefl -- cgit v1.2.3-70-g09d2