base: make Foreign.Marshal.Pool use RTS internal arena for allocation

`Foreign.Marshal.Pool` used to call `malloc` once for each allocation
request. Each `Pool` maintained a list of allocated pointers, and
traverses the list to `free` each one of those pointers. The extra O(n)
overhead is apparently bad for a `Pool` that serves a lot of small
allocation requests.

This patch uses the RTS internal arena to implement `Pool`, with these
benefits:

- Gets rid of the extra O(n) overhead.
- The RTS arena is simply a bump allocator backed by the block
  allocator, each allocation request is likely faster than a libc
  `malloc` call.

Closes #14762 #18338.

3 files changed, 32 insertions, 30 deletions

diff --git a/libraries/base/Foreign/Marshal/Pool.hs b/libraries/base/Foreign/Marshal/Pool.hs
index 8d704c1a2d..7e6443b6ea 100644
--- a/libraries/base/Foreign/Marshal/Pool.hs
+++ b/libraries/base/Foreign/Marshal/Pool.hs
@@ -46,19 +46,18 @@ module Foreign.Marshal.Pool (
    pooledNewArray0
 ) where
 
-import GHC.Base              ( Int, Monad(..), (.), liftM, not )
+import GHC.Base              ( Int, Monad(..) )
 import GHC.Err               ( undefined )
 import GHC.Exception         ( throw )
 import GHC.IO                ( IO, mask, catchAny )
-import GHC.IORef             ( IORef, newIORef, readIORef, writeIORef )
-import GHC.List              ( elem, length )
+import GHC.List              ( length )
 import GHC.Num               ( Num(..) )
+import GHC.Real              ( fromIntegral )
 
-import Data.OldList          ( delete )
-import Foreign.Marshal.Alloc ( mallocBytes, reallocBytes, free )
+import Foreign.C.Types       ( CSize(..) )
 import Foreign.Marshal.Array ( pokeArray, pokeArray0 )
-import Foreign.Marshal.Error ( throwIf )
-import Foreign.Ptr           ( Ptr, castPtr )
+import Foreign.Marshal.Utils ( moveBytes )
+import Foreign.Ptr           ( Ptr )
 import Foreign.Storable      ( Storable(sizeOf, poke) )
 
 --------------------------------------------------------------------------------
@@ -68,20 +67,18 @@ import Foreign.Storable      ( Storable(sizeOf, poke) )
 
 -- | A memory pool.
 
-newtype Pool = Pool (IORef [Ptr ()])
+newtype Pool = Pool (Ptr ())
 
 -- | Allocate a fresh memory pool.
 
 newPool :: IO Pool
-newPool = liftM Pool (newIORef [])
+newPool = c_newArena
 
 -- | Deallocate a memory pool and everything which has been allocated in the
 -- pool itself.
 
 freePool :: Pool -> IO ()
-freePool (Pool pool) = readIORef pool >>= freeAll
-   where freeAll []     = return ()
-         freeAll (p:ps) = free p >> freeAll ps
+freePool = c_arenaFree
 
 -- | Execute an action with a fresh memory pool, which gets automatically
 -- deallocated (including its contents) after the action has finished.
@@ -108,11 +105,7 @@ pooledMalloc pool = pooledMallocBytes pool (sizeOf (undefined :: a))
 -- | Allocate the given number of bytes of storage in the pool.
 
 pooledMallocBytes :: Pool -> Int -> IO (Ptr a)
-pooledMallocBytes (Pool pool) size = do
-   ptr <- mallocBytes size
-   ptrs <- readIORef pool
-   writeIORef pool (ptr:ptrs)
-   return (castPtr ptr)
+pooledMallocBytes pool size = c_arenaAlloc pool (fromIntegral size)
 
 -- | Adjust the storage area for an element in the pool to the given size of
 -- the required type.
@@ -120,16 +113,15 @@ pooledMallocBytes (Pool pool) size = do
 pooledRealloc :: forall a . Storable a => Pool -> Ptr a -> IO (Ptr a)
 pooledRealloc pool ptr = pooledReallocBytes pool ptr (sizeOf (undefined :: a))
 
