1{-# LANGUAGE ForeignFunctionInterface #-} 2{-# LANGUAGE RecordWildCards #-} 3{-# LANGUAGE BangPatterns #-} 4{-# LANGUAGE FlexibleContexts #-} 5 6-- Haskell implementation of H2O's priority queue. 7-- https://github.com/h2o/h2o/blob/master/lib/http2/scheduler.c 8 9-- delete is not supported because TQueue does not support deletion. 10-- So, key is not passed to enqueue. 11 12module RingOfQueuesSTM ( 13 Entry 14 , newEntry 15 , renewEntry 16 , item 17 , PriorityQueue(..) 18 , new 19 , enqueue 20 , dequeue 21 ) where 22 23import Control.Concurrent.STM 24import Control.Monad (replicateM) 25import Data.Array (Array, listArray, (!)) 26import Data.Bits (setBit, clearBit, shiftR) 27import Data.Word (Word64) 28import Foreign.C.Types (CLLong(..)) 29 30---------------------------------------------------------------- 31 32type Weight = Int 33 34-- | Abstract data type of entries for priority queues. 35data Entry a = Entry { 36 item :: a -- ^ Extracting an item from an entry. 37 , weight :: {-# UNPACK #-} !Weight 38 , deficit :: {-# UNPACK #-} !Int 39 } deriving Show 40 41newEntry :: a -> Weight -> Entry a 42newEntry x w = Entry x w 0 43 44-- | Changing the item of an entry. 45renewEntry :: Entry a -> b -> Entry b 46renewEntry ent x = ent { item = x } 47 48---------------------------------------------------------------- 49 50data PriorityQueue a = PriorityQueue { 51 bitsRef :: TVar Word64 52 , offsetRef :: TVar Int 53 , queues :: Array Int (TQueue (Entry a)) 54 } 55 56---------------------------------------------------------------- 57 58bitWidth :: Int 59bitWidth = 64 60 61relativeIndex :: Int -> Int -> Int 62relativeIndex idx offset = (offset + idx) `mod` bitWidth 63 64---------------------------------------------------------------- 65 66deficitSteps :: Int 67deficitSteps = 65536 68 69deficitList :: [Int] 70deficitList = map calc idxs 71 where 72 idxs :: [Double] 73 idxs = [1..256] 74 calc w = round (65536 * 63 / w) 75 76deficitTable :: Array Int Int 77deficitTable = listArray (1,256) deficitList 78 79---------------------------------------------------------------- 80 81-- https://en.wikipedia.org/wiki/Find_first_set 82foreign import ccall unsafe "strings.h ffsll" 83 c_ffs :: CLLong -> CLLong 84 85-- | Finding first bit set. O(1) 86-- 87-- >>> firstBitSet $ setBit 0 63 88-- 63 89-- >>> firstBitSet $ setBit 0 62 90-- 62 91-- >>> firstBitSet $ setBit 0 1 92-- 1 93-- >>> firstBitSet $ setBit 0 0 94-- 0 95-- >>> firstBitSet 0 96-- -1 97firstBitSet :: Word64 -> Int 98firstBitSet x = ffs x - 1 99 where 100 ffs = fromIntegral . c_ffs . fromIntegral 101 102---------------------------------------------------------------- 103 104new :: STM (PriorityQueue a) 105new = PriorityQueue <$> newTVar 0 <*> newTVar 0 <*> newQueues 106 where 107 newQueues = listArray (0, bitWidth - 1) <$> replicateM bitWidth newTQueue 108 109-- | Enqueuing an entry. PriorityQueue is updated. 110enqueue :: Entry a -> PriorityQueue a -> STM () 111enqueue ent PriorityQueue{..} = do 112 let (!idx,!deficit') = calcIdxAndDeficit 113 !offidx <- getOffIdx idx 114 push offidx ent { deficit = deficit' } 115 updateBits idx 116 where 117 calcIdxAndDeficit = total `divMod` deficitSteps 118 where 119 total = deficitTable ! weight ent + deficit ent 120 getOffIdx idx = relativeIndex idx <$> readTVar offsetRef 121 push offidx ent' = writeTQueue (queues ! offidx) ent' 122 updateBits idx = modifyTVar' bitsRef $ flip setBit idx 123 124-- | Dequeuing an entry. PriorityQueue is updated. 125dequeue :: PriorityQueue a -> STM (Entry a) 126dequeue pq@PriorityQueue{..} = do 127 !idx <- getIdx 128 if idx == -1 then 129 retry 130 else do 131 !offidx <- getOffIdx idx 132 updateOffset offidx 133 queueIsEmpty <- checkEmpty offidx 134 updateBits idx queueIsEmpty 135 if queueIsEmpty then 136 dequeue pq 137 else 138 pop offidx 139 where 140 getIdx = firstBitSet <$> readTVar bitsRef 141 getOffIdx idx = relativeIndex idx <$> readTVar offsetRef 142 pop offidx = readTQueue (queues ! offidx) 143 checkEmpty offidx = isEmptyTQueue (queues ! offidx) 144 updateOffset offset' = writeTVar offsetRef offset' 145 updateBits idx isEmpty = modifyTVar' bitsRef shiftClear 146 where 147 shiftClear bits 148 | isEmpty = clearBit (shiftR bits idx) 0 149 | otherwise = shiftR bits idx 150