Add ability to shrink/slice vectors and arrays (#181)

Author: utdemir

src/Data/Array/Mutable/Linear.hs

@@ -53,6 +53,7 @@ module Data.Array.Mutable.Linear
     read,
     unsafeRead,
     size,
+    slice,
     toList,
   )
 where
@@ -167,7 +168,7 @@ resize newSize seed (Array arr :: Array a)
       doCopy (Unlifted.allocBeside newSize seed arr)
      where
       doCopy :: (# Array# a, Array# a #) #-> Array a
-      doCopy (# src, dest #) = wrap (Unlifted.copyInto src dest)
+      doCopy (# src, dest #) = wrap (Unlifted.copyInto 0 src dest)
 
       wrap :: (# Array# a, Array# a #) #-> Array a
       wrap (# old, new #) = old `Unlifted.lseq` Array new
@@ -177,6 +178,42 @@ resize newSize seed (Array arr :: Array a)
 toList :: Array a #-> Ur [a]
 toList (Array arr) = Unlifted.toList arr
 
+-- | Copy a slice of the array, starting from given offset and copying given
+-- number of elements. Returns the pair (oldArray, slice).
+--
+-- Start offset + target size should be within the input array, and both should
+-- be non-negative.
+--
+-- @
+-- let b = slice i n a,
+--   then size b = n,
+--   and b[j] = a[i+j] for 0 <= j < n
+-- @
+slice
+  :: HasCallStack
+  => Int -- ^ Start offset
+  -> Int -- ^ Target size
+  -> Array a #-> (Array a, Array a)
+slice from targetSize arr =
+  size arr & \case
+    (Array old, Ur s)
+      | s < from + targetSize ->
+          Unlifted.lseq
+            old
+            (error "Slice index out of bounds.")
+      | otherwise ->
+          doCopy
+            (Unlifted.allocBeside
+               targetSize
+               (error "invariant violation: uninitialized array index")
+               old)
+  where
+    doCopy :: (# Array# a, Array# a #) #-> (Array a, Array a)
+    doCopy (# old, new #) = wrap (Unlifted.copyInto from old new)
+
+    wrap :: (# Array# a, Array# a  #) #-> (Array a, Array a)
+    wrap (# old, new #) = (Array old, Array new)
+
 -- # Instances
 -------------------------------------------------------------------------------
 

src/Data/Array/Mutable/Unlifted/Linear.hs

@@ -101,18 +101,25 @@ write (GHC.I# i) (a :: a) = Unsafe.toLinear go
         _ -> Array# arr
 {-# NOINLINE write #-}  -- prevents the runRW# effect from being reordered
 
--- | Copy the first mutable array into the second mutable array.
--- This function is safe, it copies fewer elements if the second
--- array is smaller than the first.
-copyInto :: Array# a #-> Array# a #-> (# Array# a, Array# a #)
-copyInto = Unsafe.toLinear2 go
+-- | Copy the first mutable array into the second mutable array, starting
+-- from the given index of the source array.
+--
+-- It copies fewer elements if the second array is smaller than the
+-- first. 'n' should be within [0..size src).
+--
+-- @
+--  copyInto n src dest:
+--   dest[i] = src[n+i] for i < size dest, i < size src + n
+-- @
+copyInto :: Int -> Array# a #-> Array# a #-> (# Array# a, Array# a #)
+copyInto start@(GHC.I# start#) = Unsafe.toLinear2 go
   where
     go :: Array# a -> Array# a -> (# Array# a, Array# a #)
     go (Array# src) (Array# dst) =
       let !(GHC.I# len#) = Prelude.min
-            (GHC.I# (GHC.sizeofMutableArray# src))
+            (GHC.I# (GHC.sizeofMutableArray# src) Prelude.- start)
             (GHC.I# (GHC.sizeofMutableArray# dst))
-      in  case GHC.runRW# (GHC.copyMutableArray# src 0# dst 0# len#) of
+      in  case GHC.runRW# (GHC.copyMutableArray# src start# dst 0# len#) of
             _ -> (# Array# src, Array# dst #)
 {-# NOINLINE copyInto #-}  -- prevents the runRW# effect from being reordered
 

src/Data/Vector/Mutable/Linear.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LinearTypes #-}
 {-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StrictData #-}
@@ -11,7 +12,8 @@
 -- | Mutable vectors with a linear API.
 --
 -- Vectors are arrays that grow automatically, that you can append to with
--- 'snoc'.
+-- 'push'. They never shrink automatically to reduce unnecessary copying,
+-- use 'shrinkToFit' to get rid of the wasted space.
 --
 -- To use mutable vectors, create a linear computation of type
 -- @Vector a #-> Ur b@ and feed it to 'constant' or 'fromList'.
@@ -50,11 +52,15 @@ module Data.Vector.Mutable.Linear
     -- * Mutators
     write,
     unsafeWrite,
-    snoc,
+    push,
+    pop,
+    slice,
+    shrinkToFit,
     -- * Accessors
     read,
     unsafeRead,
     size,
+    capacity,
     toList,
   )
 where
@@ -64,6 +70,15 @@ import Prelude.Linear hiding (read)
 import Data.Array.Mutable.Linear (Array)
 import qualified Data.Array.Mutable.Linear as Array
 
+-- # Constants
+-------------------------------------------------------------------------------
+
+-- | When growing the vector, capacity will be multiplied by this number.
+--
+-- This is usually chosen between 1.5 and 2; 2 being the most common.
+constGrowthFactor :: Int
+constGrowthFactor = 2
+
 -- # Core data types
 -------------------------------------------------------------------------------
 
@@ -81,6 +96,8 @@ data Vector a where
 
 -- | Create a 'Vector' from an 'Array'. Result will have the size and capacity
 -- equal to the size of the given array.
+--
+-- This is a constant time operation.
 fromArray :: HasCallStack => Array a #-> Vector a
 fromArray arr =
   Array.size arr
@@ -104,24 +121,44 @@ constant size' x f
 fromList :: HasCallStack => [a] -> (Vector a #-> Ur b) #-> Ur b
 fromList xs f = Array.fromList xs (f . fromArray)
 
--- | Number of elements inside the vector
+-- | Number of elements inside the vector.
+--
+-- This might be different than how much actual memory the vector is using.
+-- For that, see: 'capacity'.
 size :: Vector a #-> (Vector a, Ur Int)
 size (Vec size' arr) = (Vec size' arr, Ur size')
 
--- | Insert at the end of the vector
-snoc :: HasCallStack => Vector a #-> a -> Vector a
-snoc (Vec size' arr) x =
-  Array.size arr & \(arr', Ur cap) ->
-    if size' < cap
-    then write (Vec (size' + 1) arr') size' x
-    else write (unsafeResize ((max size' 1) * 2) (Vec (size' + 1) arr')) size' x
+-- | Capacity of a vector. In other words, the number of elements
+-- the vector can contain before it is copied to a bigger array.
+capacity :: Vector a #-> (Vector a, Ur Int)
+capacity (Vec s arr) =
+  Array.size arr & \(arr', cap) -> (Vec s arr', cap)
+
+-- | Insert at the end of the vector. This will grow the vector if there
+-- is no empty space.
+push :: Vector a #-> a -> Vector a
+push vec x =
+  growToFit 1 vec & \(Vec s arr) ->
+    write (Vec (s + 1) arr) s x
+
+-- | Pop from the end of the vector. This will never shrink the vector, use
+-- 'shrinkToFit' to remove the wasted space.
+pop :: Vector a #-> (Vector a, Ur (Maybe a))
+pop vec =
+  size vec & \case
+    (vec', Ur 0) ->
+      (vec', Ur Nothing)
+    (vec', Ur s) ->
+      read vec' (s-1) & \(Vec _ arr, Ur a) ->
+        ( Vec (s-1) arr
+        , Ur (Just a)
+        )
 
 -- | Write to an element . Note: this will not write to elements beyond the
 -- current size of the vector and will error instead.
 write :: HasCallStack => Vector a #-> Int -> a -> Vector a
-write (Vec size' arr) ix val
-  | indexInRange size' ix = Vec size' (Array.unsafeWrite arr ix val)
-  | otherwise = arr `lseq` error "Write index not in range."
+write vec ix val =
+  unsafeWrite (assertIndexInRange ix vec) ix val
 
 -- | Same as 'write', but does not do bounds-checking. The behaviour is undefined
 -- when passed an invalid index.
@@ -132,11 +169,8 @@ unsafeWrite (Vec size' arr) ix val =
 -- | Read from a vector, with an in-range index and error for an index that is
 -- out of range (with the usual range @0..size-1@).
 read :: HasCallStack => Vector a #-> Int -> (Vector a, Ur a)
-read (Vec size' arr) ix
-  | indexInRange size' ix =
-      Array.unsafeRead arr ix
-        & \(arr', val) -> (Vec size' arr', val)
-  | otherwise = arr `lseq` error "Read index not in range."
+read vec ix =
+  unsafeRead (assertIndexInRange ix vec) ix
 
 -- | Same as 'read', but does not do bounds-checking. The behaviour is undefined
 -- when passed an invalid index.
@@ -153,6 +187,32 @@ toList (Vec s arr) =
   Array.toList arr & \(Ur xs) ->
     Ur (take s xs)
 
+-- | Resize the vector to not have any wasted memory (size == capacity). This
+-- returns a semantically identical vector.
+shrinkToFit :: Vector a #-> Vector a
+shrinkToFit vec =
+  capacity vec & \(vec', Ur cap) ->
+    size vec' & \(vec'', Ur s') ->
+      if cap > s'
+      then unsafeResize s' vec''
+      else vec''
+
+-- | Return a slice of the vector with given size, starting from an offset.
+--
+-- Start offset + target size should be within the input vector, and both should
+-- be non-negative.
+--
+-- This is a constant time operation if the start offset is 0. Use 'shrinkToFit'
+-- to remove the possible wasted space if necessary.
+slice :: Int -> Int -> Vector a #-> Vector a
+slice from newSize (Vec oldSize arr) =
+  if oldSize < from + newSize
+  then arr `lseq` error "Slice index out of bounds"
+  else if from == 0
+       then Vec newSize arr
+       else Array.slice from newSize arr & \(oldArr, newArr) ->
+              oldArr `lseq` fromArray newArr
+
 -- # Instances
 -------------------------------------------------------------------------------
 
@@ -162,6 +222,29 @@ instance Consumable (Vector a) where
 -- # Internal library
 -------------------------------------------------------------------------------
 
+-- | Grows the vector to the closest power of growthFactor to
+-- fit at least n more elements.
+growToFit :: HasCallStack => Int -> Vector a #-> Vector a
+growToFit n vec =
+  capacity vec & \(vec', Ur cap) ->
+    size vec' & \(vec'', Ur s') ->
+      if s' + n <= cap
+      then vec''
+      else
+        let -- Calculate the closest power of growth factor
+            -- larger than required size.
+            newSize =
+              constGrowthFactor -- This constant is defined above.
+                ^ (ceiling :: Double -> Int)
+                    (logBase
+                      (fromIntegral constGrowthFactor)
+                      (fromIntegral (s' + n))) -- this is always
+                                               -- > 0 because of
+                                               -- the if condition
+        in  unsafeResize
+              newSize
+              vec''
+
 -- | Resize the vector to a non-negative size. In-range elements are preserved,
 -- the possible new elements are bottoms.
 unsafeResize :: HasCallStack => Int -> Vector a #-> Vector a
@@ -174,6 +257,10 @@ unsafeResize newSize (Vec size' ma) =
       ma
     )
 
--- | Argument order: indexInRange size ix
-indexInRange :: Int -> Int -> Bool
-indexInRange size' ix = 0 <= ix && ix < size'
+-- | Check if given index is within the Vector, otherwise panic.
+assertIndexInRange :: HasCallStack => Int -> Vector a #-> Vector a
+assertIndexInRange i vec =
+  size vec & \(vec', Ur s) ->
+    if 0 <= i && i < s
+    then vec'
+    else vec' `lseq` error "Vector: index out of bounds"