{-# LANGUAGE LambdaCase #-}
module Control.Once.Internal where
import qualified Data.HashMap.Strict as HM
import Data.Hashable
import Data.IORef
import Control.Monad

once0 :: IO a -> IO (IO a)
once0 act = do
    ref <- newIORef Nothing
    pure $ readIORef ref >>= \case
        Nothing -> do
            value <- act
            atomicWriteIORef ref (Just value)
            pure value
        Just value -> pure value

once1 :: (Eq a, Hashable a) => (a -> IO b) -> IO (a -> IO b)
once1 fn = do
    ref <- newIORef HM.empty
    pure $ \arg -> do
        action <- once0 $ fn arg
        let modify' map = case HM.lookup arg map of
                              Just was -> (map, was)
                              Nothing  -> (HM.insert arg action map, action)
        join $ atomicModifyIORef' ref modify'