core/src/Streamly/Internal/Data/Array/Generic.hs · streamly

git.iko.soy/streamly

on fa9fa09 help sign in

streamly/core/src/Streamly/Internal/Data/Array/Generic.hs

326 lines 10 kB history raw

1-- |

2-- Module : Streamly.Internal.Data.Array.Generic

4--

5-- License : BSD-3-Clause

6-- Maintainer : streamly@composewell.com

7-- Portability : GHC

8--

9module Streamly.Internal.Data.Array.Generic

10 ( Array(..)

12 -- * Construction

13 , nil

14 , writeN

15 , write

16 , writeWith

17 , writeLastN

19 , fromStreamN

20 , fromStream

21 , fromPureStream

22 , fromByteStr#

24 , fromListN

25 , fromList

27 , chunksOf

29 -- * Elimination

30 , length

31 , reader

33 , toList

34 , read

35 , readRev

37 , foldl'

38 , foldr

39 , streamFold

40 , fold

42 -- * Random Access

43 , getIndexUnsafe

44 , getIndex

45 , getSliceUnsafe

46 , strip

47 )

48where

50#include "inline.hs"

52import Control.Monad (replicateM)

53import Control.Monad.IO.Class (MonadIO)

54import Data.Functor.Identity (Identity(..))

55import Data.Word (Word8)

