Use functions from newer primitive and primitive-unlifted

The implementation UnliftedArray in primitive-unlifted-2.1
penalizes the creation of an uninitialized unlifted array. When
shrinking and resizing unlifted arrays, there are primitives that
we can use to avoid this.

Also, primitive itself now has shims for common operations on
PrimArray, so this commit also cleans up the Shim module.

contiguous.cabal

@@ -38,7 +38,7 @@ library
   build-depends:
     , base                >=4.14    && <5
     , deepseq             >=1.4
-    , primitive           >=0.7.2   && <0.10
+    , primitive           >=0.9     && <0.10
     , primitive-unlifted  >=2.1
     , run-st              >=0.1.3.2
 

src/Data/Primitive/Contiguous.hs

@@ -273,7 +273,7 @@ module Data.Primitive.Contiguous
   ) where
 
 import Control.Monad.Primitive
-import Data.Primitive hiding (fromList, fromListN)
+import Data.Primitive
 import Data.Primitive.Unlifted.Array
 import Prelude hiding (Foldable (..), all, any, filter, map, mapM, mapM_, read, replicate, reverse, scanl, sequence, sequence_, traverse, zip, zipWith, (<$))
 

src/Data/Primitive/Contiguous/Class.hs

@@ -27,7 +27,7 @@ module Data.Primitive.Contiguous.Class
   , Always
   ) where
 
-import Data.Primitive hiding (fromList, fromListN)
+import Data.Primitive
 import Data.Primitive.Contiguous.Shim
 import Data.Primitive.Unlifted.Array
 import Prelude hiding
@@ -569,7 +569,12 @@ class (Contiguous arr) => ContiguousU arr where
   -- | The unifted version of the mutable array type (i.e. eliminates an indirection through a thunk).
   type UnliftedMut arr = (r :: Type -> Type -> TYPE UnliftedRep) | r -> arr
 
-  -- | Resize an array into one with the given size.
+  -- | Resize an array into one with the given size. If the array is grown,
+  -- then reading from any newly introduced element before writing to it is undefined behavior.
+  -- The current behavior is that anything backed by @MutableByteArray#@ ends
+  -- uninitialized memory at these indices. But for @SmallMutableArray@, these
+  -- are set to an error thunk, so reading from them and forcing the result
+  -- causes the program to crash.
   resize ::
     (PrimMonad m, Element arr b) =>
     Mutable arr (PrimState m) b ->
