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

git.iko.soy/streamly

on 63407dc help sign in

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

304 lines 9.3 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 -- * Elimination

28 , length

29 , reader

31 , toList

32 , read

33 , readRev

35 , foldl'

36 , foldr

37 , streamFold

38 , fold

40 -- * Random Access

41 , getIndexUnsafe

42 , getIndex

43 , getSliceUnsafe

44 , strip

45 )

46where

48#include "inline.hs"

50import Control.Monad (replicateM)

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

52import Data.Functor.Identity (Identity(..))

53import Data.Word (Word8)

54import GHC.Base (MutableArray#, RealWorld)

55import GHC.Exts (Addr#)

56import GHC.IO (unsafePerformIO)

57import Text.Read (readPrec)

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

60import Streamly.Internal.Data.Stream.StreamD.Type (Stream)

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

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

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

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

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

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

68import qualified Streamly.Internal.Data.Ring as RB

69import qualified Streamly.Internal.Data.Stream.StreamD.Type as D

70import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D

71import qualified Text.ParserCombinators.ReadPrec as ReadPrec

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

75-------------------------------------------------------------------------------

76-- Array Data Type

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

79data Array a =

80 Array

81 { arrContents# :: MutableArray# RealWorld a

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

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

85 -- ^ The starting index of this slice.

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

88 -- ^ The length of this slice.

89 }

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

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

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

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

97{-# NOINLINE nil #-}

98nil :: Array a

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

100

101-------------------------------------------------------------------------------

102-- Construction - Folds

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

104

105{-# INLINE_NORMAL writeN #-}

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

107writeN = fmap unsafeFreeze <$> MArray.writeN

108

109{-# INLINE_NORMAL writeWith #-}

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

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

112

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

114--

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

116--

117{-# INLINE write #-}

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

119write = fmap unsafeFreeze MArray.write

120

121fromPureStream :: Stream Identity a -> Array a

122fromPureStream x =

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

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

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

126

127fromByteStr# :: Addr# -> Array Word8

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

129

130-------------------------------------------------------------------------------

131-- Construction - from streams

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

133

134{-# INLINE fromStreamN #-}

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

136fromStreamN n = D.fold (writeN n)

137

138{-# INLINE fromStream #-}

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

140fromStream = D.fold write

141

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

143

144{-# INLINABLE fromListN #-}

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

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

147

148{-# INLINABLE fromList #-}

149fromList :: [a] -> Array a

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

151

152-------------------------------------------------------------------------------

153-- Elimination - Unfolds

154-------------------------------------------------------------------------------

155

156{-# INLINE length #-}

157length :: Array a -> Int

158length = arrLen

159

160{-# INLINE_NORMAL reader #-}

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

162reader =

163 Producer.simplify

164 $ Producer.translate unsafeThaw unsafeFreeze

165 $ MArray.producerWith (return . unsafeInlineIO)

166

167-------------------------------------------------------------------------------

168-- Elimination - to streams

169-------------------------------------------------------------------------------

170

171{-# INLINE_NORMAL toList #-}

172toList :: Array a -> [a]

173toList arr = loop 0

174

175 where

176

177 len = length arr

178 loop i | i == len = []

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

180

181{-# INLINE_NORMAL read #-}

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

183read arr@Array{..} =

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

185

186{-# INLINE_NORMAL readRev #-}

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

188readRev arr@Array{..} =

189 D.map (`getIndexUnsafe` arr)

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

191

192-------------------------------------------------------------------------------

193-- Elimination - using Folds

194-------------------------------------------------------------------------------

195

196{-# INLINE_NORMAL foldl' #-}

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

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

199

200{-# INLINE_NORMAL foldr #-}

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

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

203

204{-# INLINE fold #-}

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

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

207

208{-# INLINE streamFold #-}

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

210streamFold f arr = f (read arr)

211

212-------------------------------------------------------------------------------

213-- Random reads and writes

214-------------------------------------------------------------------------------

215

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

217-- not check the bounds.

218--

219-- @since 0.8.0

220{-# INLINE getIndexUnsafe #-}

221getIndexUnsafe :: Int -> Array a -> a

222getIndexUnsafe i arr =

223 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)

224

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

226--

227{-# INLINE getIndex #-}

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

229getIndex i arr@Array {..} =

230 if i >= 0 && i < arrLen

231 then Just $ getIndexUnsafe i arr

232 else Nothing

233

234{-# INLINE writeLastN #-}

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

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

237

238 where

239

240 f rb = do

241 arr <- RB.toMutArray 0 n rb

242 return $ unsafeFreeze arr

243

244{-# INLINE getSliceUnsafe #-}

245getSliceUnsafe :: Int -> Int -> Array a -> Array a

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

247

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

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

250-- copying.

251

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

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

254{-# INLINE strip #-}

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

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

257

258-------------------------------------------------------------------------------

259-- Instances

260-------------------------------------------------------------------------------

261

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

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

264 arr1 == arr2 =

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

266

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

268 {-# INLINE compare #-}

269 compare arr1 arr2 =

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

271

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

273 -- replicate them here with an INLINE.

274 {-# INLINE (<) #-}

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

276

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

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

279

280 {-# INLINE (>) #-}

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

282

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

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

285

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

287 -- because the latter is often more expensive

288 {-# INLINE max #-}

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

290

291 {-# INLINE min #-}

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

293

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

295 {-# INLINE show #-}

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

297

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

299 {-# INLINE readPrec #-}

300 readPrec = do

301 fromListWord <- replicateM 9 ReadPrec.get

302 if fromListWord == "fromList "

303 then fromList <$> readPrec

304 else ReadPrec.pfail