56import GHC.Base (MutableArray#, RealWorld)

57import GHC.Exts (Addr#)

58import GHC.IO (unsafePerformIO)

59import Text.Read (readPrec)

61import Streamly.Internal.Data.Fold.Type (Fold(..))

62import Streamly.Internal.Data.Stream.Type (Stream)

63import Streamly.Internal.Data.Unfold.Type (Unfold(..))

64import Streamly.Internal.System.IO (unsafeInlineIO)

66import qualified Streamly.Internal.Data.MutArray.Generic as MArray

67import qualified Streamly.Internal.Data.Fold.Type as FL

68import qualified Streamly.Internal.Data.Producer.Type as Producer

69import qualified Streamly.Internal.Data.Producer as Producer

70import qualified Streamly.Internal.Data.Ring.Generic as RB

71import qualified Streamly.Internal.Data.Stream.Type as D

72import qualified Streamly.Internal.Data.Stream.Generate as D

73import qualified Text.ParserCombinators.ReadPrec as ReadPrec

75import Prelude hiding (Foldable(..), read)

77-------------------------------------------------------------------------------

78-- Array Data Type

79-------------------------------------------------------------------------------

81data Array a =

82 Array

83 { arrContents# :: MutableArray# RealWorld a

84 -- ^ The internal contents of the array representing the entire array.

86 , arrStart :: {-# UNPACK #-}!Int

87 -- ^ The starting index of this slice.

89 , arrLen :: {-# UNPACK #-}!Int

90 -- ^ The length of this slice.

91 }

93unsafeFreeze :: MArray.MutArray a -> Array a

94unsafeFreeze (MArray.MutArray cont# arrS arrL _) = Array cont# arrS arrL

96unsafeThaw :: Array a -> MArray.MutArray a

97unsafeThaw (Array cont# arrS arrL) = MArray.MutArray cont# arrS arrL arrL

99{-# NOINLINE nil #-}

100nil :: Array a

101nil = unsafePerformIO $ unsafeFreeze <$> MArray.nil

102

103-------------------------------------------------------------------------------

104-- Construction - Folds

105-------------------------------------------------------------------------------

106

107{-# INLINE_NORMAL writeN #-}

108writeN :: MonadIO m => Int -> Fold m a (Array a)

109writeN = fmap unsafeFreeze <$> MArray.writeN

110

111{-# INLINE_NORMAL writeWith #-}

112writeWith :: MonadIO m => Int -> Fold m a (Array a)

113writeWith elemCount = unsafeFreeze <$> MArray.writeWith elemCount

114

115-- | Fold the whole input to a single array.

116--

117-- /Caution! Do not use this on infinite streams./

118--

119{-# INLINE write #-}

120write :: MonadIO m => Fold m a (Array a)

121write = fmap unsafeFreeze MArray.write

122

123fromPureStream :: Stream Identity a -> Array a

124fromPureStream x =

125 unsafePerformIO $ fmap (unsafeFreeze) (MArray.fromPureStream x)

126-- fromPureStream = runIdentity . D.fold (unsafeMakePure write)

127-- fromPureStream = fromList . runIdentity . D.toList

128

129fromByteStr# :: Addr# -> Array Word8

130fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

131

132-------------------------------------------------------------------------------

133-- Stream Ops

134-------------------------------------------------------------------------------

135

136{-# INLINE_NORMAL chunksOf #-}

137chunksOf :: forall m a. MonadIO m

138 => Int -> Stream m a -> Stream m (Array a)

139chunksOf n strm = fmap unsafeFreeze $ MArray.chunksOf n strm

140

141-------------------------------------------------------------------------------

142-- Construction - from streams

143-------------------------------------------------------------------------------

144

145{-# INLINE fromStreamN #-}

146fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)

147fromStreamN n = D.fold (writeN n)

148

149{-# INLINE fromStream #-}

150fromStream :: MonadIO m => Stream m a -> m (Array a)

151fromStream = D.fold write

152

153-- XXX Consider foldr/build fusion in toList/fromList

154

155{-# INLINABLE fromListN #-}

156fromListN :: Int -> [a] -> Array a

157fromListN n xs = unsafePerformIO $ fromStreamN n $ D.fromList xs

158

159{-# INLINABLE fromList #-}

160fromList :: [a] -> Array a

161fromList xs = unsafePerformIO $ fromStream $ D.fromList xs

162

163-------------------------------------------------------------------------------

164-- Elimination - Unfolds

165-------------------------------------------------------------------------------

166

167{-# INLINE length #-}

168length :: Array a -> Int

169length = arrLen

170

171{-# INLINE_NORMAL reader #-}

172reader :: Monad m => Unfold m (Array a) a

173reader =

174 Producer.simplify

175 $ Producer.translate unsafeThaw unsafeFreeze

176 $ MArray.producerWith (return . unsafeInlineIO)

177

178-------------------------------------------------------------------------------

179-- Elimination - to streams

180-------------------------------------------------------------------------------

181

182{-# INLINE_NORMAL toList #-}

183toList :: Array a -> [a]

184toList arr = loop 0

185

186 where

187

188 len = length arr

189 loop i | i == len = []

190 loop i = getIndexUnsafe i arr : loop (i + 1)

191

192{-# INLINE_NORMAL read #-}

193read :: Monad m => Array a -> Stream m a

194read arr@Array{..} =

195 D.map (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)

196

197{-# INLINE_NORMAL readRev #-}

198readRev :: Monad m => Array a -> Stream m a

199readRev arr@Array{..} =

200 D.map (`getIndexUnsafe` arr)

201 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0

202

203-------------------------------------------------------------------------------

204-- Elimination - using Folds

205-------------------------------------------------------------------------------

206

207{-# INLINE_NORMAL foldl' #-}

208foldl' :: (b -> a -> b) -> b -> Array a -> b

209foldl' f z arr = unsafePerformIO $ D.foldl' f z $ read arr

210

211{-# INLINE_NORMAL foldr #-}

212foldr :: (a -> b -> b) -> b -> Array a -> b

213foldr f z arr = unsafePerformIO $ D.foldr f z $ read arr

214

215{-# INLINE fold #-}

216fold :: Monad m => Fold m a b -> Array a -> m b

217fold f arr = D.fold f (read arr)

218

219{-# INLINE streamFold #-}

220streamFold :: Monad m => (Stream m a -> m b) -> Array a -> m b

221streamFold f arr = f (read arr)

222

223-------------------------------------------------------------------------------

224-- Random reads and writes

225-------------------------------------------------------------------------------

226

227-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does

228-- not check the bounds.

229--

230-- @since 0.8.0

231{-# INLINE getIndexUnsafe #-}

232getIndexUnsafe :: Int -> Array a -> a

233getIndexUnsafe i arr =

234 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)

235

236-- | Lookup the element at the given index. Index starts from 0.

237--

238{-# INLINE getIndex #-}

239getIndex :: Int -> Array a -> Maybe a

240getIndex i arr@Array {..} =

241 if i >= 0 && i < arrLen

242 then Just $ getIndexUnsafe i arr

243 else Nothing

244

245-- >>> import qualified Streamly.Data.Stream as Stream

246-- >>> import qualified Streamly.Data.Fold as Fold

247-- >>> import qualified Streamly.Internal.Data.Array.Generic as Array

248-- >>> import Data.Function ((&))

249-- >>> :{

250-- Stream.fromList [1,2,3,4,5::Int]

251-- & Stream.scan (Array.writeLastN 2)

252-- & Stream.fold Fold.toList

253-- :}

254-- [fromList [],fromList [1],fromList [1,2],fromList [2,3],fromList [3,4],fromList [4,5]]

255--

256{-# INLINE writeLastN #-}

257writeLastN :: MonadIO m => Int -> Fold m a (Array a)

258writeLastN n = FL.rmapM f (RB.writeLastN n)

259

260 where

261

262 f rb = do

263 arr <- RB.copyToMutArray 0 n rb

264 return $ unsafeFreeze arr

265

266{-# INLINE getSliceUnsafe #-}

267getSliceUnsafe :: Int -> Int -> Array a -> Array a

268getSliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len

269

270-- XXX This is not efficient as it copies the array. We should support array

271-- slicing so that we can just refer to the underlying array memory instead of

272-- copying.

273

274-- | Truncate the array at the beginning and end as long as the predicate

275-- holds true. Returns a slice of the original array.

276{-# INLINE strip #-}

277strip :: (a -> Bool) -> Array a -> Array a

278strip p arr = unsafeFreeze $ unsafePerformIO $ MArray.strip p (unsafeThaw arr)

279

280-------------------------------------------------------------------------------

281-- Instances

282-------------------------------------------------------------------------------

283

284instance Eq a => Eq (Array a) where

285 {-# INLINE (==) #-}

286 arr1 == arr2 =

287 unsafeInlineIO $! unsafeThaw arr1 `MArray.eq` unsafeThaw arr2

288

289instance Ord a => Ord (Array a) where

290 {-# INLINE compare #-}

291 compare arr1 arr2 =

292 unsafeInlineIO $! unsafeThaw arr1 `MArray.cmp` unsafeThaw arr2

293

294 -- Default definitions defined in base do not have an INLINE on them, so we

295 -- replicate them here with an INLINE.

296 {-# INLINE (<) #-}

297 x < y = case compare x y of { LT -> True; _ -> False }

298

299 {-# INLINE (<=) #-}

300 x <= y = case compare x y of { GT -> False; _ -> True }

301

302 {-# INLINE (>) #-}

303 x > y = case compare x y of { GT -> True; _ -> False }

304

305 {-# INLINE (>=) #-}

306 x >= y = case compare x y of { LT -> False; _ -> True }

307

308 -- These two default methods use '<=' rather than 'compare'

309 -- because the latter is often more expensive

310 {-# INLINE max #-}

311 max x y = if x <= y then y else x

312

313 {-# INLINE min #-}

314 min x y = if x <= y then x else y

315

316instance Show a => Show (Array a) where

317 {-# INLINE show #-}

318 show arr = "fromList " ++ show (toList arr)

319

320instance Read a => Read (Array a) where

321 {-# INLINE readPrec #-}

322 readPrec = do

323 fromListWord <- replicateM 9 ReadPrec.get

324 if fromListWord == "fromList "

325 then fromList <$> readPrec

326 else ReadPrec.pfail