    {-# 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 = 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 = 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
            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 = 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)
      (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