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