    {-# 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