diff --git a/changelog.md b/changelog.md index f2681bd..7455533 100644 --- a/changelog.md +++ b/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 diff --git a/src/Data/Chimera.hs b/src/Data/Chimera.hs index f0e1cbf..927df12 100644 --- a/src/Data/Chimera.hs +++ b/src/Data/Chimera.hs @@ -48,6 +48,7 @@ module Data.Chimera ( -- * Subvectors -- $subvectors mapSubvectors, + imapSubvectors, traverseSubvectors, zipWithSubvectors, zipWithMSubvectors, diff --git a/src/Data/Chimera/Internal.hs b/src/Data/Chimera/Internal.hs index e3a169c..8dfdad0 100644 --- a/src/Data/Chimera/Internal.hs +++ b/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 diff --git a/test/Test.hs b/test/Test.hs index b5d8c63..c40926c 100644 --- a/test/Test.hs +++ b/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