From c4c8fe59dd233faf00916a7c0eba5d382ccdc235 Mon Sep 17 00:00:00 2001
From: Bodigrim <andrew.lelechenko@gmail.com>
Date: Thu, 28 Dec 2023 22:11:57 +0000
Subject: [PATCH] #35 Add prependVector

---
 changelog.md                 |  1 +
 src/Data/Chimera.hs          |  1 +
 src/Data/Chimera/Internal.hs | 41 ++++++++++++++++++++++++++++++++++++
 test/Test.hs                 |  7 ++++++
 4 files changed, 50 insertions(+)

diff --git a/changelog.md b/changelog.md
index 7455533..423fde8 100644
--- a/changelog.md
+++ b/changelog.md
@@ -5,6 +5,7 @@
   change types of `toZCurve`, `fromZCurve`, `toZCurve3`, `fromZCurve3` accordingly.
 * Add `throughZCurveFix` and `throughZCurveFix3`.
 * Add `imapSubvectors`.
+* Add `prependVector`.
 
 # 0.3.4.0
 
diff --git a/src/Data/Chimera.hs b/src/Data/Chimera.hs
index 927df12..23a8a8f 100644
--- a/src/Data/Chimera.hs
+++ b/src/Data/Chimera.hs
@@ -29,6 +29,7 @@ module Data.Chimera (
 
   -- * Manipulation
   interleave,
+  prependVector,
 
   -- * Elimination
   index,
diff --git a/src/Data/Chimera/Internal.hs b/src/Data/Chimera/Internal.hs
index 8dfdad0..fea0ea3 100644
--- a/src/Data/Chimera/Internal.hs
+++ b/src/Data/Chimera/Internal.hs
@@ -10,6 +10,7 @@
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE ViewPatterns #-}
 
 -- |
 -- Module:      Data.Chimera.Internal
@@ -36,6 +37,7 @@ module Data.Chimera.Internal (
 
   -- * Manipulation
   interleave,
+  prependVector,
 
   -- * Elimination
   index,
@@ -600,6 +602,45 @@ fromVectorWithDef a = Chimera . fromListN (bits + 1) . go0
       where
         kk = 1 `shiftL` k
 
+-- | Prepend a given vector to a stream of values.
+--
+-- @since 0.4.0.0
+prependVector
+  :: forall v a
+   . G.Vector v a
+  => v a
+  -> Chimera v a
+  -> Chimera v a
+prependVector (G.uncons -> Nothing) ch = ch
+prependVector (G.uncons -> Just (pref0, pref)) (Chimera as) =
+  Chimera $
+    fromListN (bits + 1) $
+      fmap sliceAndConcat $
+        [LazySlice 0 1 $ G.singleton pref0] : go 0 1 0 inputs
+  where
+    inputs :: [(Word, v a)]
+    inputs =
+      (int2word $ G.length pref, pref)
+        : zip (1 : map (1 `unsafeShiftL`) [0 .. bits - 1]) (F.toList as)
+
+    go :: Int -> Word -> Word -> [(Word, t)] -> [[LazySlice t]]
+    go _ _ _ [] = []
+    go n need off orig@((lt, t) : rest)
+      | n >= bits = []
+      | otherwise = case compare (off + need) lt of
+          LT -> [LazySlice off need t] : go (n + 1) (1 `shiftL` (n + 1)) (off + need) orig
+          EQ -> [LazySlice off need t] : go (n + 1) (1 `shiftL` (n + 1)) 0 rest
+          GT -> case go n (off + need - lt) 0 rest of
+            [] -> error "prependVector: the stream should not get exhausted prematurely"
+            hd : tl -> (LazySlice off (lt - off) t : hd) : tl
+
+data LazySlice a = LazySlice !Word !Word a
+
+sliceAndConcat :: G.Vector v a => [LazySlice (v a)] -> v a
+sliceAndConcat =
+  G.concat
+    . map (\(LazySlice from len vec) -> G.slice (word2int from) (word2int len) vec)
+
 -- | Return an infinite repetition of a given vector.
 -- Throw an error on an empty vector.
 --
diff --git a/test/Test.hs b/test/Test.hs
index c40926c..cf0956b 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -165,6 +165,13 @@ 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 "prependVector" $
+    \(Blind bs) xs ix ->
+      let jx = ix `mod` 65536 in
+        let vs = G.fromList xs in
+          (if fromIntegral jx < length xs then vs G.! fromIntegral jx else Ch.index bs (min 65555 $ jx - fromIntegral (length xs))) ===
+            Ch.index (Ch.prependVector vs bs :: UChimera Bool) jx
+
   , QC.testProperty "mapSubvectors" $
     \(Blind bs) (Fun _ (g :: Word -> Word)) ix ->
       let jx = ix `mod` 65536 in