on master help sign in
A tiny bit more efficient getLocation
commit6506958cce216db87aed2f7e1295e51a64ef8868 authorAndrzej Rybczak <andrzej@rybczak.net> dateTue 28 Dec 2021, 16:04 (5 years ago) parent3b965ed
1 file changed+21−16
effectful-core/src/Effectful/Internal/Env.hs+21−16
effectful-core/src/Effectful/Internal/Env.hs+21−16
@@ -1,3 +1,4 @@
1+{-# LANGUAGE MagicHash #-}
12{-# OPTIONS_HADDOCK not-home #-}
23-- | The environment for the 'Effectful.Internal.Monad.Eff' monad.
34--
@@ -42,7 +43,7 @@ import Control.Monad
4243import Control.Monad.Primitive
4344import Data.IORef
4445import Data.Primitive.SmallArray
45−import GHC.Exts (Any)
46+import GHC.Exts (Any, SmallMutableArray#)
4647import GHC.Stack (HasCallStack)
4748import Unsafe.Coerce
4849import qualified Data.IntMap.Strict as IM
@@ -421,7 +422,7 @@ getEnv
421422 => Env es
422423 -> IO (adapter e)
423424getEnv env = do
424− (i, es) <- getLocation (reifyIndex @e @es) env
425+ Location i es <- getLocation (reifyIndex @e @es) env
425426 fromAny <$> readSmallArray es i
426427
427428-- | Replace the data type in the environment with a new value (in place).
@@ -431,7 +432,7 @@ putEnv
431432 -> adapter e
432433 -> IO ()
433434putEnv env e = do
434− (i, es) <- getLocation (reifyIndex @e @es) env
435+ Location i es <- getLocation (reifyIndex @e @es) env
435436 e `seq` writeSmallArray es i (toAny e)
436437
437438-- | Modify the data type in the environment (in place) and return a value.
@@ -441,7 +442,7 @@ stateEnv
441442 -> (adapter e -> (a, adapter e))
442443 -> IO a
443444stateEnv env f = do
444− (i, es) <- getLocation (reifyIndex @e @es) env
445+ Location i es <- getLocation (reifyIndex @e @es) env
445446 (a, e) <- f . fromAny <$> readSmallArray es i
446447 e `seq` writeSmallArray es i (toAny e)
447448 pure a
@@ -453,7 +454,7 @@ modifyEnv
453454 -> (adapter e -> adapter e)
454455 -> IO ()
455456modifyEnv env f = do
456− (i, es) <- getLocation (reifyIndex @e @es) env
457+ Location i es <- getLocation (reifyIndex @e @es) env
457458 e <- f . fromAny <$> readSmallArray es i
458459 e `seq` writeSmallArray es i (toAny e)
459460
@@ -469,13 +470,17 @@ emptyEnvRef = do
469470 fs <- newSmallArray 0 undefinedData
470471 newIORef $ EnvRef 0 es fs
471472
473+-- | Location with unboxed fields to make GHC generate Core without unnecessary
474+-- reboxing around the local 'go' function from 'getLocation'.
475+data Location = Location !Int !(SmallMutableArray RealWorld Any)
476+
472477-- | Determine location of the data type in the environment.
473478getLocation
474479 :: Int
475480 -> Env es
476− -> IO (Int, SmallMutableArray RealWorld Any)
481+ -> IO Location
477482getLocation ixE (Env fork ref _) = do
478− EnvRef n es _ <- readIORef ref
483+ EnvRef n es@(SmallMutableArray es#) _ <- readIORef ref
479484 -- Optimized for the common access pattern:
480485 --
481486 -- - Most application code has no access to forks, in which case we look at
@@ -484,25 +489,25 @@ getLocation ixE (Env fork ref _) = do
484489 -- - Effect handlers will most likely access the newest fork with handler
485490 -- specific effects for reinterpretation.
486491 case fork of
487− NoFork -> pure (ix n, es)
492+ NoFork -> pure $ Location (ix n) es
488493 Forks _ baseIx lref forks -> do
489494 EnvRef ln les _ <- readIORef lref
490495 let i = ix (baseIx + ln)
491496 if i >= baseIx
492− then pure (i - baseIx, les)
493− else go es i forks
497+ then pure $ Location (i - baseIx) les
498+ else go es# i forks
494499 where
495− go :: SmallMutableArray RealWorld Any
500+ go :: SmallMutableArray# RealWorld Any
496501 -> Int
497502 -> Forks
498− -> IO (Int, SmallMutableArray RealWorld Any)
499− go es i = \case
500− NoFork -> pure (i, es)
503+ -> IO Location
504+ go es# i = \case
505+ NoFork -> pure $ Location i (SmallMutableArray es#)
501506 Forks _ baseIx lref forks -> do
502507 EnvRef _ les _ <- readIORef lref
503508 if i >= baseIx
504− then pure (i - baseIx, les)
505− else go es i forks
509+ then pure $ Location (i - baseIx) les
510+ else go es# i forks
506511
507512 ix :: Int -> Int
508513 ix n = n - ixE - 1