From a158f928c6d617fa950f74c517e94c5062422efd Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Tue, 18 Nov 2025 12:57:17 +0100 Subject: Generalise sh?Enum functions to arbitrary Num The explicit usage at Int in the non-primed versions of the functions results in an auto-generated SPEC rule at Int, so even without any SPECIALIZE pragmas, using sh?Enum' at Int should not (need to) result in multiple specialisations. --- src/Data/Array/Nested/Ranked/Shape.hs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) (limited to 'src/Data/Array/Nested/Ranked') diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs index a1fa247..200fac5 100644 --- a/src/Data/Array/Nested/Ranked/Shape.hs +++ b/src/Data/Array/Nested/Ranked/Shape.hs @@ -331,11 +331,15 @@ shrPermutePrefix :: forall n i. [Int] -> ShR n i -> ShR n i shrPermutePrefix = coerce (listrPermutePrefix @i) shrEnum :: IShR sh -> [IIxR sh] -shrEnum = \sh -> go sh id [] +shrEnum = shrEnum' + +{-# INLINABLE shrEnum' #-} -- ensure this can be specialised at use site +shrEnum' :: Num i => IShR sh -> [IxR sh i] +shrEnum' = \sh -> go sh id [] where - go :: IShR sh -> (IIxR sh -> a) -> [a] -> [a] + go :: Num i => IShR sh -> (IxR sh i -> a) -> [a] -> [a] go ZSR f = (f ZIR :) - go (n :$: sh) f = foldr (.) id [go sh (f . (i :.:)) | i <- [0 .. n - 1]] + go (n :$: sh) f = foldr (.) id [go sh (f . (fromIntegral i :.:)) | i <- [0 .. n - 1]] -- | Untyped: length is checked at runtime. -- cgit v1.2.3-70-g09d2