Use functions from newer primitive and primitive-unlifted

src/Data/Primitive/Contiguous/Class.hs

@@ -148,16 +148,6 @@ class Contiguous (arr :: Type -> Type) where
     -- | new length
     Int ->
     m (Mutable arr (PrimState m) a)
-  default shrink ::
-    ( ContiguousU arr
-    , PrimMonad m
-    , Element arr a
-    ) =>
-    Mutable arr (PrimState m) a ->
-    Int ->
-    m (Mutable arr (PrimState m) a)
-  {-# INLINE shrink #-}
-  shrink = resize
 
   -- | The empty array.
   empty :: arr a
@@ -380,26 +370,13 @@ class Contiguous (arr :: Type -> Type) where
     (PrimMonad m, Element arr b) =>
     Mutable arr (PrimState m) b ->
     m (arr b)
-  unsafeFreeze xs = unsafeShrinkAndFreeze xs =<< sizeMut xs
-  {-# INLINE unsafeFreeze #-}
 
   unsafeShrinkAndFreeze ::
     (PrimMonad m, Element arr a) =>
     Mutable arr (PrimState m) a ->
     -- | final size
     Int ->
     m (arr a)
-  default unsafeShrinkAndFreeze ::
-    ( ContiguousU arr
-    , PrimMonad m
-    , Element arr a
-    ) =>
-    Mutable arr (PrimState m) a ->
-    Int ->
-    m (arr a)
-  {-# INLINE unsafeShrinkAndFreeze #-}
-  unsafeShrinkAndFreeze arr0 len' =
-    resize arr0 len' >>= unsafeFreeze
 
   -- | Copy a slice of an immutable array into a new mutable array.
   thaw ::
@@ -572,9 +549,10 @@ class (Contiguous arr) => ContiguousU arr where
   -- | 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
+  -- uninitialized memory at these indices. But for @SmallMutableArray@ or @Array@, these
   -- are set to an error thunk, so reading from them and forcing the result
-  -- causes the program to crash.
+  -- causes the program to crash. For @UnliftedArray@, the new elements have undefined values of an unknown type.
+  -- If the array is not grown, it may (or may not) be modified in place.
   resize ::
     (PrimMonad m, Element arr b) =>
     Mutable arr (PrimState m) b ->
@@ -623,6 +601,10 @@ instance (ContiguousU arr) => Contiguous (Slice arr) where
   replicateMut len x = do
     baseMut <- replicateMut len x
     pure MutableSlice {offsetMut = 0, lengthMut = len, baseMut = unliftMut baseMut}
+  {-# INLINE unsafeFreeze #-}
+  unsafeFreeze (MutableSlice off len base) = do
+    base' <- unsafeFreeze (liftMut base)
+    pure (Slice off len (unlift base'))
   {-# INLINE shrink #-}
   shrink xs len' = pure $ case compare len' (lengthMut xs) of
     LT -> xs {lengthMut = len'}
@@ -884,6 +866,9 @@ instance Contiguous SmallArray where
   shrink !arr !n = do
     shrinkSmallMutableArray arr n
     pure arr
+  unsafeShrinkAndFreeze !arr !n = do
+    shrinkSmallMutableArray arr n
+    unsafeFreezeSmallArray arr
 
 instance ContiguousU SmallArray where
   type Unlifted SmallArray = SmallArray#
@@ -1012,6 +997,14 @@ instance Contiguous PrimArray where
     unsafeFreeze dst
   {-# INLINE run #-}
   run = runPrimArrayST
+  {-# INLINE shrink #-}
+  shrink !arr !n = do
+    shrinkMutablePrimArray arr n
+    pure arr
+  {-# INLINE unsafeShrinkAndFreeze #-}
+  unsafeShrinkAndFreeze !arr !n = do
+    shrinkMutablePrimArray arr n
+    unsafeFreezePrimArray arr
 
 newtype PrimArray# a = PrimArray# ByteArray#
 newtype MutablePrimArray# s a = MutablePrimArray# (MutableByteArray# s)
@@ -1136,6 +1129,14 @@ instance Contiguous Array where
     unsafeFreezeArray m
   {-# INLINE run #-}
   run = runArrayST
+  {-# INLINE shrink #-}
+  shrink !arr !n = do
+    -- See Note [Shrinking Arrays Without a Shrink Primop]
+    cloneMutableArray arr 0 n
+  {-# INLINE unsafeShrinkAndFreeze #-}
+  unsafeShrinkAndFreeze !arr !n =
+    -- See Note [Shrinking Arrays Without a Shrink Primop]
+    freezeArray arr 0 n
 
 instance ContiguousU Array where
   type Unlifted Array = Array#
@@ -1261,18 +1262,20 @@ instance Contiguous (UnliftedArray_ unlifted_a) where
   run = runUnliftedArrayST
   {-# INLINE shrink #-}
   shrink !arr !n = do
-    -- See Note [Shrinking Unlifted Arrays]
+    -- See Note [Shrinking Arrays Without a Shrink Primop]
     cloneMutableUnliftedArray arr 0 n
   {-# INLINE unsafeShrinkAndFreeze #-}
   unsafeShrinkAndFreeze !arr !n =
-    -- See Note [Shrinking Unlifted Arrays]
+    -- See Note [Shrinking Arrays Without a Shrink Primop]
     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.
+-- Note [Shrinking Arrays Without a Shrink Primop]
+-- ===============================================
+-- GHC's Array# type has a card table and cannot currently be shrunk in place.
+-- (SmallArray#, however, can be shrunk in place). These implementations copy
+-- the array rather than freezing it in place. But at least they are 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)

src/Data/Primitive/Contiguous/Shim.hs

@@ -8,7 +8,6 @@ module Data.Primitive.Contiguous.Shim
   , replicateMutablePrimArray
   ) where
 
-import Control.Monad (when)
 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, (<$))
@@ -22,20 +21,25 @@ errorThunk = error "Contiguous typeclass: unitialized element"
 
 resizeArray :: (PrimMonad m) => MutableArray (PrimState m) a -> Int -> m (MutableArray (PrimState m) a)
 resizeArray !src !sz = do
-  dst <- newArray sz errorThunk
-  copyMutableArray dst 0 src 0 (min sz (sizeofMutableArray src))
-  pure dst
+  let !srcSz = sizeofMutableArray src
+  case compare sz srcSz of
+    EQ -> pure src
+    LT -> cloneMutableArray src 0 sz
+    GT -> do
+      dst <- newArray sz errorThunk
+      copyMutableArray dst 0 src 0 sz
+      pure dst
 {-# INLINE resizeArray #-}
 
 resizeUnliftedArray :: (PrimMonad m, PrimUnlifted a) => MutableUnliftedArray (PrimState m) a -> Int -> m (MutableUnliftedArray (PrimState m) a)
-resizeUnliftedArray !src !sz =
-  let !srcSz = sizeofMutableUnliftedArray src in
+resizeUnliftedArray !src !sz = do
+  let !srcSz = sizeofMutableUnliftedArray src
   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))
+      copyMutableUnliftedArray dst 0 src 0 sz
       pure dst
 {-# INLINE resizeUnliftedArray #-}