--- | Adjust the storage area for an element in the pool to the given size.
+-- | Adjust the storage area for an element in the pool to the given size. Note
+-- that the previously allocated space is still retained in the same 'Pool' and
+-- will only be freed when the entire 'Pool' is freed.
 
 pooledReallocBytes :: Pool -> Ptr a -> Int -> IO (Ptr a)
-pooledReallocBytes (Pool pool) ptr size = do
-   let cPtr = castPtr ptr
-   _ <- throwIf (not . (cPtr `elem`)) (\_ -> "pointer not in pool") (readIORef pool)
-   newPtr <- reallocBytes cPtr size
-   ptrs <- readIORef pool
-   writeIORef pool (newPtr : delete cPtr ptrs)
-   return (castPtr newPtr)
+pooledReallocBytes pool ptr size = do
+   newPtr <- pooledMallocBytes pool size
+   moveBytes newPtr ptr size
+   return newPtr
 
 -- | Allocate storage for the given number of elements of a storable type in the
 -- pool.
@@ -185,3 +177,9 @@ pooledNewArray0 pool marker vals = do
    ptr <- pooledMallocArray0 pool (length vals)
    pokeArray0 marker ptr vals
    return ptr
+
+foreign import ccall unsafe "newArena" c_newArena :: IO Pool
+
+foreign import ccall unsafe "arenaAlloc" c_arenaAlloc :: Pool -> CSize -> IO (Ptr a)
+
+foreign import ccall unsafe "arenaFree" c_arenaFree :: Pool -> IO ()
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index d955cd59a4..cc5f44ef96 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -31,8 +31,8 @@
     as well as [the migration guide](https://github.com/haskell/core-libraries-committee/blob/main/guides/export-lifta2-prelude.md)
   * Update to [Unicode 15.0.0](https://www.unicode.org/versions/Unicode15.0.0/).
   * Add standard Unicode case predicates `isUpperCase` and `isLowerCase` to
-    `GHC.Unicode` and `Data.Char`. These predicates use the standard Unicode 
-    case properties and are more intuitive than `isUpper` and `isLower`. See 
+    `GHC.Unicode` and `Data.Char`. These predicates use the standard Unicode
+    case properties and are more intuitive than `isUpper` and `isLower`. See
     [CLC proposal #90](https://github.com/haskell/core-libraries-committee/issues/90).
   * Add `Eq` and `Ord` instances for `Generically1`.
   * Relax instances for Functor combinators; put superclass on Class1 and Class2
@@ -50,6 +50,10 @@
   * The `Enum` instance of `Down a` now enumerates values in the opposite
     order as the `Enum a` instance, per
     [CLC proposal #51](https://github.com/haskell/core-libraries-committee/issues/51).
+  * `Foreign.Marshal.Pool` now uses the RTS internal arena instead of libc
+    `malloc` for allocation. It avoids the O(n) overhead of maintaining a list
+    of individually allocated pointers as well as freeing each one of them when
+    freeing a `Pool`. (#14762) (#18338)
 
 ## 4.17.0.0 *August 2022*
 
diff --git a/rts/Arena.h b/rts/Arena.h
index 49298713ab..8a7b304a7a 100644
--- a/rts/Arena.h
+++ b/rts/Arena.h
@@ -10,13 +10,13 @@
 typedef struct _Arena Arena;
 
 // Start a new arena
-RTS_PRIVATE Arena * newArena   ( void );
+Arena * newArena   ( void );
 
 // Allocate memory in an arena
-RTS_PRIVATE void  * arenaAlloc ( Arena *, size_t );
+void  * arenaAlloc ( Arena *, size_t );
 
 // Free an entire arena
-RTS_PRIVATE void arenaFree  ( Arena * );
+void arenaFree  ( Arena * );
 
 // For internal use only:
 RTS_PRIVATE unsigned long arenaBlocks( void );