cca647c · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Lift the API in Foreign.Mut.Types to a higher level MonadIO m => m

cca647cAdithya Kumar5 years ago

src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+50−31

compare with master

unifiedsplit

@@ -176,6 +176,13 @@ memcmp p1 p2 len = do

176 r <- c_memcmp p1 p2 (fromIntegral len)

177 return $ r == 0

178

179+{-# INLINE unsafeWithForeignPtrM #-}

180+unsafeWithForeignPtrM :: MonadIO m => ForeignPtr a -> (Ptr a -> m b) -> m b

181+unsafeWithForeignPtrM fo f = do

182+ r <- f (unsafeForeignPtrToPtr fo)

183+ liftIO $ touchForeignPtr fo

184+ return r

185+

179-------------------------------------------------------------------------------

186-------------------------------------------------------------------------------

180-- Array Data Type

187-- Array Data Type

181-------------------------------------------------------------------------------

188-------------------------------------------------------------------------------

@@ -213,8 +220,8 @@ mutableArray = Array

213

220

214-- | allocate a new array using the provided allocator function.

221-- | allocate a new array using the provided allocator function.

215{-# INLINE newArrayAlignedAllocWith #-}

222{-# INLINE newArrayAlignedAllocWith #-}

216−newArrayAlignedAllocWith :: forall a. Storable a

223+newArrayAlignedAllocWith :: forall m a. (MonadIO m, Storable a)

217− => (Int -> Int -> IO (ForeignPtr a)) -> Int -> Int -> IO (Array a)

224+ => (Int -> Int -> m (ForeignPtr a)) -> Int -> Int -> m (Array a)

218newArrayAlignedAllocWith alloc alignSize count = do

225newArrayAlignedAllocWith alloc alignSize count = do

219 let size = count * sizeOf (undefined :: a)

226 let size = count * sizeOf (undefined :: a)

220 fptr <- alloc size alignSize

227 fptr <- alloc size alignSize

@@ -230,13 +237,23 @@ newArrayAlignedAllocWith alloc alignSize count = do

230-- This could be useful in allocate once global data structures. Use carefully

237-- This could be useful in allocate once global data structures. Use carefully

231-- as incorrect use can lead to memory leak.

238-- as incorrect use can lead to memory leak.

232{-# INLINE newArrayAlignedUnmanaged #-}

239{-# INLINE newArrayAlignedUnmanaged #-}

233−newArrayAlignedUnmanaged :: forall a. Storable a => Int -> Int -> IO (Array a)

240+newArrayAlignedUnmanaged :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)

234newArrayAlignedUnmanaged =

241newArrayAlignedUnmanaged =

242+ newArrayAlignedAllocWith mallocForeignPtrAlignedUnmanagedBytes

243+

244+ where

245+

246+ mallocForeignPtrAlignedUnmanagedBytes size alignSize =

235− newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedUnmanagedBytes

247+ liftIO $ Malloc.mallocForeignPtrAlignedUnmanagedBytes size alignSize

236

248

237{-# INLINE newArrayAligned #-}

249{-# INLINE newArrayAligned #-}

238−newArrayAligned :: forall a. Storable a => Int -> Int -> IO (Array a)

250+newArrayAligned :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)

239−newArrayAligned = newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedBytes

251+newArrayAligned = newArrayAlignedAllocWith mallocForeignPtrAlignedBytes

252+

253+ where

254+

255+ mallocForeignPtrAlignedBytes size alignSize =

256+ liftIO $ Malloc.mallocForeignPtrAlignedBytes size alignSize

240

257

241-- XXX can unaligned allocation be more efficient when alignment is not needed?

258-- XXX can unaligned allocation be more efficient when alignment is not needed?

242--

259--

@@ -246,52 +263,54 @@ newArrayAligned = newArrayAlignedAllocWith Malloc.mallocForeignPtrAlignedBytes

246-- Note that this is internal routine, the reference to this array cannot be

263-- Note that this is internal routine, the reference to this array cannot be

247-- given out until the array has been written to and frozen.

264-- given out until the array has been written to and frozen.

248{-# INLINE newArray #-}

265{-# INLINE newArray #-}

249−newArray :: forall a. Storable a => Int -> IO (Array a)

266+newArray :: forall m a. (MonadIO m, Storable a) => Int -> m (Array a)

250newArray = newArrayAligned (alignment (undefined :: a))

267newArray = newArrayAligned (alignment (undefined :: a))

251

268

252-- | Allocate an Array of the given size and run an IO action passing the array

269-- | Allocate an Array of the given size and run an IO action passing the array

253-- start pointer.

270-- start pointer.

254{-# INLINE unsafeWithNewArray #-}

271{-# INLINE unsafeWithNewArray #-}

272+unsafeWithNewArray ::

255−unsafeWithNewArray :: forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Array a)

273+ (MonadIO m, Storable a) => Int -> (Ptr a -> m ()) -> m (Array a)

256unsafeWithNewArray count f = do

274unsafeWithNewArray count f = do

257 arr <- newArray count

275 arr <- newArray count

258− unsafeWithForeignPtr (aStart arr) $ \p -> f p >> return arr

276+ unsafeWithForeignPtrM (aStart arr) $ \p -> f p >> return arr

259

277

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

278-------------------------------------------------------------------------------

261-- Mutation

279-- Mutation

262-------------------------------------------------------------------------------

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

263

281

264{-# INLINE unsafeWriteIndex #-}

282{-# INLINE unsafeWriteIndex #-}

265−unsafeWriteIndex :: forall a. Storable a => Array a -> Int -> a -> IO ()

283+unsafeWriteIndex :: forall m a. (MonadIO m, Storable a)

284+ => Array a -> Int -> a -> m ()

266unsafeWriteIndex Array {..} i x =

285unsafeWriteIndex Array {..} i x =

267− unsafeWithForeignPtr aStart

286+ unsafeWithForeignPtrM aStart

268 $ \begin -> do

287 $ \begin -> do

269− poke (begin `plusPtr` (i * sizeOf (undefined :: a))) x

288+ liftIO $ poke (begin `plusPtr` (i * sizeOf (undefined :: a))) x

270

289

271-- Internal routine for when the array is being created. Appends one item at

290-- Internal routine for when the array is being created. Appends one item at

272-- the end of the array. Useful when sequentially writing a stream to the

291-- the end of the array. Useful when sequentially writing a stream to the

273-- array.

292-- array.

274{-# INLINE unsafeSnoc #-}

293{-# INLINE unsafeSnoc #-}

275−unsafeSnoc :: forall a. Storable a => Array a -> a -> IO (Array a)

294+unsafeSnoc :: forall m a. (MonadIO m, Storable a) => Array a -> a -> m (Array a)

276−unsafeSnoc arr@Array {..} x = do

295+unsafeSnoc arr@Array {..} x = liftIO $ do

277 when (aEnd == aBound) $ error "BUG: unsafeSnoc: writing beyond array bounds"

296 when (aEnd == aBound) $ error "BUG: unsafeSnoc: writing beyond array bounds"

278 poke aEnd x

297 poke aEnd x

279 touchForeignPtr aStart

298 touchForeignPtr aStart

280 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

299 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

281

300

282{-# INLINE snoc #-}

301{-# INLINE snoc #-}

283−snoc :: forall a. Storable a => Array a -> a -> IO (Array a)

302+snoc :: forall m a. (MonadIO m, Storable a) => Array a -> a -> m (Array a)

284snoc arr@Array {..} x =

303snoc arr@Array {..} x =

285 if aEnd == aBound

304 if aEnd == aBound

286 then do

305 then do

287 let oldStart = unsafeForeignPtrToPtr aStart

306 let oldStart = unsafeForeignPtrToPtr aStart

288 size = aEnd `minusPtr` oldStart

307 size = aEnd `minusPtr` oldStart

289 newSize = size + sizeOf (undefined :: a)

308 newSize = size + sizeOf (undefined :: a)

290− newPtr <-

309+ newPtr <- liftIO $

291 Malloc.mallocForeignPtrAlignedBytes

310 Malloc.mallocForeignPtrAlignedBytes

292 newSize

311 newSize

293 (alignment (undefined :: a))

312 (alignment (undefined :: a))

294− unsafeWithForeignPtr newPtr $ \pNew -> do

313+ unsafeWithForeignPtrM newPtr $ \pNew -> liftIO $ do

295 memcpy (castPtr pNew) (castPtr oldStart) size

314 memcpy (castPtr pNew) (castPtr oldStart) size

296 poke (pNew `plusPtr` size) x

315 poke (pNew `plusPtr` size) x

297 touchForeignPtr aStart

316 touchForeignPtr aStart

@@ -301,7 +320,7 @@ snoc arr@Array {..} x =

301 , aEnd = pNew `plusPtr` (size + sizeOf (undefined :: a))

320 , aEnd = pNew `plusPtr` (size + sizeOf (undefined :: a))

302 , aBound = pNew `plusPtr` newSize

321 , aBound = pNew `plusPtr` newSize

303 }

322 }

304− else do

323+ else liftIO $ do

305 poke aEnd x

324 poke aEnd x

306 touchForeignPtr aStart

325 touchForeignPtr aStart

307 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

326 return $ arr {aEnd = aEnd `plusPtr` sizeOf (undefined :: a)}

@@ -313,13 +332,13 @@ snoc arr@Array {..} x =

313-- | Reallocate the array to the specified size in bytes. If the size is less

332-- | Reallocate the array to the specified size in bytes. If the size is less

314-- than the original array the array gets truncated.

333-- than the original array the array gets truncated.

315{-# NOINLINE reallocAligned #-}

334{-# NOINLINE reallocAligned #-}

316−reallocAligned :: Int -> Int -> Array a -> IO (Array a)

335+reallocAligned :: MonadIO m => Int -> Int -> Array a -> m (Array a)

317reallocAligned alignSize newSize Array{..} = do

336reallocAligned alignSize newSize Array{..} = do

318 assert (aEnd <= aBound) (return ())

337 assert (aEnd <= aBound) (return ())

319 let oldStart = unsafeForeignPtrToPtr aStart

338 let oldStart = unsafeForeignPtrToPtr aStart

320 let size = aEnd `minusPtr` oldStart

339 let size = aEnd `minusPtr` oldStart

321− newPtr <- Malloc.mallocForeignPtrAlignedBytes newSize alignSize

340+ newPtr <- liftIO $ Malloc.mallocForeignPtrAlignedBytes newSize alignSize

322− unsafeWithForeignPtr newPtr $ \pNew -> do

341+ unsafeWithForeignPtrM newPtr $ \pNew -> liftIO $ do

323 memcpy (castPtr pNew) (castPtr oldStart) size

342 memcpy (castPtr pNew) (castPtr oldStart) size

324 touchForeignPtr aStart

343 touchForeignPtr aStart

325 return $ Array

344 return $ Array

@@ -330,11 +349,11 @@ reallocAligned alignSize newSize Array{..} = do

330

349

331-- XXX can unaligned allocation be more efficient when alignment is not needed?

350-- XXX can unaligned allocation be more efficient when alignment is not needed?

332{-# INLINABLE realloc #-}

351{-# INLINABLE realloc #-}

333−realloc :: forall a. Storable a => Int -> Array a -> IO (Array a)

352+realloc :: forall m a. (MonadIO m, Storable a) => Int -> Array a -> m (Array a)

334−realloc = reallocAligned (alignment (undefined :: a))

353+realloc i arr = liftIO $ reallocAligned (alignment (undefined :: a)) i arr

335

354

336-- | Remove the free space from an Array.

355-- | Remove the free space from an Array.

337−shrinkToFit :: forall a. Storable a => Array a -> IO (Array a)

356+shrinkToFit :: (MonadIO m, Storable a) => Array a -> m (Array a)

338shrinkToFit arr@Array{..} = do

357shrinkToFit arr@Array{..} = do

339 assert (aEnd <= aBound) (return ())

358 assert (aEnd <= aBound) (return ())

340 let start = unsafeForeignPtrToPtr aStart

359 let start = unsafeForeignPtrToPtr aStart

@@ -354,14 +373,14 @@ shrinkToFit arr@Array{..} = do

354--

373--

355-- Unsafe because it does not check the bounds of the array.

374-- Unsafe because it does not check the bounds of the array.

356{-# INLINE_NORMAL unsafeIndex #-}

375{-# INLINE_NORMAL unsafeIndex #-}

357−unsafeIndex :: forall a. Storable a => Array a -> Int -> IO a

376+unsafeIndex :: forall m a. (MonadIO m, Storable a) => Array a -> Int -> m a

358unsafeIndex Array {..} i =

377unsafeIndex Array {..} i =

359− unsafeWithForeignPtr aStart $ \p -> do

378+ unsafeWithForeignPtrM aStart $ \p -> do

360 let elemSize = sizeOf (undefined :: a)

379 let elemSize = sizeOf (undefined :: a)

361 elemOff = p `plusPtr` (elemSize * i)

380 elemOff = p `plusPtr` (elemSize * i)

362 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)

381 assert (i >= 0 && elemOff `plusPtr` elemSize <= aEnd)

363 (return ())

382 (return ())

364− peek elemOff

383+ liftIO $ peek elemOff

365

384

366-------------------------------------------------------------------------------

385-------------------------------------------------------------------------------

367-- Size

386-- Size

@@ -741,12 +760,12 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr

741

760

742{-# INLINE_NORMAL writeNAllocWith #-}

761{-# INLINE_NORMAL writeNAllocWith #-}

743writeNAllocWith :: forall m a. (MonadIO m, Storable a)

762writeNAllocWith :: forall m a. (MonadIO m, Storable a)

744− => (Int -> IO (Array a)) -> Int -> Fold m a (Array a)

763+ => (Int -> m (Array a)) -> Int -> Fold m a (Array a)

745writeNAllocWith alloc n = Fold step initial extract

764writeNAllocWith alloc n = Fold step initial extract

746

765

747 where

766 where

748

767

749− initial = FL.Partial <$> liftIO (alloc (max n 0))

768+ initial = FL.Partial <$> alloc (max n 0)

750 step arr@(Array _ end bound) _ | end == bound = return $ FL.Done arr

769 step arr@(Array _ end bound) _ | end == bound = return $ FL.Done arr

751 step (Array start end bound) x = do

770 step (Array start end bound) x = do

752 liftIO $ poke end x

771 liftIO $ poke end x

@@ -929,7 +948,7 @@ fromStreamDN limit str = do

929--

948--

930-- @since 0.7.0

949-- @since 0.7.0

931{-# INLINABLE fromListN #-}

950{-# INLINABLE fromListN #-}

932−fromListN :: Storable a => Int -> [a] -> IO (Array a)

951+fromListN :: (MonadIO m, Storable a) => Int -> [a] -> m (Array a)

933fromListN n xs = fromStreamDN n $ D.fromList xs

952fromListN n xs = fromStreamDN n $ D.fromList xs

934

953

935-------------------------------------------------------------------------------

954-------------------------------------------------------------------------------

@@ -962,7 +981,7 @@ fromStreamD m = do

962--

981--

963-- @since 0.7.0

982-- @since 0.7.0

964{-# INLINABLE fromList #-}

983{-# INLINABLE fromList #-}

965−fromList :: Storable a => [a] -> IO (Array a)

984+fromList :: (MonadIO m, Storable a) => [a] -> m (Array a)

966fromList xs = fromStreamD $ D.fromList xs

985fromList xs = fromStreamD $ D.fromList xs

967

986

968-------------------------------------------------------------------------------

987-------------------------------------------------------------------------------