1 files changed, 24 insertions, 20 deletions
diff --git a/src/ImmutGrowVector.hs b/src/ImmutGrowVector.hs
index d36209d..ee38bc4 100644
--- a/src/ImmutGrowVector.hs
+++ b/src/ImmutGrowVector.hs
@@ -1,43 +1,47 @@
 module ImmutGrowVector where
 
-import Foreign.Storable
-import Data.Vector.Storable qualified as VS
+import Data.Vector.Unboxed (Unbox)
+import Data.Vector.Unboxed qualified as VU
 
 
 -- | Acts like an immutable storable vector, except that it's split in a long
 -- prefix and a short suffix so that modifications at the end are cheap. If the
 -- suffix gets longer (by appending elements), the suffix elements are promoted
 -- to prefix elements once in a while, resulting in a big copy at those times.
-data ImmutGrowVector a = ImmutGrowVector (VS.Vector a) (VS.Vector a)
+data ImmutGrowVector a = ImmutGrowVector (VU.Vector a) (VU.Vector a)
   deriving (Show)
 
-empty :: Storable a => ImmutGrowVector a
-empty = ImmutGrowVector VS.empty VS.empty
+empty :: Unbox a => ImmutGrowVector a
+empty = ImmutGrowVector VU.empty VU.empty
 
-fromListN :: Storable a => Int -> [a] -> ImmutGrowVector a
+fromListN :: Unbox a => Int -> [a] -> ImmutGrowVector a
 fromListN n l
   | n > 2 =
       let (l1, l2) = splitAt (n - 2) l
-      in ImmutGrowVector (VS.fromListN (n - 2) l1) (VS.fromListN 2 l2)
+      in ImmutGrowVector (VU.fromListN (n - 2) l1) (VU.fromListN 2 l2)
   | otherwise =
-      ImmutGrowVector VS.empty (VS.fromListN n l)
+      ImmutGrowVector VU.empty (VU.fromListN n l)
 
-(!) :: Storable a => ImmutGrowVector a -> Int -> a
+(!) :: Unbox a => ImmutGrowVector a -> Int -> a
 ImmutGrowVector prefix suffix ! i
-  | i < VS.length prefix = prefix VS.! i
-  | otherwise = suffix VS.! (i - VS.length prefix)
+  | i < VU.length prefix = prefix VU.! i
+  | otherwise = suffix VU.! (i - VU.length prefix)
 
-length :: Storable a => ImmutGrowVector a -> Int
-length (ImmutGrowVector prefix suffix) = VS.length prefix + VS.length suffix
+length :: Unbox a => ImmutGrowVector a -> Int
+length (ImmutGrowVector prefix suffix) = VU.length prefix + VU.length suffix
 
-set :: Storable a => ImmutGrowVector a -> Int -> a -> ImmutGrowVector a
+set :: Unbox a => ImmutGrowVector a -> Int -> a -> ImmutGrowVector a
 set (ImmutGrowVector prefix suffix) idx value
-  | idx < VS.length prefix = error "ImmutGrowVector: mutation in slow part"
-  | otherwise = ImmutGrowVector prefix (suffix VS.// [(idx - VS.length prefix, value)])
+  | idx < VU.length prefix = error "ImmutGrowVector: mutation in slow part"
+  | otherwise = ImmutGrowVector prefix (suffix VU.// [(idx - VU.length prefix, value)])
 
-append :: Storable a => ImmutGrowVector a -> a -> ImmutGrowVector a
+append :: Unbox a => ImmutGrowVector a -> a -> ImmutGrowVector a
 append (ImmutGrowVector prefix suffix) value
-  | VS.length suffix < 8 = ImmutGrowVector prefix (suffix `VS.snoc` value)
+  | VU.length suffix < 8 = ImmutGrowVector prefix (suffix `VU.snoc` value)
   | otherwise =
-      let n = VS.length suffix
-      in ImmutGrowVector (prefix <> VS.take (n - 1) suffix) (VS.drop (n - 1) suffix `VS.snoc` value)
+      let n = VU.length suffix
+      in ImmutGrowVector (prefix <> VU.take (n - 1) suffix) (VU.drop (n - 1) suffix `VU.snoc` value)
+
+-- | Transform the vector by mapping the underlying unboxed vectors
+mapUVector :: (VU.Vector a -> VU.Vector b) -> ImmutGrowVector a -> ImmutGrowVector b
+mapUVector f (ImmutGrowVector prefix suffix) = ImmutGrowVector (f prefix) (f suffix)