@@ -42,6 +42,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
4242 -- * Mutation
4343 , putIndexUnsafe
4444 , snocUnsafe
45+ , snocWith
4546 , snoc
4647 , snocLinear
4748
@@ -57,6 +58,12 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
5758 , cmp
5859
5960 -- * Composable Folds
61+ -- XXX These should be Refolds
62+ , appendNUnsafe
63+ , appendN
64+ , appendWith
65+ , append
66+
6067 , writeWith
6168 , writeNAllocWith
6269 , writeN
@@ -122,6 +129,11 @@ import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
122129import Foreign.Ptr (plusPtr, minusPtr, castPtr, nullPtr)
123130import Foreign.Storable (Storable(..))
124131import GHC.Base (build)
132+#if __GLASGOW_HASKELL__ < 802
133+#define noinline
134+#else
135+import GHC.Base (noinline)
136+#endif
125137import GHC.ForeignPtr (ForeignPtr(..))
126138import GHC.Ptr (Ptr(..))
127139
@@ -155,8 +167,12 @@ plusForeignPtr (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c
155167#endif
156168
157169-- $setup
158−-- >>> import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
170+-- >>> :m
159171-- >>> import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as Array
172+-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream
173+-- >>> import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
174+-- >>> import qualified Streamly.Internal.Data.Fold as Fold
175+-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold
160176
161177-------------------------------------------------------------------------------
162178-- Array Data Type
@@ -490,7 +506,7 @@ snocLinear = snocWith (+ arrayChunkSize)
490506--
491507-- Note that the returned array may be a mutated version of the original array.
492508--
493−-- >>> snoc = snocWith (* 2)
509+-- >>> snoc = Array.snocWith (* 2)
494510--
495511-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
496512--
@@ -974,6 +990,95 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
974990-- Folds
975991-------------------------------------------------------------------------------
976992
993+data ArrayUnsafe a = ArrayUnsafe
994+ {-# UNPACK #-} !(ForeignPtr a) -- first address
995+ {-# UNPACK #-} !(Ptr a) -- first unused address
996+
997+-- Note: Arrays may be allocated with a specific alignment at the beginning of
998+-- the array. If you need to maintain that alignment on reallocations then you
999+-- can resize the array manually before append, using an aligned resize
1000+-- operation.
1001+
1002+-- XXX Keep the bound intact to not lose any free space? Perf impact?
1003+
1004+-- | Append up to @n@ input items to the supplied array.
1005+--
1006+-- Unsafe: Do not drive the fold beyond @n@ elements, it will lead to memory
1007+-- corruption or segfault.
1008+--
1009+-- Any free space left in the array after appending @n@ elements is lost.
1010+--
1011+-- /Internal/
1012+{-# INLINE_NORMAL appendNUnsafe #-}
1013+appendNUnsafe :: forall m a. (MonadIO m, Storable a) =>
1014+ m (Array a)
1015+ -> Int
1016+ -> Fold m a (Array a)
1017+appendNUnsafe action n =
1018+ fmap extract $ FL.foldlM' step initial
1019+
1020+ where
1021+
1022+ initial = do
1023+ assert (n >= 0) (return ())
1024+ arr@(Array _ end bound) <- action
1025+ let free = bound `minusPtr` end
1026+ elemSize = sizeOf (undefined :: a)
1027+ needed = n * elemSize
1028+ -- XXX We can also reallocate if the array has too much free space,
1029+ -- otherwise we lose that space.
1030+ (Array start1 end1 _) <-
1031+ if free < needed
1032+ then noinline reallocWith "appendNUnsafeWith" (+ needed) needed arr
1033+ else return arr
1034+ return $ ArrayUnsafe start1 end1
1035+
1036+ step (ArrayUnsafe start end) x = do
1037+ liftIO $ poke end x
1038+ return $ ArrayUnsafe start (end `plusPtr` sizeOf (undefined :: a))
1039+
1040+ extract (ArrayUnsafe start end) = Array start end end
1041+
1042+-- | Append @n@ elements to an existing array. Any free space left in the array
1043+-- after appending @n@ elements is lost.
1044+--
1045+-- >>> appendN initial n = Fold.take n (Array.appendNUnsafe initial n)
1046+--
1047+-- /Pre-release/
1048+{-# INLINE_NORMAL appendN #-}
1049+appendN :: forall m a. (MonadIO m, Storable a) =>
1050+ m (Array a) -> Int -> Fold m a (Array a)
1051+appendN initial n = FL.take n (appendNUnsafe initial n)
1052+
1053+-- | @appendWith realloc action@ mutates the array generated by @action@ to
1054+-- append the input stream. If there is no reserved space available in the
1055+-- array it is reallocated to a size in bytes determined by @realloc oldSize@,
1056+-- where @oldSize@ is the current size of the array in bytes.
1057+--
1058+-- Note that the returned array may be a mutated version of original array.
1059+--
1060+-- >>> appendWith sizer = Fold.foldlM' (Array.snocWith sizer)
1061+--
1062+-- /Pre-release/
1063+{-# INLINE appendWith #-}
1064+appendWith :: forall m a. (MonadIO m, Storable a) =>
1065+ (Int -> Int) -> m (Array a) -> Fold m a (Array a)
1066+appendWith sizer = FL.foldlM' (snocWith sizer)
1067+
1068+-- | @append action@ mutates the array generated by @action@ to append the
1069+-- input stream. If there is no reserved space available in the array it is
1070+-- reallocated to double the size.
1071+--
1072+-- Note that the returned array may be a mutated version of original array.
1073+--
1074+-- >>> append = Array.appendWith (* 2)
1075+--
1076+-- /Pre-release/
1077+{-# INLINE append #-}
1078+append :: forall m a. (MonadIO m, Storable a) =>
1079+ m (Array a) -> Fold m a (Array a)
1080+append = appendWith (* 2)
1081+
9771082{-# INLINE_NORMAL writeNAllocWith #-}
9781083writeNAllocWith :: forall m a. (MonadIO m, Storable a)
9791084 => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
@@ -982,8 +1087,8 @@ writeNAllocWith alloc n = FL.take n (writeNUnsafeWith alloc n)
9821087-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
9831088-- 'Array'.
9841089--
985−-- >>> import qualified Streamly.Data.Fold as Fold
986−-- >>> writeN n = Fold.take n (writeNUnsafe n)
1090+-- >>> writeN n = Fold.take n (Array.writeNUnsafe n)
1091+-- >>> writeN n = Array.appendN (Array.newArray n) n
9871092--
9881093-- @since 0.7.0
9891094{-# INLINE_NORMAL writeN #-}
@@ -1013,10 +1118,6 @@ writeNAlignedUnmanaged :: forall m a. (MonadIO m, Storable a)
10131118writeNAlignedUnmanaged alignSize =
10141119 writeNAllocWith (newArrayAlignedUnmanaged alignSize)
10151120
1016−data ArrayUnsafe a = ArrayUnsafe
1017− {-# UNPACK #-} !(ForeignPtr a) -- first address
1018− {-# UNPACK #-} !(Ptr a) -- first unused address
1019−
10201121-- XXX We can carry bound as well in the state to make sure we do not lose the
10211122-- remaining capacity. Need to check perf impact.
10221123--
@@ -1058,7 +1159,7 @@ writeNUnsafe = writeNUnsafeWith newArray
10581159--
10591160-- | Buffer a stream into a stream of arrays.
10601161--
1061−-- >>> writeChunks n = FL.many (writeN n) FL.toStreamK
1162+-- >>> writeChunks n = Fold.many (Array.writeN n) Fold.toStreamK
10621163--
10631164-- Breaking an array into an array stream can be useful to consume a large
10641165-- array sequentially such that memory of the array is released incrementatlly.
@@ -1150,7 +1251,7 @@ writeAligned alignSize =
11501251
11511252-- | Use the 'writeN' fold instead.
11521253--
1153−-- >>> fromStreamDN n = StreamD.fold (writeN n)
1254+-- >>> fromStreamDN n = StreamD.fold (Array.writeN n)
11541255--
11551256{-# INLINE_NORMAL fromStreamDN #-}
11561257fromStreamDN :: forall m a. (MonadIO m, Storable a)
@@ -1205,9 +1306,7 @@ fromArrayStreamK as = do
12051306-- the chunking approach. However, if we have to shrink in the end then it may
12061307-- result in an extra copy of the entire data.
12071308--
1208−-- @
1209−-- fromStreamD = StreamD.fold Array.write
1210−-- @
1309+-- >>> fromStreamD = StreamD.fold Array.write
12111310--
12121311{-# INLINE fromStreamD #-}
12131312fromStreamD :: (MonadIO m, Storable a) => D.Stream m a -> m (Array a)