compiler/cmm/OptimizationFuel.hs

   1 module OptimizationFuel
   2     ( OptimizationFuel, canRewriteWithFuel, maybeRewriteWithFuel, oneLessFuel
   3     , tankFilledTo, diffFuel
   4     , FuelConsumer
   5     , FuelUsingMonad, FuelState
   6     , lastFuelPass, fuelExhausted, fuelRemaining, fuelDecrement
   7     , lastFuelPassInState, fuelExhaustedInState, fuelRemainingInState
   8     , fuelDecrementState
   9     , runFuel, runFuelIO, runFuelWithLastPass, fuelConsumingPass
  10     , runWithInfiniteFuel
  11     , FuelMonad(..)
  12     )
  13 where
  14
  15 --import GHC.Exts (State#)
  16 import Panic
  17
  18 import Data.IORef
  19
  20 #include "HsVersions.h"
  21
  22 type FuelConsumer a = OptimizationFuel -> (a, OptimizationFuel)
  23
  24 canRewriteWithFuel :: OptimizationFuel -> Bool
  25 oneLessFuel :: OptimizationFuel -> OptimizationFuel
  26 maybeRewriteWithFuel :: OptimizationFuel -> Maybe a -> Maybe a
  27 diffFuel :: OptimizationFuel -> OptimizationFuel -> Int
  28    -- to measure consumption during compilation
  29 tankFilledTo :: Int -> OptimizationFuel
  30
  31 #ifdef DEBUG
  32 newtype OptimizationFuel = OptimizationFuel Int
  33   deriving Show
  34
  35 tankFilledTo = OptimizationFuel
  36 canRewriteWithFuel (OptimizationFuel f) = f > 0
  37 maybeRewriteWithFuel fuel ma = if canRewriteWithFuel fuel then ma else Nothing
  38 oneLessFuel (OptimizationFuel f) = ASSERT (f > 0) (OptimizationFuel (f - 1))
  39 diffFuel (OptimizationFuel f) (OptimizationFuel f') = f - f'
  40 #else
  41 -- type OptimizationFuel = State# () -- would like this, but it won't work
  42 data OptimizationFuel = OptimizationFuel
  43   deriving Show
  44 tankFilledTo _ = undefined -- should be impossible to evaluate
  45   -- realWorld# might come in handy, too...
  46 canRewriteWithFuel OptimizationFuel = True
  47 maybeRewriteWithFuel _ ma = ma
  48 oneLessFuel f = f
  49 diffFuel _ _ = 0
  50 #endif
  51
  52 data FuelState = FuelState { fs_fuellimit :: OptimizationFuel, fs_lastpass :: String }
  53 newtype FuelMonad a = FuelMonad (FuelState -> (a, FuelState))
  54
  55 fuelConsumingPass :: String -> FuelConsumer a -> FuelMonad a
  56 fuelConsumingPass name f = do fuel <- fuelRemaining
  57                               let (a, fuel') = f fuel
  58                               fuelDecrement name fuel fuel'
  59                               return a
  60
  61 runFuel             :: FuelMonad a -> FuelConsumer a
  62 runFuelWithLastPass :: FuelMonad a -> FuelConsumer (a, String)
  63 runWithInfiniteFuel :: FuelMonad a -> a
  64
  65
  66 runFuelIO :: IORef String -> IORef OptimizationFuel -> FuelMonad a -> IO a
  67 runFuelIO pass_ref fuel_ref (FuelMonad f) =
  68     do { pass <- readIORef pass_ref
  69        ; fuel <- readIORef fuel_ref
  70        ; let (a, FuelState fuel' pass') = f (FuelState fuel pass)
  71        ; writeIORef pass_ref pass'
  72        ; writeIORef fuel_ref fuel'
  73        ; return a
  74        }
  75
  76 initialFuelState :: OptimizationFuel -> FuelState
  77 initialFuelState fuel = FuelState fuel "unoptimized program"
  78
  79 runFuel             (FuelMonad f) fuel = let (a, s) = f $ initialFuelState fuel
  80                                          in (a, fs_fuellimit s)
  81 runFuelWithLastPass (FuelMonad f) fuel = let (a, s) = f $ initialFuelState fuel
  82                                          in ((a, fs_lastpass s), fs_fuellimit s)
  83
  84 runWithInfiniteFuel (FuelMonad f) = fst $ f $ initialFuelState $ tankFilledTo maxBound
  85
  86 lastFuelPassInState :: FuelState -> String
  87 lastFuelPassInState = fs_lastpass
  88
  89 fuelExhaustedInState :: FuelState -> Bool
  90 fuelExhaustedInState = canRewriteWithFuel . fs_fuellimit
  91
  92 fuelRemainingInState :: FuelState -> OptimizationFuel
  93 fuelRemainingInState = fs_fuellimit
  94
  95 fuelDecrementState
  96     :: String -> OptimizationFuel -> OptimizationFuel -> FuelState -> FuelState
  97 fuelDecrementState new_optimizer old new s =
  98     FuelState { fs_fuellimit = lim, fs_lastpass = optimizer }
  99   where lim = if diffFuel old (fs_fuellimit s) == 0 then new
 100               else panic $
 101                    concat ["lost track of ", new_optimizer, "'s transactions"]
 102         optimizer = if diffFuel old new > 0 then new_optimizer else fs_lastpass s
 103
 104 class Monad m => FuelUsingMonad m where
 105   fuelRemaining :: m OptimizationFuel
 106   fuelDecrement :: String -> OptimizationFuel -> OptimizationFuel -> m ()
 107   fuelExhausted :: m Bool
 108   lastFuelPass  :: m String
 109
 110
 111 instance Monad FuelMonad where
 112   FuelMonad f >>= k = FuelMonad (\s -> let (a, s') = f s
 113                                            FuelMonad f' = k a
 114                                        in  f' s')
 115   return a = FuelMonad (\s -> (a, s))
 116
 117 instance FuelUsingMonad FuelMonad where
 118   fuelRemaining = extract fuelRemainingInState
 119   lastFuelPass  = extract lastFuelPassInState
 120   fuelExhausted = extract fuelExhaustedInState
 121   fuelDecrement p f f' = FuelMonad (\s -> ((), fuelDecrementState p f f' s))
 122
 123 extract :: (FuelState -> a) -> FuelMonad a
 124 extract f = FuelMonad (\s -> (f s, s))