From b63642a41f3bddc991d92f2f59b9e3ad53c1f15e Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Tue, 2 Dec 2025 15:03:10 +0100
Subject: Provide ix*FromLinear for all three shape kinds

This speeds up {r,s}generatePrim
---
 src/Data/Array/Nested/Ranked.hs | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

(limited to 'src/Data/Array/Nested/Ranked.hs')

diff --git a/src/Data/Array/Nested/Ranked.hs b/src/Data/Array/Nested/Ranked.hs
index b77b529..d687983 100644
--- a/src/Data/Array/Nested/Ranked.hs
+++ b/src/Data/Array/Nested/Ranked.hs
@@ -70,16 +70,13 @@ rgenerate sh f
   , Refl <- lemRankReplicate sn
   = Ranked (mgenerate (shxFromShR sh) (f . ixrFromIxX))
 
--- TODO: this would be shorter and faster written with rfromVector,
--- but unfortunately we don't have ixrFromLinear
+-- | See 'mgeneratePrim'.
 {-# INLINE rgeneratePrim #-}
 rgeneratePrim :: forall n a i. (PrimElt a, Num i)
               => IShR n -> (IxR n i -> a) -> Ranked n a
-rgeneratePrim sh f
-  | sn@SNat <- shrRank sh
-  , Dict <- lemKnownReplicate sn
-  , Refl <- lemRankReplicate sn
-  = Ranked (mgeneratePrim (shxFromShR sh) (f . ixrFromIxX))
+rgeneratePrim sh f =
+  let g i = f (ixrFromLinear sh i)
+  in rfromVector sh $ VS.generate (shrSize sh) g
 
 -- | See the documentation of 'mlift'.
 rlift :: forall n1 n2 a. Elt a
-- 
cgit v1.2.3-70-g09d2