Ticket #8793: IntTable.diff

File IntTable.diff, 8.4 KB (added by cdk, 16 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