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

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

import GHC.Internal.Data.Bits ((.&.), shiftL, shiftR)
import GHC.Internal.Data.IORef (IORef, newIORef, readIORef, writeIORef)
import GHC.Internal.Data.Maybe (Maybe(..), isJust)
import GHC.Internal.Base (Monad(..), (=<<), ($), ($!), const, liftM, otherwise, when)
import GHC.Classes (Eq(..), Ord(..))
import GHC.Internal.Event.Arr (Arr)
import GHC.Internal.Event.IntVar
import GHC.Internal.Num (Num(..))
import GHC.Prim (seq)
import GHC.Types (Bool(..), IO(..), Int(..))
import qualified GHC.Internal.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{..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  bkt <- Arr.read tabArr (indexOf k it)
  return $! go 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 -> Int -> IO (Arr (Bucket a))
forall a. a -> Int -> IO (Arr a)
Arr.new Bucket a
forall a. Bucket a
Empty Int
capacity
  size <- newIntVar 0
  return IT { tabArr = arr
            , tabSize = 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
  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
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
                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.write (tabArr newit) idx bkt { bucketNext = next }
                copyBucket (m+1) 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
  copySlot 0 0
  writeIntVar (tabSize newit) size
  writeIORef ref 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{..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx = Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
k IT a
it
      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 = bucketNext <> 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 = seen } Bucket a
bucketNext
      go Bucket a
seen Bucket a
_ = do
        size <- IntVar -> IO Int
readIntVar IntVar
tabSize
        if size + 1 >= Arr.size tabArr - (Arr.size tabArr `shiftR` 2)
          then grow it ref size >> insertWith f k v inttable
          else do
            v `seq` Arr.write tabArr idx (Bucket k v seen)
            writeIntVar tabSize (size + 1)
            return Nothing
  go Empty =<< Arr.read tabArr 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{..} <- IORef (IT a) -> IO (IT a)
forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx = Int -> IT a -> Int
forall a. Int -> IT a -> Int
indexOf Int
k IT a
it
      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 = 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 = nb })
      go Bucket a
e = (Bool
False, Maybe a
forall a. Maybe a
Nothing, Bucket a
e)
  (del, oldVal, newBucket) <- go `liftM` Arr.read tabArr idx
  when (isJust oldVal) $ do
    Arr.write tabArr idx newBucket
    when del $ do
      size <- readIntVar tabSize
      writeIntVar tabSize (size - 1)
  return oldVal