{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# OPTIONS_GHC -fno-full-laziness -fno-cse #-}
module Data.StableName.Extra (
  StableName,
  makeStableName',
  showStableName,
) where

import Data.Hashable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import Data.IORef
import GHC.StableName
import System.IO.Unsafe


-- | This function evaluates its argument to WHNF and returns a stable name for
-- the evaluation result. This function is not referentially transparent and is
-- implemented using 'unsafePerformIO'.
{-# NOINLINE makeStableName' #-}
makeStableName' :: a -> StableName a
makeStableName' !x = unsafePerformIO (makeStableName x)


data SomeStableName = forall a. SomeStableName (StableName a)

instance Eq SomeStableName where
  SomeStableName a == SomeStableName b = eqStableName a b

instance Hashable SomeStableName where
  hashWithSalt salt (SomeStableName name) = hashWithSalt salt name

{-# NOINLINE showStableNameCache #-}
showStableNameCache :: IORef (HashMap SomeStableName Int, Int)
showStableNameCache = unsafePerformIO $ newIORef (mempty, 0)

{-# NOINLINE showStableName #-}
showStableName :: StableName a -> String
showStableName name =
  unsafePerformIO $ atomicModifyIORef' showStableNameCache $ \tup@(mp, nexti) ->
    case HM.lookup (SomeStableName name) mp of
      Just res -> (tup, '$' : show res)
      Nothing -> ((HM.insert (SomeStableName name) nexti mp, nexti + 1), '$' : show nexti)