Ticket #8793: IntTable.diff

File IntTable.diff, 8.4 KB (added by cdk, 19 months ago)
  • GHC/Event/IntTable.hs

    diff --git a/GHC/Event/IntTable.hs b/GHC/Event/IntTable.hs
    index d8cbcc0..835df05 100644
    a b  
    1 {-# LANGUAGE BangPatterns, NoImplicitPrelude, RecordWildCards, Trustworthy #-}
    2 {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
     1{-# LANGUAGE BangPatterns, MagicHash, NoImplicitPrelude, RecordWildCards #-}
    32
    43module GHC.Event.IntTable
    5     (
    6       IntTable
     4    ( IntTable
    75    , new
    86    , lookup
    97    , insertWith
    module GHC.Event.IntTable 
    1210    , updateWith
    1311    ) where
    1412
    15 import Control.Monad ((=<<), liftM, unless, when)
    16 import Data.Bits ((.&.), shiftL, shiftR)
    17 import Data.IORef (IORef, newIORef, readIORef, writeIORef)
    18 import Data.Maybe (Maybe(..), isJust)
    19 import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtr, withForeignPtr)
    20 import Foreign.Storable (peek, poke)
    21 import GHC.Base (Monad(..), ($), const, otherwise)
    22 import GHC.Classes (Eq(..), Ord(..))
     13import Prelude (($!))
     14import Control.Monad
     15import Data.Bits
     16import Data.IORef
     17import Data.Maybe
     18import Foreign.ForeignPtr
     19import Foreign.Storable
    2320import GHC.Event.Arr (Arr)
    24 import GHC.Num (Num(..))
    25 import GHC.Prim (seq)
    26 import GHC.Types (Bool(..), IO(..), Int(..))
    27 import qualified GHC.Event.Arr as Arr
     21import GHC.Base
     22import GHC.Num
     23import qualified GHC.Event.Arr as A
    2824
    29 -- A very simple chained integer-keyed mutable hash table. We use
    30 -- power-of-two sizing, grow at a load factor of 0.75, and never
    31 -- shrink. The "hash function" is the identity function.
     25-- A simple chained int-keyed mutable hash table.
     26-- Uses power-of-two sizing, grows at a load factor of 0.75,
     27-- and never shrinks. The "hash function" is the identity function.
    3228
    33 newtype IntTable a = IntTable (IORef (IT a))
    34 
    35 data IT a = IT {
    36       tabArr  :: {-# UNPACK #-} !(Arr (Bucket a))
    37     , tabSize :: {-# UNPACK #-} !(ForeignPtr Int)
     29data IntTable a = IntTable
     30    { tblRef  :: {-# UNPACK #-} !(IORef (Arr (Bucket a)))
     31    , tblSize :: {-# UNPACK #-} !(ForeignPtr Int)
    3832    }
    3933
    40 data Bucket a = Empty
    41               | Bucket {
    42       bucketKey   :: {-# UNPACK #-} !Int
    43     , bucketValue :: a
    44     , bucketNext  :: Bucket a
     34data Bucket a
     35    = Empty
     36    | Bucket
     37    { bucketKey  :: {-# UNPACK #-} !Int
     38    , bucketVal  :: a
     39    , bucketNext :: Bucket a
    4540    }
    4641
     42indexOf :: Int -> Arr a -> Int
     43indexOf k arr = k .&. (A.size arr - 1)
     44
    4745lookup :: Int -> IntTable a -> IO (Maybe a)
    48 lookup k (IntTable ref) = do
    49   let go Bucket{..}
    50         | bucketKey == k = return (Just bucketValue)
    51         | otherwise      = go bucketNext
    52       go _ = return Nothing
    53   it@IT{..} <- readIORef ref
    54   go =<< Arr.read tabArr (indexOf k it)
     46lookup k IntTable{..} = do
     47    arr <- readIORef tblRef
     48    bkt <- A.read arr (indexOf k arr)
     49    return (go bkt)
     50    where go Empty = Nothing
     51          go Bucket{..}
     52            | bucketKey == k = Just bucketVal
     53            | otherwise      = go bucketNext
    5554
    5655new :: Int -> IO (IntTable a)
    57 new capacity = IntTable `liftM` (newIORef =<< new_ capacity)
    58 
    59 new_ :: Int -> IO (IT a)
    60 new_ capacity = do
    61   arr <- Arr.new Empty capacity
    62   size <- mallocForeignPtr
    63   withForeignPtr size $ \ptr -> poke ptr 0
    64   return IT { tabArr = arr
    65             , tabSize = size
    66             }
     56new size = do
     57    arr <- A.new Empty size
     58    ref <- newIORef arr
     59    ptr <- mallocForeignPtr
     60    withForeignPtr ptr (`poke` 0)
     61    return (IntTable ref ptr)
    6762
    68 grow :: IT a -> IORef (IT a) -> Int -> IO ()
    69 grow oldit ref size = do
    70   newit <- new_ (Arr.size (tabArr oldit) `shiftL` 1)
    71   let copySlot n !i
    72         | n == size = return ()
    73         | otherwise = do
    74           let copyBucket !m Empty          = copySlot m (i+1)
    75               copyBucket  m bkt@Bucket{..} = do
    76                 let idx = indexOf bucketKey newit
    77                 next <- Arr.read (tabArr newit) idx
    78                 Arr.write (tabArr newit) idx bkt { bucketNext = next }
    79                 copyBucket (m+1) bucketNext
    80           copyBucket n =<< Arr.read (tabArr oldit) i
    81   copySlot 0 0
    82   withForeignPtr (tabSize newit) $ \ptr -> poke ptr size
    83   writeIORef ref newit
     63copy :: Arr (Bucket a) -> Arr (Bucket a) -> Int -> IO ()
     64copy src dst size = A.read src 0 >>= go 1 0
     65    where go i n Empty
     66            | n == size = return ()
     67            | otherwise = A.read src i >>= go (i + 1) n
     68          go i n bkt@Bucket{..} = do
     69            let ix = bucketKey .&. (A.size dst - 1)
     70            next <- A.read dst ix
     71            A.write dst ix (bkt { bucketNext = next })
     72            go i (n + 1) bucketNext
    8473
    8574insertWith :: (a -> a -> a) -> Int -> a -> IntTable a -> IO (Maybe a)
    86 insertWith f k v inttable@(IntTable ref) = do
    87   it@IT{..} <- readIORef ref
    88   let idx = indexOf k it
    89       go seen bkt@Bucket{..}
    90         | bucketKey == k = do
    91           let !v' = f v bucketValue
    92               !next = seen <> bucketNext
    93               Empty        <> bs = bs
    94               b@Bucket{..} <> bs = b { bucketNext = bucketNext <> bs }
    95           Arr.write tabArr idx (Bucket k v' next)
    96           return (Just bucketValue)
    97         | otherwise = go bkt { bucketNext = seen } bucketNext
    98       go seen _ = withForeignPtr tabSize $ \ptr -> do
    99         size <- peek ptr
    100         if size + 1 >= Arr.size tabArr - (Arr.size tabArr `shiftR` 2)
    101           then grow it ref size >> insertWith f k v inttable
    102           else do
    103             v `seq` Arr.write tabArr idx (Bucket k v seen)
    104             poke ptr (size + 1)
    105             return Nothing
    106   go Empty =<< Arr.read tabArr idx
    107 {-# INLINABLE insertWith #-}
     75insertWith f k v tbl@IntTable{..} = do
     76    arr  <- readIORef tblRef
     77    let ix = indexOf k arr
     78        go cont Empty = do
     79            size <- withForeignPtr tblSize peek
     80            if size + 1 >= A.size arr - (A.size arr `unsafeShiftR` 2)
     81            then do arr' <- A.new Empty (A.size arr `unsafeShiftL` 1)
     82                    copy arr arr' size
     83                    writeIORef tblRef arr'
     84                    insertWith f k v tbl
     85            else do A.write arr ix $! cont (Bucket k v Empty)
     86                    withForeignPtr tblSize (`poke` (size + 1))
     87                    return Nothing
     88        go cont bkt@Bucket{..}
     89            | bucketKey == k = do
     90                let !v' = f v bucketVal
     91                A.write arr ix $! cont (bkt { bucketVal = v' })
     92                return (Just bucketVal)
     93            | otherwise = go (\x -> cont (bkt { bucketNext = x })) bucketNext
     94    A.read arr ix >>= go id
    10895
    109 -- | Used to undo the effect of a prior insertWith.
    11096reset :: Int -> Maybe a -> IntTable a -> IO ()
    111 reset k (Just v) tbl = insertWith const k v tbl >> return ()
    11297reset k Nothing  tbl = delete k tbl >> return ()
    113 
    114 indexOf :: Int -> IT a -> Int
    115 indexOf k IT{..} = k .&. (Arr.size tabArr - 1)
     98reset k (Just v) tbl = insertWith const k v tbl >> return ()
    11699
    117100delete :: Int -> IntTable a -> IO (Maybe a)
    118 delete k t = updateWith (const Nothing) k t
     101delete = updateWith (const Nothing)
    119102
    120103updateWith :: (a -> Maybe a) -> Int -> IntTable a -> IO (Maybe a)
    121 updateWith f k (IntTable ref) = do
    122   it@IT{..} <- readIORef ref
    123   let idx = indexOf k it
    124       go changed bkt@Bucket{..}
    125         | bucketKey == k =
    126             let fbv = f bucketValue
    127                 !nb = case fbv of
    128                         Just val -> bkt { bucketValue = val }
    129                         Nothing  -> bucketNext
    130             in (fbv, Just bucketValue, nb)
    131         | otherwise = case go changed bucketNext of
    132                         (fbv, ov, nb) -> (fbv, ov, bkt { bucketNext = nb })
    133       go _ e = (Nothing, Nothing, e)
    134   (fbv, oldVal, newBucket) <- go False `liftM` Arr.read tabArr idx
    135   when (isJust oldVal) $ do
    136     Arr.write tabArr idx newBucket
    137     unless (isJust fbv) $
    138       withForeignPtr tabSize $ \ptr -> do
    139         size <- peek ptr
    140         poke ptr (size - 1)
    141   return oldVal
     104updateWith f k IntTable{..} = do
     105    arr <- readIORef tblRef
     106    let ix = indexOf k arr
     107        go _ Empty = (False, Nothing, Empty)
     108        go cont bkt@Bucket{..}
     109            | bucketKey == k = case f bucketVal of
     110                Nothing -> (True,  Just bucketVal, cont bucketNext)
     111                Just v  -> (False, Just bucketVal, cont (bkt { bucketVal = v }))
     112            | otherwise = go (\x -> cont (bkt { bucketNext = x })) bucketNext
     113    (del, old, bkt) <- A.read arr ix >>= return . go id
     114    when (isJust old) $ do
     115        A.write arr ix bkt
     116        when del . withForeignPtr tblSize $ \ptr -> do
     117            size <- peek ptr
     118            poke ptr (size - 1)
     119    return old