@@ -807,7 +812,7 @@ instance Contiguous SmallArray where
   {-# INLINE size #-}
   size = sizeofSmallArray
   {-# INLINE sizeMut #-}
-  sizeMut = (\x -> pure $! sizeofSmallMutableArray x)
+  sizeMut = getSizeofSmallMutableArray
   {-# INLINE thaw_ #-}
   thaw_ = thawSmallArray
   {-# INLINE equals #-}
@@ -872,15 +877,19 @@ instance Contiguous SmallArray where
   {-# INLINE copyMut_ #-}
   copyMut_ = copySmallMutableArray
   {-# INLINE replicateMut #-}
-  replicateMut = replicateSmallMutableArray
+  replicateMut = newSmallArray
   {-# INLINE run #-}
   run = runSmallArrayST
+  {-# INLINE shrink #-}
+  shrink !arr !n = do
+    shrinkSmallMutableArray arr n
+    pure arr
 
 instance ContiguousU SmallArray where
   type Unlifted SmallArray = SmallArray#
   type UnliftedMut SmallArray = SmallMutableArray#
   {-# INLINE resize #-}
-  resize = resizeSmallArray
+  resize !arr !n = resizeSmallMutableArray arr n resizeSmallMutableArrayUninitializedElement
   {-# INLINE unlift #-}
   unlift (SmallArray x) = x
   {-# INLINE unliftMut #-}
@@ -926,7 +935,7 @@ instance Contiguous PrimArray where
     lengthMut <- sizeMut baseMut
     pure MutableSlice {offsetMut = 0, lengthMut, baseMut = unliftMut baseMut}
   {-# INLINE freeze_ #-}
-  freeze_ = freezePrimArrayShim
+  freeze_ = freezePrimArray
   {-# INLINE unsafeFreeze #-}
   unsafeFreeze = unsafeFreezePrimArray
   {-# INLINE thaw_ #-}
@@ -936,9 +945,9 @@ instance Contiguous PrimArray where
   {-# INLINE copyMut_ #-}
   copyMut_ = copyMutablePrimArray
   {-# INLINE clone_ #-}
-  clone_ = clonePrimArrayShim
+  clone_ = clonePrimArray
   {-# INLINE cloneMut_ #-}
-  cloneMut_ = cloneMutablePrimArrayShim
+  cloneMut_ = cloneMutablePrimArray
   {-# INLINE equals #-}
   equals = (==)
   {-# INLINE null #-}
@@ -1250,6 +1259,20 @@ instance Contiguous (UnliftedArray_ unlifted_a) where
     unsafeFreezeUnliftedArray m
   {-# INLINE run #-}
   run = runUnliftedArrayST
+  {-# INLINE shrink #-}
+  shrink !arr !n = do
+    -- See Note [Shrinking Unlifted Arrays]
+    cloneMutableUnliftedArray arr 0 n
+  {-# INLINE unsafeShrinkAndFreeze #-}
+  unsafeShrinkAndFreeze !arr !n =
+    -- See Note [Shrinking Unlifted Arrays]
+    freezeUnliftedArray arr 0 n
+
+-- Note [Shrinking Unlifted Arrays]
+-- ================================
+-- This implementation copies the array rather than freezing it in place.
+-- But at least it is able to avoid assigning all of the elements to
+-- a nonsense value before replacing them with memcpy.
 
 newtype UnliftedArray## (u :: TYPE UnliftedRep) (a :: Type)
   = UnliftedArray## (Exts.Array# u)
@@ -1269,3 +1292,7 @@ instance ContiguousU (UnliftedArray_ unlifted_a) where
   lift (UnliftedArray## x) = UnliftedArray (UnliftedArray# x)
   {-# INLINE liftMut #-}
   liftMut (MutableUnliftedArray## x) = MutableUnliftedArray (MutableUnliftedArray# x)
+
+resizeSmallMutableArrayUninitializedElement :: a
+{-# noinline resizeSmallMutableArrayUninitializedElement #-}
+resizeSmallMutableArrayUninitializedElement = errorWithoutStackTrace "uninitialized element of resizeSmallMutableArray"

src/Data/Primitive/Contiguous/Shim.hs

@@ -4,18 +4,12 @@
 module Data.Primitive.Contiguous.Shim
   ( errorThunk
   , resizeArray
-  , resizeSmallArray
-  , replicateSmallMutableArray
   , resizeUnliftedArray
   , replicateMutablePrimArray
-  , clonePrimArrayShim
-  , cloneMutablePrimArrayShim
-  , freezePrimArrayShim
   ) where
 
 import Control.Monad (when)
-import Control.Monad.ST.Run (runPrimArrayST)
-import Data.Primitive hiding (fromList, fromListN)
+import Data.Primitive
 import Data.Primitive.Unlifted.Array
 import Prelude hiding (all, any, elem, filter, foldMap, foldl, foldr, map, mapM, mapM_, maximum, minimum, null, read, replicate, reverse, scanl, sequence, sequence_, traverse, zip, zipWith, (<$))
 
@@ -33,32 +27,16 @@ resizeArray !src !sz = do
   pure dst
 {-# INLINE resizeArray #-}
 
-resizeSmallArray :: (PrimMonad m) => SmallMutableArray (PrimState m) a -> Int -> m (SmallMutableArray (PrimState m) a)
-resizeSmallArray !src !sz = do
-  dst <- newSmallArray sz errorThunk
-  copySmallMutableArray dst 0 src 0 (min sz (sizeofSmallMutableArray src))
-  pure dst
-{-# INLINE resizeSmallArray #-}
-
-replicateSmallMutableArray ::
-  (PrimMonad m) =>
-  Int ->
-  a ->
-  m (SmallMutableArray (PrimState m) a)
-replicateSmallMutableArray len a = do
-  marr <- newSmallArray len errorThunk
-  let go !ix = when (ix < len) $ do
-        writeSmallArray marr ix a
-        go (ix + 1)
-  go 0
-  pure marr
-{-# INLINE replicateSmallMutableArray #-}
-
 resizeUnliftedArray :: (PrimMonad m, PrimUnlifted a) => MutableUnliftedArray (PrimState m) a -> Int -> m (MutableUnliftedArray (PrimState m) a)
-resizeUnliftedArray !src !sz = do
-  dst <- unsafeNewUnliftedArray sz
-  copyMutableUnliftedArray dst 0 src 0 (min sz (sizeofMutableUnliftedArray src))
-  pure dst
+resizeUnliftedArray !src !sz =
+  let !srcSz = sizeofMutableUnliftedArray src in
+  case compare sz srcSz of
+    EQ -> pure src
+    LT -> cloneMutableUnliftedArray src 0 sz
+    GT -> do
+      dst <- unsafeNewUnliftedArray sz
+      copyMutableUnliftedArray dst 0 src 0 (min sz (sizeofMutableUnliftedArray src))
+      pure dst
 {-# INLINE resizeUnliftedArray #-}
 
 replicateMutablePrimArray ::
@@ -73,24 +51,3 @@ replicateMutablePrimArray len a = do
   setPrimArray marr 0 len a
   pure marr
 {-# INLINE replicateMutablePrimArray #-}
-
-clonePrimArrayShim :: (Prim a) => PrimArray a -> Int -> Int -> PrimArray a
-clonePrimArrayShim !arr !off !len = runPrimArrayST $ do
-  marr <- newPrimArray len
-  copyPrimArray marr 0 arr off len
-  unsafeFreezePrimArray marr
-{-# INLINE clonePrimArrayShim #-}
-
-cloneMutablePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (MutablePrimArray (PrimState m) a)
-cloneMutablePrimArrayShim !arr !off !len = do
-  marr <- newPrimArray len
-  copyMutablePrimArray marr 0 arr off len
-  pure marr
-{-# INLINE cloneMutablePrimArrayShim #-}
-
-freezePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
-freezePrimArrayShim !src !off !len = do
-  dst <- newPrimArray len
-  copyMutablePrimArray dst 0 src off len
-  unsafeFreezePrimArray dst
-{-# INLINE freezePrimArrayShim #-}