{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE BangPatterns, NoImplicitPrelude, RecordWildCards #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}

module GHC.Event.IntTable
    (
      IntTable
    , new
    , lookup
    , insertWith
    , reset
    , delete
    , updateWith
    ) where

import Data.Bits ((.&.), shiftL, shiftR)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Maybe (Maybe(..), isJust)
import GHC.Base (Monad(..), (=<<), ($), ($!), const, liftM, otherwise, when)
import GHC.Classes (Eq(..), Ord(..))
import GHC.Event.Arr (Arr)
import GHC.Event.IntVar
import GHC.Num (Num(..))
import GHC.Prim (seq)
import GHC.Types (Bool(..), IO(..), Int(..))
import qualified GHC.Event.Arr as Arr

-- A very simple chained integer-keyed mutable hash table. We use
-- power-of-two sizing, grow at a load factor of 0.75, and never
-- shrink. The "hash function" is the identity function.

newtype IntTable a = IntTable (IORef (IT a))

data IT a = IT {
      forall a. IT a -> Arr (Bucket a)
tabArr  :: {-# UNPACK #-} !(Arr (Bucket a))
    , forall a. IT a -> IntVar
tabSize :: {-# UNPACK #-} !IntVar
    }

data Bucket a = Empty
              | Bucket {
      forall a. Bucket a -> Int
bucketKey   :: {-# UNPACK #-} !Int
    , forall a. Bucket a -> a
bucketValue :: a
    , forall a. Bucket a -> Bucket a
bucketNext  :: Bucket a
    }

lookup :: Int -> IntTable a -> IO (Maybe a)
lookup :: forall a. Int -> IntTable a -> IO (Maybe a)
lookup Int
k (IntTable IORef (IT a)
ref) = do
  let go :: Bucket a -> Maybe a
go Bucket{a
Int
Bucket a
bucketKey :: forall a. Bucket a -> Int
bucketValue :: forall a. Bucket a -> a
bucketNext :: forall a. Bucket a -> Bucket a
bucketKey :: Int
bucketValue :: a
bucketNext :: Bucket a
..}
        | Int
bucketKey Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k = a -> Maybe a
forall a. a -> Maybe a
Just a
bucketValue
        | Bool
otherwise      = Bucket a -> Maybe a
go Bucket a
bucketNext
      go Bucket a
_ = Maybe a
forall a. Maybe a
Nothing
  it :: IT a
it@IT{IntVar
Arr (Bucket a)
tabArr :: forall a. IT a -> Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: Arr (Bucket a)
tabSize :: IntVar
..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  Bucket a
bkt <- Arr (Bucket a) -> Int -> IO (Bucket a)
forall a. Arr a -> Int -> IO a
Arr.read Arr (Bucket a)
tabArr (Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
k IT a
it)
  Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$! Bucket a -> Maybe a
forall {a}. Bucket a -> Maybe a
go Bucket a
bkt

new :: Int -> IO (IntTable a)
new :: forall a. Int -> IO (IntTable a)
new Int
capacity = IORef (IT a) -> IntTable a
forall a. IORef (IT a) -> IntTable a
IntTable (IORef (IT a) -> IntTable a)
-> IO (IORef (IT a)) -> IO (IntTable a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` (IT a -> IO (IORef (IT a))
forall a. a -> IO (IORef a)
newIORef (IT a -> IO (IORef (IT a))) -> IO (IT a) -> IO (IORef (IT a))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> IO (IT a)
forall a. Int -> IO (IT a)
new_ Int
capacity)

new_ :: Int -> IO (IT a)
new_ :: forall a. Int -> IO (IT a)
new_ Int
capacity = do
  Arr (Bucket a)
arr <- Bucket a -> Int -> IO (Arr (Bucket a))
forall a. a -> Int -> IO (Arr a)
Arr.new Bucket a
forall a. Bucket a
Empty Int
capacity
  IntVar
size <- Int -> IO IntVar
newIntVar Int
0
  IT a -> IO (IT a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IT { tabArr :: Arr (Bucket a)
tabArr = Arr (Bucket a)
arr
            , tabSize :: IntVar
tabSize = IntVar
size
            }

grow :: IT a -> IORef (IT a) -> Int -> IO ()
grow :: forall a. IT a -> IORef (IT a) -> Int -> IO ()
grow IT a
oldit IORef (IT a)
ref Int
size = do
  IT a
newit <- Int -> IO (IT a)
forall a. Int -> IO (IT a)
new_ (Arr (Bucket a) -> Int
forall a. Arr a -> Int
Arr.size (IT a -> Arr (Bucket a)
forall a. IT a -> Arr (Bucket a)
tabArr IT a
oldit) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
1)
  let copySlot :: Int -> Int -> IO ()
copySlot Int
n !Int
i
        | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
size = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
          let copyBucket :: Int -> Bucket a -> IO ()
copyBucket !Int
m Bucket a
Empty          = Int -> Int -> IO ()
copySlot Int
m (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
              copyBucket  Int
m bkt :: Bucket a
bkt@Bucket{a
Int
Bucket a
bucketKey :: forall a. Bucket a -> Int
bucketValue :: forall a. Bucket a -> a
bucketNext :: forall a. Bucket a -> Bucket a
bucketKey :: Int
bucketValue :: a
bucketNext :: Bucket a
..} = do
                let idx :: Int
idx = Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
bucketKey IT a
newit
                Bucket a
next <- Arr (Bucket a) -> Int -> IO (Bucket a)
forall a. Arr a -> Int -> IO a
Arr.read (IT a -> Arr (Bucket a)
forall a. IT a -> Arr (Bucket a)
tabArr IT a
newit) Int
idx
                Arr (Bucket a) -> Int -> Bucket a -> IO ()
forall a. Arr a -> Int -> a -> IO ()
Arr.write (IT a -> Arr (Bucket a)
forall a. IT a -> Arr (Bucket a)
tabArr IT a
newit) Int
idx Bucket a
bkt { bucketNext :: Bucket a
bucketNext = Bucket a
next }
                Int -> Bucket a -> IO ()
copyBucket (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Bucket a
bucketNext
          Int -> Bucket a -> IO ()
copyBucket Int
n (Bucket a -> IO ()) -> IO (Bucket a) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Arr (Bucket a) -> Int -> IO (Bucket a)
forall a. Arr a -> Int -> IO a
Arr.read (IT a -> Arr (Bucket a)
forall a. IT a -> Arr (Bucket a)
tabArr IT a
oldit) Int
i
  Int -> Int -> IO ()
copySlot Int
0 Int
0
  IntVar -> Int -> IO ()
writeIntVar (IT a -> IntVar
forall a. IT a -> IntVar
tabSize IT a
newit) Int
size
  IORef (IT a) -> IT a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IT a)
ref IT a
newit

-- | @insertWith f k v table@ inserts @k@ into @table@ with value @v@.
-- If @k@ already appears in @table@ with value @v0@, the value is updated
-- to @f v0 v@ and @Just v0@ is returned.
insertWith :: (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith :: forall a. (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith a -> a -> a
f Int
k a
v inttable :: IntTable a
inttable@(IntTable IORef (IT a)
ref) = do
  it :: IT a
it@IT{IntVar
Arr (Bucket a)
tabArr :: forall a. IT a -> Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: Arr (Bucket a)
tabSize :: IntVar
..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx :: Int
idx = Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
k IT a
it
      go :: Bucket a -> Bucket a -> IO (Maybe a)
go Bucket a
seen bkt :: Bucket a
bkt@Bucket{a
Int
Bucket a
bucketKey :: forall a. Bucket a -> Int
bucketValue :: forall a. Bucket a -> a
bucketNext :: forall a. Bucket a -> Bucket a
bucketKey :: Int
bucketValue :: a
bucketNext :: Bucket a
..}
        | Int
bucketKey Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k = do
          let !v' :: a
v' = a -> a -> a
f a
v a
bucketValue
              !next :: Bucket a
next = Bucket a
seen Bucket a -> Bucket a -> Bucket a
forall {a}. Bucket a -> Bucket a -> Bucket a
<> Bucket a
bucketNext
              Bucket a
Empty        <> :: Bucket a -> Bucket a -> Bucket a
<> Bucket a
bs = Bucket a
bs
              b :: Bucket a
b@Bucket{a
Int
Bucket a
bucketKey :: forall a. Bucket a -> Int
bucketValue :: forall a. Bucket a -> a
bucketNext :: forall a. Bucket a -> Bucket a
bucketKey :: Int
bucketValue :: a
bucketNext :: Bucket a
..} <> Bucket a
bs = Bucket a
b { bucketNext :: Bucket a
bucketNext = Bucket a
bucketNext Bucket a -> Bucket a -> Bucket a
<> Bucket a
bs }
          Arr (Bucket a) -> Int -> Bucket a -> IO ()
forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx (Int -> a -> Bucket a -> Bucket a
forall a. Int -> a -> Bucket a -> Bucket a
Bucket Int
k a
v' Bucket a
next)
          Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
bucketValue)
        | Bool
otherwise = Bucket a -> Bucket a -> IO (Maybe a)
go Bucket a
bkt { bucketNext :: Bucket a
bucketNext = Bucket a
seen } Bucket a
bucketNext
      go Bucket a
seen Bucket a
_ = do
        Int
size <- IntVar -> IO Int
readIntVar IntVar
tabSize
        if Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Arr (Bucket a) -> Int
forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Arr (Bucket a) -> Int
forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
2)
          then IT a -> IORef (IT a) -> Int -> IO ()
forall a. IT a -> IORef (IT a) -> Int -> IO ()
grow IT a
it IORef (IT a)
ref Int
size IO () -> IO (Maybe a) -> IO (Maybe a)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
forall a. (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith a -> a -> a
f Int
k a
v IntTable a
inttable
          else do
            a
v a -> IO () -> IO ()
forall a b. a -> b -> b
`seq` Arr (Bucket a) -> Int -> Bucket a -> IO ()
forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx (Int -> a -> Bucket a -> Bucket a
forall a. Int -> a -> Bucket a -> Bucket a
Bucket Int
k a
v Bucket a
seen)
            IntVar -> Int -> IO ()
writeIntVar IntVar
tabSize (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
  Bucket a -> Bucket a -> IO (Maybe a)
go Bucket a
forall a. Bucket a
Empty (Bucket a -> IO (Maybe a)) -> IO (Bucket a) -> IO (Maybe a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Arr (Bucket a) -> Int -> IO (Bucket a)
forall a. Arr a -> Int -> IO a
Arr.read Arr (Bucket a)
tabArr Int
idx
{-# INLINABLE insertWith #-}

-- | Used to undo the effect of a prior insertWith.
reset :: Int -> Maybe a -> IntTable a -> IO ()
reset :: forall a. Int -> Maybe a -> IntTable a -> IO ()
reset Int
k (Just a
v) IntTable a
tbl = (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
forall a. (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith a -> a -> a
forall a b. a -> b -> a
const Int
k a
v IntTable a
tbl IO (Maybe a) -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
reset Int
k Maybe a
Nothing  IntTable a
tbl = Int -> IntTable a -> IO (Maybe a)
forall a. Int -> IntTable a -> IO (Maybe a)
delete Int
k IntTable a
tbl IO (Maybe a) -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

indexOf :: Int -> IT a -> Int
indexOf :: forall a. Int -> IT a -> Int
indexOf Int
k IT{IntVar
Arr (Bucket a)
tabArr :: forall a. IT a -> Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: Arr (Bucket a)
tabSize :: IntVar
..} = Int
k Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (Arr (Bucket a) -> Int
forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Remove the given key from the table and return its associated value.
delete :: Int -> IntTable a -> IO (Maybe a)
delete :: forall a. Int -> IntTable a -> IO (Maybe a)
delete Int
k IntTable a
t = (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
forall a. (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
updateWith (Maybe a -> a -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) Int
k IntTable a
t

updateWith :: (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
updateWith :: forall a. (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
updateWith a -> Maybe a
f Int
k (IntTable IORef (IT a)
ref) = do
  it :: IT a
it@IT{IntVar
Arr (Bucket a)
tabArr :: forall a. IT a -> Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: Arr (Bucket a)
tabSize :: IntVar
..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx :: Int
idx = Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
k IT a
it
      go :: Bucket a -> (Bool, Maybe a, Bucket a)
go bkt :: Bucket a
bkt@Bucket{a
Int
Bucket a
bucketKey :: forall a. Bucket a -> Int
bucketValue :: forall a. Bucket a -> a
bucketNext :: forall a. Bucket a -> Bucket a
bucketKey :: Int
bucketValue :: a
bucketNext :: Bucket a
..}
        | Int
bucketKey Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k = case a -> Maybe a
f a
bucketValue of
            Just a
val -> let !nb :: Bucket a
nb = Bucket a
bkt { bucketValue :: a
bucketValue = a
val }
                        in (Bool
False, a -> Maybe a
forall a. a -> Maybe a
Just a
bucketValue, Bucket a
nb)
            Maybe a
Nothing  -> (Bool
True, a -> Maybe a
forall a. a -> Maybe a
Just a
bucketValue, Bucket a
bucketNext)
        | Bool
otherwise = case Bucket a -> (Bool, Maybe a, Bucket a)
go Bucket a
bucketNext of
                        (Bool
fbv, Maybe a
ov, Bucket a
nb) -> (Bool
fbv, Maybe a
ov, Bucket a
bkt { bucketNext :: Bucket a
bucketNext = Bucket a
nb })
      go Bucket a
e = (Bool
False, Maybe a
forall a. Maybe a
Nothing, Bucket a
e)
  (Bool
del, Maybe a
oldVal, Bucket a
newBucket) <- Bucket a -> (Bool, Maybe a, Bucket a)
go (Bucket a -> (Bool, Maybe a, Bucket a))
-> IO (Bucket a) -> IO (Bool, Maybe a, Bucket a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Arr (Bucket a) -> Int -> IO (Bucket a)
forall a. Arr a -> Int -> IO a
Arr.read Arr (Bucket a)
tabArr Int
idx
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
oldVal) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    Arr (Bucket a) -> Int -> Bucket a -> IO ()
forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx Bucket a
newBucket
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
del (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      Int
size <- IntVar -> IO Int
readIntVar IntVar
tabSize
      IntVar -> Int -> IO ()
writeIntVar IntVar
tabSize (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
oldVal