{-# 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
bucketNext :: Bucket a
bucketValue :: a
bucketKey :: Int
bucketNext :: forall a. Bucket a -> Bucket a
bucketValue :: forall a. Bucket a -> a
bucketKey :: forall a. Bucket a -> Int
..}
        | Int
bucketKey forall a. Eq a => a -> a -> Bool
== Int
k = forall a. a -> Maybe a
Just a
bucketValue
        | Bool
otherwise      = Bucket a -> Maybe a
go Bucket a
bucketNext
      go Bucket a
_ = forall a. Maybe a
Nothing
  it :: IT a
it@IT{Arr (Bucket a)
IntVar
tabSize :: IntVar
tabArr :: Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: forall a. IT a -> Arr (Bucket a)
..} <- forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  Bucket a
bkt <- forall a. Arr a -> Int -> IO a
Arr.read Arr (Bucket a)
tabArr (forall a. Int -> IT a -> Int
indexOf Int
k IT a
it)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! 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 = forall a. IORef (IT a) -> IntTable a
IntTable forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` (forall a. a -> IO (IORef a)
newIORef forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 <- forall a. a -> Int -> IO (Arr a)
Arr.new forall a. Bucket a
Empty Int
capacity
  IntVar
size <- Int -> IO IntVar
newIntVar Int
0
  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 <- forall a. Int -> IO (IT a)
new_ (forall a. Arr a -> Int
Arr.size (forall a. IT a -> Arr (Bucket a)
tabArr IT a
oldit) forall a. Bits a => a -> Int -> a
`shiftL` Int
1)
  let copySlot :: Int -> Int -> IO ()
copySlot Int
n !Int
i
        | Int
n forall a. Eq a => a -> a -> Bool
== Int
size = 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
iforall a. Num a => a -> a -> a
+Int
1)
              copyBucket  Int
m bkt :: Bucket a
bkt@Bucket{a
Int
Bucket a
bucketNext :: Bucket a
bucketValue :: a
bucketKey :: Int
bucketNext :: forall a. Bucket a -> Bucket a
bucketValue :: forall a. Bucket a -> a
bucketKey :: forall a. Bucket a -> Int
..} = do
                let idx :: Int
idx = forall a. Int -> IT a -> Int
indexOf Int
bucketKey IT a
newit
                Bucket a
next <- forall a. Arr a -> Int -> IO a
Arr.read (forall a. IT a -> Arr (Bucket a)
tabArr IT a
newit) Int
idx
                forall a. Arr a -> Int -> a -> IO ()
Arr.write (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
mforall a. Num a => a -> a -> a
+Int
1) Bucket a
bucketNext
          Int -> Bucket a -> IO ()
copyBucket Int
n forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. Arr a -> Int -> IO a
Arr.read (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 (forall a. IT a -> IntVar
tabSize IT a
newit) Int
size
  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{Arr (Bucket a)
IntVar
tabSize :: IntVar
tabArr :: Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: forall a. IT a -> Arr (Bucket a)
..} <- forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx :: Int
idx = 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
bucketNext :: Bucket a
bucketValue :: a
bucketKey :: Int
bucketNext :: forall a. Bucket a -> Bucket a
bucketValue :: forall a. Bucket a -> a
bucketKey :: forall a. Bucket a -> Int
..}
        | Int
bucketKey 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 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
bucketNext :: Bucket a
bucketValue :: a
bucketKey :: Int
bucketNext :: forall a. Bucket a -> Bucket a
bucketValue :: forall a. Bucket a -> a
bucketKey :: forall a. Bucket a -> Int
..} <> Bucket a
bs = Bucket a
b { bucketNext :: Bucket a
bucketNext = Bucket a
bucketNext Bucket a -> Bucket a -> Bucket a
<> Bucket a
bs }
          forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx (forall a. Int -> a -> Bucket a -> Bucket a
Bucket Int
k a
v' Bucket a
next)
          forall (m :: * -> *) a. Monad m => a -> m a
return (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 forall a. Num a => a -> a -> a
+ Int
1 forall a. Ord a => a -> a -> Bool
>= forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr forall a. Num a => a -> a -> a
- (forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr forall a. Bits a => a -> Int -> a
`shiftR` Int
2)
          then forall a. IT a -> IORef (IT a) -> Int -> IO ()
grow IT a
it IORef (IT a)
ref Int
size forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> 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 seq :: forall a b. a -> b -> b
`seq` forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx (forall a. Int -> a -> Bucket a -> Bucket a
Bucket Int
k a
v Bucket a
seen)
            IntVar -> Int -> IO ()
writeIntVar IntVar
tabSize (Int
size forall a. Num a => a -> a -> a
+ Int
1)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
  Bucket a -> Bucket a -> IO (Maybe a)
go forall a. Bucket a
Empty forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 = forall a. (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith forall a b. a -> b -> a
const Int
k a
v IntTable a
tbl forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return ()
reset Int
k Maybe a
Nothing  IntTable a
tbl = forall a. Int -> IntTable a -> IO (Maybe a)
delete Int
k IntTable a
tbl forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> 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{Arr (Bucket a)
IntVar
tabSize :: IntVar
tabArr :: Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: forall a. IT a -> Arr (Bucket a)
..} = Int
k forall a. Bits a => a -> a -> a
.&. (forall a. Arr a -> Int
Arr.size Arr (Bucket a)
tabArr 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 = forall a. (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
updateWith (forall a b. a -> b -> a
const 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{Arr (Bucket a)
IntVar
tabSize :: IntVar
tabArr :: Arr (Bucket a)
tabSize :: forall a. IT a -> IntVar
tabArr :: forall a. IT a -> Arr (Bucket a)
..} <- forall a. IORef a -> IO a
readIORef IORef (IT a)
ref
  let idx :: Int
idx = 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
bucketNext :: Bucket a
bucketValue :: a
bucketKey :: Int
bucketNext :: forall a. Bucket a -> Bucket a
bucketValue :: forall a. Bucket a -> a
bucketKey :: forall a. Bucket a -> Int
..}
        | Int
bucketKey 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, forall a. a -> Maybe a
Just a
bucketValue, Bucket a
nb)
            Maybe a
Nothing  -> (Bool
True, 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, forall a. Maybe a
Nothing, Bucket a
e)
  (Bool
del, Maybe a
oldVal, Bucket a
newBucket) <- Bucket a -> (Bool, Maybe a, Bucket a)
go forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall a. Arr a -> Int -> IO a
Arr.read Arr (Bucket a)
tabArr Int
idx
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Maybe a -> Bool
isJust Maybe a
oldVal) forall a b. (a -> b) -> a -> b
$ do
    forall a. Arr a -> Int -> a -> IO ()
Arr.write Arr (Bucket a)
tabArr Int
idx Bucket a
newBucket
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
del forall a b. (a -> b) -> a -> b
$ do
      Int
size <- IntVar -> IO Int
readIntVar IntVar
tabSize
      IntVar -> Int -> IO ()
writeIntVar IntVar
tabSize (Int
size forall a. Num a => a -> a -> a
- Int
1)
  forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
oldVal