{-# LANGUAGE TypeFamilies
           , FlexibleInstances
           , FlexibleContexts
           , BangPatterns
           , CPP
           , GeneralizedNewtypeDeriving #-}

module Data.Interned.Internal
  ( Interned(..)
  , Uninternable(..)
  , mkCache
  , Cache(..)
  , CacheState(..)
  , cacheSize
  , Id
  , intern
  , recover
  ) where

import Data.Array
import Data.Hashable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Foldable
import Data.IORef
import GHC.IO (unsafeDupablePerformIO, unsafePerformIO)

-- tuning parameter
defaultCacheWidth :: Int
defaultCacheWidth :: Int
defaultCacheWidth = Int
1024

data CacheState t = CacheState
   { forall t. CacheState t -> Int
fresh :: {-# UNPACK #-} !Id
   , forall t. CacheState t -> HashMap (Description t) t
content :: !(HashMap (Description t) t)
   }

newtype Cache t = Cache { forall t. Cache t -> Array Int (IORef (CacheState t))
getCache :: Array Int (IORef (CacheState t)) }

cacheSize :: Cache t -> IO Int
cacheSize :: forall t. Cache t -> IO Int
cacheSize (Cache Array Int (IORef (CacheState t))
t) = (IORef (CacheState t) -> Int -> IO Int)
-> Int -> Array Int (IORef (CacheState t)) -> IO Int
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
foldrM
   (\IORef (CacheState t)
a Int
b -> do
       CacheState t
v <- IORef (CacheState t) -> IO (CacheState t)
forall a. IORef a -> IO a
readIORef IORef (CacheState t)
a
       Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$! HashMap (Description t) t -> Int
forall k v. HashMap k v -> Int
HashMap.size (CacheState t -> HashMap (Description t) t
forall t. CacheState t -> HashMap (Description t) t
content CacheState t
v) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b
   ) Int
0 Array Int (IORef (CacheState t))
t

mkCache :: Interned t => Cache t
mkCache :: forall t. Interned t => Cache t
mkCache   = Cache t
result where
  element :: CacheState t
element = Int -> HashMap (Description t) t -> CacheState t
forall t. Int -> HashMap (Description t) t -> CacheState t
CacheState (Cache t -> Int
forall t (p :: * -> *). Interned t => p t -> Int
forall (p :: * -> *). p t -> Int
seedIdentity Cache t
result) HashMap (Description t) t
forall k v. HashMap k v
HashMap.empty
  w :: Int
w       = Cache t -> Int
forall t (p :: * -> *). Interned t => p t -> Int
forall (p :: * -> *). p t -> Int
cacheWidth Cache t
result
  result :: Cache t
result  = Array Int (IORef (CacheState t)) -> Cache t
forall t. Array Int (IORef (CacheState t)) -> Cache t
Cache
          (Array Int (IORef (CacheState t)) -> Cache t)
-> Array Int (IORef (CacheState t)) -> Cache t
forall a b. (a -> b) -> a -> b
$ IO (Array Int (IORef (CacheState t)))
-> Array Int (IORef (CacheState t))
forall a. IO a -> a
unsafePerformIO
          (IO (Array Int (IORef (CacheState t)))
 -> Array Int (IORef (CacheState t)))
-> IO (Array Int (IORef (CacheState t)))
-> Array Int (IORef (CacheState t))
forall a b. (a -> b) -> a -> b
$ (CacheState t -> IO (IORef (CacheState t)))
-> Array Int (CacheState t)
-> IO (Array Int (IORef (CacheState t)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array Int a -> f (Array Int b)
traverse CacheState t -> IO (IORef (CacheState t))
forall a. a -> IO (IORef a)
newIORef
          (Array Int (CacheState t) -> IO (Array Int (IORef (CacheState t))))
-> Array Int (CacheState t)
-> IO (Array Int (IORef (CacheState t)))
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> [CacheState t] -> Array Int (CacheState t)
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0,Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          ([CacheState t] -> Array Int (CacheState t))
-> [CacheState t] -> Array Int (CacheState t)
forall a b. (a -> b) -> a -> b
$ Int -> CacheState t -> [CacheState t]
forall a. Int -> a -> [a]
replicate Int
w CacheState t
element

type Id = Int

class ( Eq (Description t)
      , Hashable (Description t)
      ) => Interned t where
  data Description t
  type Uninterned t
  describe :: Uninterned t -> Description t
  identify :: Id -> Uninterned t -> t
  -- identity :: t -> Id
  seedIdentity :: p t -> Id
  seedIdentity p t
_ = Int
0
  cacheWidth :: p t -> Int
  cacheWidth p t
_ = Int
defaultCacheWidth
  modifyAdvice :: IO t -> IO t
  modifyAdvice = IO t -> IO t
forall a. a -> a
id
  cache        :: Cache t

class Interned t => Uninternable t where
  unintern :: t -> Uninterned t

intern :: Interned t => Uninterned t -> t
intern :: forall t. Interned t => Uninterned t -> t
intern !Uninterned t
bt = IO t -> t
forall a. IO a -> a
unsafeDupablePerformIO (IO t -> t) -> IO t -> t
forall a b. (a -> b) -> a -> b
$ IO t -> IO t
forall t. Interned t => IO t -> IO t
modifyAdvice (IO t -> IO t) -> IO t -> IO t
forall a b. (a -> b) -> a -> b
$ IORef (CacheState t) -> (CacheState t -> (CacheState t, t)) -> IO t
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (CacheState t)
slot CacheState t -> (CacheState t, t)
go
  where
  slot :: IORef (CacheState t)
slot = Cache t -> Array Int (IORef (CacheState t))
forall t. Cache t -> Array Int (IORef (CacheState t))
getCache Cache t
forall t. Interned t => Cache t
cache Array Int (IORef (CacheState t)) -> Int -> IORef (CacheState t)
forall i e. Ix i => Array i e -> i -> e
! Int
r
  !dt :: Description t
dt = Uninterned t -> Description t
forall t. Interned t => Uninterned t -> Description t
describe Uninterned t
bt
  !hdt :: Int
hdt = Description t -> Int
forall a. Hashable a => a -> Int
hash Description t
dt
  !wid :: Int
wid = Description t -> Int
forall t (p :: * -> *). Interned t => p t -> Int
forall (p :: * -> *). p t -> Int
cacheWidth Description t
dt
  r :: Int
r = Int
hdt Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
wid
  go :: CacheState t -> (CacheState t, t)
go (CacheState Int
i HashMap (Description t) t
m) = case Description t -> HashMap (Description t) t -> Maybe t
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Description t
dt HashMap (Description t) t
m of
    Maybe t
Nothing -> let t :: t
t = Int -> Uninterned t -> t
forall t. Interned t => Int -> Uninterned t -> t
identify (Int
wid Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Uninterned t
bt in (Int -> HashMap (Description t) t -> CacheState t
forall t. Int -> HashMap (Description t) t -> CacheState t
CacheState (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Description t
-> t -> HashMap (Description t) t -> HashMap (Description t) t
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert Description t
dt t
t HashMap (Description t) t
m), t
t)
    Just t
t -> (Int -> HashMap (Description t) t -> CacheState t
forall t. Int -> HashMap (Description t) t -> CacheState t
CacheState Int
i HashMap (Description t) t
m, t
t)

-- given a description, go hunting for an entry in the cache
recover :: Interned t => Description t -> IO (Maybe t)
recover :: forall t. Interned t => Description t -> IO (Maybe t)
recover !Description t
dt = do
  CacheState Int
_ HashMap (Description t) t
m <- IORef (CacheState t) -> IO (CacheState t)
forall a. IORef a -> IO a
readIORef (IORef (CacheState t) -> IO (CacheState t))
-> IORef (CacheState t) -> IO (CacheState t)
forall a b. (a -> b) -> a -> b
$ Cache t -> Array Int (IORef (CacheState t))
forall t. Cache t -> Array Int (IORef (CacheState t))
getCache Cache t
forall t. Interned t => Cache t
cache Array Int (IORef (CacheState t)) -> Int -> IORef (CacheState t)
forall i e. Ix i => Array i e -> i -> e
! (Description t -> Int
forall a. Hashable a => a -> Int
hash Description t
dt Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Description t -> Int
forall t (p :: * -> *). Interned t => p t -> Int
forall (p :: * -> *). p t -> Int
cacheWidth Description t
dt)
  Maybe t -> IO (Maybe t)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe t -> IO (Maybe t)) -> Maybe t -> IO (Maybe t)
forall a b. (a -> b) -> a -> b
$ Description t -> HashMap (Description t) t -> Maybe t
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Description t
dt HashMap (Description t) t
m