0569ab5 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add unsafe array read without bound check

0569ab5Harendra Kumar7 years ago

src/Streamly/Internal/Memory/Array.hs+38

compare with master

unifiedsplit

@@ -67,6 +67,7 @@ module Streamly.Internal.Memory.Array

67 , toStream

68 , toStreamRev

69 , read

70+ , unsafeRead

70 -- , readChunksOf

71 -- , readChunksOf

72 -- * Random Access

73 -- * Random Access

@@ -135,6 +136,8 @@ import Prelude hiding (length, null, last, map, (!!), read, concat)

135

136

136import GHC.ForeignPtr (ForeignPtr(..))

137import GHC.ForeignPtr (ForeignPtr(..))

137import GHC.Ptr (Ptr(..))

138import GHC.Ptr (Ptr(..))

139+import GHC.Prim (touch#)

140+import GHC.IO (IO(..))

138

141

139import Streamly.Internal.Data.Fold.Types (Fold(..))

142import Streamly.Internal.Data.Fold.Types (Fold(..))

140import Streamly.Internal.Data.Unfold.Types (Unfold(..))

143import Streamly.Internal.Data.Unfold.Types (Unfold(..))

@@ -239,6 +242,41 @@ read = Unfold step inject

239 return $ D.Yield x

242 return $ D.Yield x

240 (ReadUState fp (p `plusPtr` (sizeOf (undefined :: a))))

243 (ReadUState fp (p `plusPtr` (sizeOf (undefined :: a))))

241

244

245+-- | Unfold an array into a stream, does not check the end of the array, the

246+-- user is responsible for terminating the stream within the array bounds. For

247+-- high performance application where the end condition can be determined by

248+-- a terminating fold.

249+--

250+-- Written in the hope that it may be faster than "read", however, in the case

251+-- for which this was written, "read" proves to be faster even though the core

252+-- generated with unsafeRead looks simpler.

253+--

254+-- /Internal/

255+--

256+{-# INLINE_NORMAL unsafeRead #-}

257+unsafeRead :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

258+unsafeRead = Unfold step inject

259+ where

260+

261+ inject (Array fp _ _) = return fp

262+

263+ {-# INLINE_LATE step #-}

264+ step (ForeignPtr p contents) = do

265+ -- unsafeInlineIO allows us to run this in Identity monad for pure

266+ -- toList/foldr case which makes them much faster due to not

267+ -- accumulating the list and fusing better with the pure consumers.

268+ --

269+ -- This should be safe as the array contents are guaranteed to be

270+ -- evaluated/written to before we peek at them.

271+ let !x = A.unsafeInlineIO $ do

272+ r <- peek (Ptr p)

273+ touch contents

274+ return r

275+ let !(Ptr p1) = Ptr p `plusPtr` (sizeOf (undefined :: a))

276+ return $ D.Yield x (ForeignPtr p1 contents)

277+

278+ touch r = IO $ \s -> case touch# r s of s' -> (# s', () #)

279+

242-- | > null arr = length arr == 0

280-- | > null arr = length arr == 0

243--

281--

244-- /Internal/

282-- /Internal/