{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-cse #-}

-- |
-- Module      :   Grisette.IR.SymPrim.Data.Prim.InternedTerm.Caches
-- Copyright   :   (c) Sirui Lu 2021-2023
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.IR.SymPrim.Data.Prim.InternedTerm.Caches (typeMemoizedCache) where

import Control.Concurrent
  ( forkIO,
    newEmptyMVar,
    putMVar,
    readMVar,
    takeMVar,
    tryPutMVar,
  )
import Data.Data (Proxy (Proxy), TypeRep, Typeable, typeRep)
import qualified Data.HashMap.Strict as M
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Interned (Cache, Interned, mkCache)
import GHC.Base (Any)
import GHC.IO (unsafeDupablePerformIO, unsafePerformIO)
import Unsafe.Coerce (unsafeCoerce)

mkOnceIO :: IO a -> IO (IO a)
mkOnceIO :: forall a. IO a -> IO (IO a)
mkOnceIO IO a
io = do
  MVar a
mv <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
  MVar ()
demand <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  IO () -> IO ThreadId
forkIO (MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
demand IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO a
io IO a -> (a -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
mv)
  IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MVar () -> () -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar ()
demand () IO Bool -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MVar a -> IO a
forall a. MVar a -> IO a
readMVar MVar a
mv)

termCacheCell :: IO (IORef (M.HashMap TypeRep Any))
termCacheCell :: IO (IORef (HashMap TypeRep Any))
termCacheCell = IO (IO (IORef (HashMap TypeRep Any)))
-> IO (IORef (HashMap TypeRep Any))
forall a. IO a -> a
unsafePerformIO (IO (IO (IORef (HashMap TypeRep Any)))
 -> IO (IORef (HashMap TypeRep Any)))
-> IO (IO (IORef (HashMap TypeRep Any)))
-> IO (IORef (HashMap TypeRep Any))
forall a b. (a -> b) -> a -> b
$ IO (IORef (HashMap TypeRep Any))
-> IO (IO (IORef (HashMap TypeRep Any)))
forall a. IO a -> IO (IO a)
mkOnceIO (IO (IORef (HashMap TypeRep Any))
 -> IO (IO (IORef (HashMap TypeRep Any))))
-> IO (IORef (HashMap TypeRep Any))
-> IO (IO (IORef (HashMap TypeRep Any)))
forall a b. (a -> b) -> a -> b
$ HashMap TypeRep Any -> IO (IORef (HashMap TypeRep Any))
forall a. a -> IO (IORef a)
newIORef HashMap TypeRep Any
forall k v. HashMap k v
M.empty
{-# NOINLINE termCacheCell #-}

typeMemoizedCache :: forall a. (Interned a, Typeable a) => Cache a
typeMemoizedCache :: forall a. (Interned a, Typeable a) => Cache a
typeMemoizedCache = IO (Cache a) -> Cache a
forall a. IO a -> a
unsafeDupablePerformIO (IO (Cache a) -> Cache a) -> IO (Cache a) -> Cache a
forall a b. (a -> b) -> a -> b
$ do
  IORef (HashMap TypeRep Any)
c <- IO (IORef (HashMap TypeRep Any))
termCacheCell
  IORef (HashMap TypeRep Any)
-> (HashMap TypeRep Any -> (HashMap TypeRep Any, Cache a))
-> IO (Cache a)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (HashMap TypeRep Any)
c ((HashMap TypeRep Any -> (HashMap TypeRep Any, Cache a))
 -> IO (Cache a))
-> (HashMap TypeRep Any -> (HashMap TypeRep Any, Cache a))
-> IO (Cache a)
forall a b. (a -> b) -> a -> b
$ \HashMap TypeRep Any
m ->
    case TypeRep -> HashMap TypeRep Any -> Maybe Any
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a)) HashMap TypeRep Any
m of
      Just Any
d -> (HashMap TypeRep Any
m, Any -> Cache a
forall a b. a -> b
unsafeCoerce Any
d)
      Maybe Any
Nothing -> (TypeRep -> Any -> HashMap TypeRep Any -> HashMap TypeRep Any
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a)) (Cache a -> Any
forall a b. a -> b
unsafeCoerce Cache a
r1) HashMap TypeRep Any
m, Cache a
r1)
        where
          r1 :: Cache a
          !r1 :: Cache a
r1 = Cache a
forall t. Interned t => Cache t
mkCache
          {-# NOINLINE r1 #-}