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

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 Foreign.ForeignPtr (ForeignPtr, mallocForeignPtr, withForeignPtr)
import Foreign.Storable (peek, poke)
import GHC.Base (Monad(..), (=<<), ($), ($!), const, liftM, otherwise, when)
import GHC.Classes (Eq(..), Ord(..))
import GHC.Event.Arr (Arr)
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 {
      IT a -> Arr (Bucket a)
tabArr  :: {-# UNPACK #-} !(Arr (Bucket a))
    , IT a -> ForeignPtr Int
tabSize :: {-# UNPACK #-} !(ForeignPtr Int)
    }

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

lookup :: Int -> IntTable a -> IO (Maybe a)
lookup :: Int -> IntTable a -> IO (Maybe a)
lookup k :: Int
k (IntTable ref :: IORef (IT a)
ref) = do
  let go :: Bucket a -> Maybe a
go Bucket{..}
        | 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 _ = Maybe a
forall a. Maybe a
Nothing
  it :: IT a
it@IT{..} <- 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 (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 :: Int -> IO (IntTable a)
new capacity :: 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_ :: Int -> IO (IT a)
new_ capacity :: 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
  ForeignPtr Int
size <- IO (ForeignPtr Int)
forall a. Storable a => IO (ForeignPtr a)
mallocForeignPtr
  ForeignPtr Int -> (Ptr Int -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Int
size ((Ptr Int -> IO ()) -> IO ()) -> (Ptr Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ptr :: Ptr Int
ptr -> Ptr Int -> Int -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Int
ptr 0
  IT a -> IO (IT a)
forall (m :: * -> *) a. Monad m => a -> m a
return $WIT :: forall a. Arr (Bucket a) -> ForeignPtr Int -> IT a
IT { tabArr :: Arr (Bucket a)
tabArr = Arr (Bucket a)
arr
            , tabSize :: ForeignPtr Int
tabSize = ForeignPtr Int
size
            }

grow :: IT a -> IORef (IT a) -> Int -> IO ()
grow :: IT a -> IORef (IT a) -> Int -> IO ()
grow oldit :: IT a
oldit ref :: IORef (IT a)
ref size :: 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` 1)
  let copySlot :: Int -> Int -> IO ()
copySlot n :: Int
n !Int
i
        | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
size = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise = do
          let copyBucket :: Int -> Bucket a -> IO ()
copyBucket !Int
m Empty          = Int -> Int -> IO ()
copySlot Int
m (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
              copyBucket  m :: Int
m bkt :: Bucket a
bkt@Bucket{..} = 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
+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 0 0
  ForeignPtr Int -> (Ptr Int -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (IT a -> ForeignPtr Int
forall a. IT a -> ForeignPtr Int
tabSize IT a
newit) ((Ptr Int -> IO ()) -> IO ()) -> (Ptr Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ptr :: Ptr Int
ptr -> Ptr Int -> Int -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Int
ptr 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 :: (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
insertWith f :: a -> a -> a
f k :: Int
k v :: a
v inttable :: IntTable a
inttable@(IntTable ref :: IORef (IT a)
ref) = do
  it :: IT a
it@IT{..} <- 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 seen :: Bucket a
seen bkt :: Bucket a
bkt@Bucket{..}
        | 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
              Empty        <> :: Bucket a -> Bucket a -> Bucket a
<> bs :: Bucket a
bs = Bucket a
bs
              b :: Bucket a
b@Bucket{..} <> bs :: 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 (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 seen :: Bucket a
seen _ = ForeignPtr Int -> (Ptr Int -> IO (Maybe a)) -> IO (Maybe a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Int
tabSize ((Ptr Int -> IO (Maybe a)) -> IO (Maybe a))
-> (Ptr Int -> IO (Maybe a)) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \ptr :: Ptr Int
ptr -> do
        Int
size <- Ptr Int -> IO Int
forall a. Storable a => Ptr a -> IO a
peek Ptr Int
ptr
        if Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 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` 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 (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)
            Ptr Int -> Int -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Int
ptr (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1)
            Maybe a -> IO (Maybe 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 :: Int -> Maybe a -> IntTable a -> IO ()
reset k :: Int
k (Just v :: a
v) tbl :: 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 (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
reset k :: Int
k Nothing  tbl :: 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 (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

indexOf :: Int -> IT a -> Int
indexOf :: Int -> IT a -> Int
indexOf k :: Int
k IT{..} = 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
- 1)

-- | Remove the given key from the table and return its associated value.
delete :: Int -> IntTable a -> IO (Maybe a)
delete :: Int -> IntTable a -> IO (Maybe a)
delete k :: Int
k t :: 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 :: (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
updateWith f :: a -> Maybe a
f k :: Int
k (IntTable ref :: IORef (IT a)
ref) = do
  it :: IT a
it@IT{..} <- 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{..}
        | Int
bucketKey Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k = case a -> Maybe a
f a
bucketValue of
            Just val :: 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)
            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
                        (fbv :: Bool
fbv, ov :: Maybe a
ov, nb :: Bucket a
nb) -> (Bool
fbv, Maybe a
ov, Bucket a
bkt { bucketNext :: Bucket a
bucketNext = Bucket a
nb })
      go e :: Bucket a
e = (Bool
False, Maybe a
forall a. Maybe a
Nothing, Bucket a
e)
  (del :: Bool
del, oldVal :: Maybe a
oldVal, newBucket :: 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
$
      ForeignPtr Int -> (Ptr Int -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Int
tabSize ((Ptr Int -> IO ()) -> IO ()) -> (Ptr Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ptr :: Ptr Int
ptr -> do
        Int
size <- Ptr Int -> IO Int
forall a. Storable a => Ptr a -> IO a
peek Ptr Int
ptr
        Ptr Int -> Int -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Int
ptr (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
  Maybe a -> IO (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
oldVal