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
  IORef (UniqFM k a)
ref <- UniqFM k a -> IO (IORef (UniqFM k a))
forall a. a -> IO (IORef a)
newIORef UniqFM k a
forall key elt. UniqFM key elt
emptyUFM
  (k -> a) -> IO (k -> a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((k -> a) -> IO (k -> a)) -> (k -> a) -> IO (k -> a)
forall a b. (a -> b) -> a -> b
$ \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
    UniqFM k a
m <- IORef (UniqFM k a) -> IO (UniqFM k a)
forall a. IORef a -> IO a
readIORef IORef (UniqFM k a)
ref
    case UniqFM k a -> k -> Maybe a
forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
lookupUFM UniqFM k a
m k
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