module GHC.Types.Unique.MemoFun (memoiseUniqueFun) where

import GHC.Prelude
import GHC.Types.Unique
import GHC.Types.Unique.FM

import Data.IORef
import System.IO.Unsafe

memoiseUniqueFun :: Uniquable k => (k -> a) -> k -> a
memoiseUniqueFun :: forall k a. Uniquable k => (k -> a) -> k -> a
memoiseUniqueFun k -> a
fun = IO (k -> a) -> k -> a
forall a. IO a -> a
unsafePerformIO (IO (k -> a) -> k -> a) -> IO (k -> a) -> k -> a
forall a b. (a -> b) -> a -> b
$ do
  ref <- UniqFM k a -> IO (IORef (UniqFM k a))
forall a. a -> IO (IORef a)
newIORef UniqFM k a
forall {k} (key :: k) elt. UniqFM key elt
emptyUFM
  return $ \k
k -> IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ do
    m <- IORef (UniqFM k a) -> IO (UniqFM k a)
forall a. IORef a -> IO a
readIORef IORef (UniqFM k a)
ref
    case lookupUFM m k of
      Just a
a  -> a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a
      Maybe a
Nothing -> do
        let !a :: a
a  = k -> a
fun k
k
            !m' :: UniqFM k a
m' = UniqFM k a -> k -> a -> UniqFM k a
forall key elt.
Uniquable key =>
UniqFM key elt -> key -> elt -> UniqFM key elt
addToUFM UniqFM k a
m k
k a
a
        IORef (UniqFM k a) -> UniqFM k a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (UniqFM k a)
ref UniqFM k a
m'
        a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a