#36 Add imapSubvectors

changelog.md

@@ -4,6 +4,7 @@
 * Add `HalfWord` and `ThirdWord` types,
   change types of `toZCurve`, `fromZCurve`, `toZCurve3`, `fromZCurve3` accordingly.
 * Add `throughZCurveFix` and `throughZCurveFix3`.
+* Add `imapSubvectors`.
 
 # 0.3.4.0
 

src/Data/Chimera.hs

@@ -48,6 +48,7 @@ module Data.Chimera (
   -- * Subvectors
   -- $subvectors
   mapSubvectors,
+  imapSubvectors,
   traverseSubvectors,
   zipWithSubvectors,
   zipWithMSubvectors,

src/Data/Chimera/Internal.hs

@@ -53,6 +53,7 @@ module Data.Chimera.Internal (
 
   -- * Subvectors
   mapSubvectors,
+  imapSubvectors,
   traverseSubvectors,
   zipWithSubvectors,
   zipWithMSubvectors,
@@ -624,6 +625,27 @@ mapSubvectors
   -> Chimera v b
 mapSubvectors f = runIdentity . traverseSubvectors (coerce f)
 
+-- | Map subvectors of a stream, using a given length-preserving function.
+-- The first argument of the function is the index of the first element of subvector
+-- in the 'Chimera'.
+--
+-- @since 0.4.0.0
+imapSubvectors
+  :: (G.Vector u a, G.Vector v b)
+  => (Word -> u a -> v b)
+  -> Chimera u a
+  -> Chimera v b
+imapSubvectors f (Chimera bs) = Chimera $ mzipWith safeF (fromListN (bits + 1) [0 .. bits]) bs
+  where
+    -- Computing vector length is cheap, so let's check that @f@ preserves length.
+    safeF i x =
+      if xLen == G.length fx
+        then fx
+        else error "imapSubvectors: the function is not length-preserving"
+      where
+        xLen = G.length x
+        fx = f (if i == 0 then 0 else 1 `unsafeShiftL` (i - 1)) x
+
 -- | Traverse subvectors of a stream, using a given length-preserving function.
 --
 -- Be careful, because similar to 'tabulateM', only lazy monadic effects can

test/Test.hs

@@ -165,12 +165,19 @@ chimeraTests = testGroup "Chimera"
           (if fromIntegral jx < length xs then vs G.! fromIntegral jx else x) ===
             Ch.index (Ch.fromVectorWithDef x vs :: UChimera Bool) jx
 
-  , QC.testProperty "mapWithKey" $
+  , QC.testProperty "mapSubvectors" $
     \(Blind bs) (Fun _ (g :: Word -> Word)) ix ->
       let jx = ix `mod` 65536 in
         g (Ch.index bs jx) === Ch.index (Ch.mapSubvectors (G.map g) bs :: UChimera Word) jx
 
-  , QC.testProperty "zipWithKey" $
+  , QC.testProperty "imapSubvectors" $
+    \(Blind bs) (Fun _ (g :: (Word, Int) -> Char)) ix ->
+      let jx = ix `mod` 65536 in
+        curry g jx (Ch.index bs jx) ===
+          Ch.index (Ch.imapSubvectors (\off ->
+            G.imap (curry g . (+ off) . fromIntegral)) bs :: UChimera Char) jx
+
+  , QC.testProperty "zipWithSubvectors" $
     \(Blind bs1) (Blind bs2) (Fun _ (g :: (Word, Word) -> Word)) ix ->
       let jx = ix `mod` 65536 in
         g (Ch.index bs1 jx, Ch.index bs2 jx) === Ch.index (Ch.zipWithSubvectors (G.zipWith (curry g)) bs1 bs2 :: UChimera Word) jx