diff options
Diffstat (limited to 'src/Data/Array/Nested/Ranked/Shape.hs')
| -rw-r--r-- | src/Data/Array/Nested/Ranked/Shape.hs | 19 |
1 files changed, 17 insertions, 2 deletions
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs index 74f4c5c..2f20e1a 100644 --- a/src/Data/Array/Nested/Ranked/Shape.hs +++ b/src/Data/Array/Nested/Ranked/Shape.hs @@ -390,12 +390,27 @@ shrRank (ShR sh) | Refl <- lemRankReplicate (Proxy @n) = shxRank sh shrSize :: IShR n -> Int shrSize (ShR sh) = shxSize sh +-- This is equivalent to but faster than @coerce (shxFromList (ssxReplicate snat))@. +-- We don't report the size of the list in case of errors in order not to retain the list. +{-# INLINEABLE shrFromList #-} shrFromList :: SNat n -> [Int] -> IShR n -shrFromList snat = coerce (shxFromList (ssxReplicate snat)) +shrFromList snat topl = ShR $ ShX $ go snat topl + where + go :: SNat n -> [Int] -> ListH (Replicate n Nothing) Int + go SZ [] = ZH + go SZ _ = error $ "shrFromList: List too long (type says " ++ show (fromSNat' snat) ++ ")" + go (SS sn :: SNat n1) (i : is) | Refl <- lemReplicateSucc2 (Proxy @n1) Refl = ConsUnknown i (go sn is) + go _ _ = error $ "shrFromList: List too short (type says " ++ show (fromSNat' snat) ++ ")" +-- This is equivalent to but faster than @coerce shxToList@. {-# INLINEABLE shrToList #-} shrToList :: IShR n -> [Int] -shrToList = coerce shxToList +shrToList (ShR (ShX l)) = build (\(cons :: i -> is -> is) (nil :: is) -> + let go :: ListH sh Int -> is + go ZH = nil + go (ConsUnknown i rest) = i `cons` go rest + go ConsKnown{} = error "shrToList: impossible case" + in go l) shrHead :: forall n i. ShR (n + 1) i -> i shrHead (ShR sh) |