1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP          #-}
3
4module Data.HashTable.Internal.Linear.Bucket
5( Bucket,
6  newBucketArray,
7  newBucketSize,
8  emptyWithSize,
9  growBucketTo,
10  snoc,
11  size,
12  lookup,
13  lookupIndex,
14  elemAt,
15  delete,
16  mutate,
17  mutateST,
18  toList,
19  fromList,
20  mapM_,
21  foldM,
22  expandBucketArray,
23  expandArray,
24  nelemsAndOverheadInWords,
25  bucketSplitSize
26) where
27
28
29------------------------------------------------------------------------------
30#if !MIN_VERSION_base(4,8,0)
31import           Control.Applicative
32#endif
33import           Control.Monad                        hiding (foldM, mapM_)
34import qualified Control.Monad
35import           Control.Monad.ST                     (ST)
36#ifdef DEBUG
37import           Data.HashTable.Internal.Utils        (unsafeIOToST)
38#endif
39import           Data.HashTable.Internal.Array
40import           Data.Maybe                           (fromMaybe)
41import           Data.STRef
42import           Prelude                              hiding (lookup, mapM_)
43------------------------------------------------------------------------------
44import           Data.HashTable.Internal.UnsafeTricks
45
46
47#ifdef DEBUG
48import           System.IO
49#endif
50
51
52type Bucket s k v = Key (Bucket_ s k v)
53
54------------------------------------------------------------------------------
55data Bucket_ s k v = Bucket { _bucketSize :: {-# UNPACK #-} !Int
56                            , _highwater  :: {-# UNPACK #-} !(STRef s Int)
57                            , _keys       :: {-# UNPACK #-} !(MutableArray s k)
58                            , _values     :: {-# UNPACK #-} !(MutableArray s v)
59                            }
60
61
62------------------------------------------------------------------------------
63bucketSplitSize :: Int
64bucketSplitSize = 16
65
66
67------------------------------------------------------------------------------
68newBucketArray :: Int -> ST s (MutableArray s (Bucket s k v))
69newBucketArray k = newArray k emptyRecord
70
71------------------------------------------------------------------------------
72nelemsAndOverheadInWords :: Bucket s k v -> ST s (Int,Int)
73nelemsAndOverheadInWords bKey = do
74    if (not $ keyIsEmpty bKey)
75      then do
76        !hw <- readSTRef hwRef
77        let !w = sz - hw
78        return (hw, constOverhead + 2*w)
79      else
80        return (0, 0)
81
82  where
83    constOverhead = 8
84    b             = fromKey bKey
85    sz            = _bucketSize b
86    hwRef         = _highwater b
87
88
89------------------------------------------------------------------------------
90emptyWithSize :: Int -> ST s (Bucket s k v)
91emptyWithSize !sz = do
92    !keys   <- newArray sz undefined
93    !values <- newArray sz undefined
94    !ref    <- newSTRef 0
95
96    return $ toKey $ Bucket sz ref keys values
97
98
99------------------------------------------------------------------------------
100newBucketSize :: Int
101newBucketSize = 4
102
103
104------------------------------------------------------------------------------
105expandArray  :: a                  -- ^ default value
106             -> Int                -- ^ new size
107             -> Int                -- ^ number of elements to copy
108             -> MutableArray s a   -- ^ old array
109             -> ST s (MutableArray s a)
110expandArray def !sz !hw !arr = do
111    newArr <- newArray sz def
112    cp newArr
113
114  where
115    cp !newArr = go 0
116      where
117        go !i
118          | i >= hw = return newArr
119          | otherwise = do
120                readArray arr i >>= writeArray newArr i
121                go (i+1)
122
123
124------------------------------------------------------------------------------
125expandBucketArray :: Int
126                  -> Int
127                  -> MutableArray s (Bucket s k v)
128                  -> ST s (MutableArray s (Bucket s k v))
129expandBucketArray = expandArray emptyRecord
130
131
132------------------------------------------------------------------------------
133growBucketTo :: Int -> Bucket s k v -> ST s (Bucket s k v)
134growBucketTo !sz bk | keyIsEmpty bk = emptyWithSize sz
135                    | otherwise = do
136    if osz >= sz
137      then return bk
138      else do
139        hw <- readSTRef hwRef
140        k' <- expandArray undefined sz hw keys
141        v' <- expandArray undefined sz hw values
142        return $ toKey $ Bucket sz hwRef k' v'
143
144  where
145    bucket = fromKey bk
146    osz    = _bucketSize bucket
147    hwRef  = _highwater bucket
148    keys   = _keys bucket
149    values = _values bucket
150
151
152------------------------------------------------------------------------------
153{-# INLINE snoc #-}
154-- Just return == new bucket object
155snoc :: Bucket s k v -> k -> v -> ST s (Int, Maybe (Bucket s k v))
156snoc bucket | keyIsEmpty bucket = mkNew
157            | otherwise         = snoc' (fromKey bucket)
158  where
159    mkNew !k !v = do
160        debug "Bucket.snoc: mkNew"
161        keys   <- newArray newBucketSize undefined
162        values <- newArray newBucketSize undefined
163
164        writeArray keys 0 k
165        writeArray values 0 v
166        ref <- newSTRef 1
167        return (1, Just $ toKey $ Bucket newBucketSize ref keys values)
168
169    snoc' (Bucket bsz hwRef keys values) !k !v =
170        readSTRef hwRef >>= check
171      where
172        check !hw
173          | hw < bsz  = bump hw
174          | otherwise = spill hw
175
176        bump hw = do
177          debug $ "Bucket.snoc: bumping hw, bsz=" ++ show bsz ++ ", hw="
178                    ++ show hw
179
180          writeArray keys hw k
181          writeArray values hw v
182          let !hw' = hw + 1
183          writeSTRef hwRef hw'
184          debug "Bucket.snoc: finished"
185          return (hw', Nothing)
186
187        doublingThreshold = bucketSplitSize `div` 2
188        growFactor = 1.5 :: Double
189        newSize z | z == 0 = newBucketSize
190                  | z < doublingThreshold = z * 2
191                  | otherwise = ceiling $ growFactor * fromIntegral z
192
193        spill !hw = do
194            let sz = newSize bsz
195            debug $ "Bucket.snoc: spilling, old size=" ++ show bsz ++ ", new size="
196                      ++ show sz
197
198            bk <- growBucketTo sz bucket
199
200            debug "Bucket.snoc: spill finished, snoccing element"
201            let (Bucket _ hwRef' keys' values') = fromKey bk
202
203            let !hw' = hw+1
204            writeArray keys' hw k
205            writeArray values' hw v
206            writeSTRef hwRef' hw'
207
208            return (hw', Just bk)
209
210
211
212------------------------------------------------------------------------------
213{-# INLINE size #-}
214size :: Bucket s k v -> ST s Int
215size b | keyIsEmpty b = return 0
216       | otherwise = readSTRef $ _highwater $ fromKey b
217
218
219------------------------------------------------------------------------------
220-- note: search in reverse order! We prefer recently snoc'd keys.
221lookup :: (Eq k) => Bucket s k v -> k -> ST s (Maybe v)
222lookup bucketKey !k | keyIsEmpty bucketKey = return Nothing
223                    | otherwise = lookup' $ fromKey bucketKey
224  where
225    lookup' (Bucket _ hwRef keys values) = do
226        hw <- readSTRef hwRef
227        go (hw-1)
228      where
229        go !i
230            | i < 0 = return Nothing
231            | otherwise = do
232                k' <- readArray keys i
233                if k == k'
234                  then do
235                    !v <- readArray values i
236                    return $! Just v
237                  else go (i-1)
238
239------------------------------------------------------------------------------
240-- note: search in reverse order! We prefer recently snoc'd keys.
241lookupIndex :: (Eq k) => Bucket s k v -> k -> ST s (Maybe Int)
242lookupIndex bucketKey !k
243  | keyIsEmpty bucketKey = return Nothing
244  | otherwise = lookup' $ fromKey bucketKey
245  where
246    lookup' (Bucket _ hwRef keys _values) = do
247        hw <- readSTRef hwRef
248        go (hw-1)
249      where
250        go !i
251            | i < 0 = return Nothing
252            | otherwise = do
253                k' <- readArray keys i
254                if k == k'
255                  then return (Just i)
256                  else go (i-1)
257
258elemAt :: Bucket s k v -> Int -> ST s (Maybe (k,v))
259elemAt bucketKey ix
260  | keyIsEmpty bucketKey = return Nothing
261  | otherwise = lookup' $ fromKey bucketKey
262  where
263    lookup' (Bucket _ hwRef keys values) = do
264        hw <- readSTRef hwRef
265        if 0 <= ix && ix < hw
266          then do k <- readArray keys ix
267                  v <- readArray values ix
268                  return (Just (k,v))
269          else return Nothing
270
271------------------------------------------------------------------------------
272{-# INLINE toList #-}
273toList :: Bucket s k v -> ST s [(k,v)]
274toList bucketKey | keyIsEmpty bucketKey = return []
275                 | otherwise = toList' $ fromKey bucketKey
276  where
277    toList' (Bucket _ hwRef keys values) = do
278        hw <- readSTRef hwRef
279        go [] hw 0
280      where
281        go !l !hw !i | i >= hw   = return l
282                     | otherwise = do
283            k <- readArray keys i
284            v <- readArray values i
285            go ((k,v):l) hw $ i+1
286
287
288------------------------------------------------------------------------------
289-- fromList needs to reverse the input in order to make fromList . toList == id
290{-# INLINE fromList #-}
291fromList :: [(k,v)] -> ST s (Bucket s k v)
292fromList l = Control.Monad.foldM f emptyRecord (reverse l)
293  where
294    f bucket (k,v) = do
295        (_,m) <- snoc bucket k v
296        return $ fromMaybe bucket m
297
298------------------------------------------------------------------------------
299delete :: (Eq k) => Bucket s k v -> k -> ST s Bool
300delete bucketKey !k | keyIsEmpty bucketKey = do
301    debug $ "Bucket.delete: empty bucket"
302    return False
303                    | otherwise = do
304    debug "Bucket.delete: start"
305    del $ fromKey bucketKey
306  where
307    del (Bucket sz hwRef keys values) = do
308        hw <- readSTRef hwRef
309        debug $ "Bucket.delete: hw=" ++ show hw ++ ", sz=" ++ show sz
310        go hw $ hw - 1
311
312      where
313        go !hw !i | i < 0 = return False
314                  | otherwise = do
315            k' <- readArray keys i
316            if k == k'
317              then do
318                  debug $ "found entry to delete at " ++ show i
319                  move (hw-1) i keys
320                  move (hw-1) i values
321                  let !hw' = hw-1
322                  writeSTRef hwRef hw'
323                  return True
324              else go hw (i-1)
325
326
327------------------------------------------------------------------------------
328mutate :: (Eq k) =>
329          Bucket s k v
330       -> k
331       -> (Maybe v -> (Maybe v, a))
332       -> ST s (Int, Maybe (Bucket s k v), a)
333mutate bucketKey !k !f = mutateST bucketKey k (pure . f)
334{-# INLINE mutate #-}
335
336
337------------------------------------------------------------------------------
338mutateST :: (Eq k) =>
339            Bucket s k v
340         -> k
341         -> (Maybe v -> ST s (Maybe v, a))
342         -> ST s (Int, Maybe (Bucket s k v), a)
343mutateST bucketKey !k !f
344    | keyIsEmpty bucketKey = do
345        fRes <- f Nothing
346        case fRes of
347            (Nothing, a) -> return (0, Nothing, a)
348            (Just v', a) -> do
349                (!hw', mbk) <- snoc bucketKey k v'
350                return (hw', mbk, a)
351    | otherwise = mutate' $ fromKey bucketKey
352  where
353    mutate' (Bucket _sz hwRef keys values) = do
354        hw <- readSTRef hwRef
355        pos <- findPosition hw (hw-1)
356        mv <- do
357            if pos < 0
358                then return Nothing
359                else readArray values pos >>= return . Just
360        fRes <- f mv
361        case (mv, fRes) of
362            (Nothing, (Nothing, a)) -> return (hw, Nothing, a)
363            (Nothing, (Just v', a)) -> do
364                (!hw', mbk) <- snoc bucketKey k v'
365                return (hw', mbk, a)
366            (Just _v, (Just v', a)) -> do
367                writeArray values pos v'
368                return (hw, Nothing, a)
369            (Just _v, (Nothing, a)) -> do
370                move (hw-1) pos keys
371                move (hw-1) pos values
372                let !hw' = hw-1
373                writeSTRef hwRef hw'
374                return (hw', Nothing, a)
375      where
376        findPosition !hw !i
377            | i < 0 = return (-1)
378            | otherwise = do
379                k' <- readArray keys i
380                if k == k'
381                  then return i
382                  else findPosition hw (i-1)
383
384
385------------------------------------------------------------------------------
386{-# INLINE mapM_ #-}
387mapM_ :: ((k,v) -> ST s a) -> Bucket s k v -> ST s ()
388mapM_ f bucketKey
389    | keyIsEmpty bucketKey = do
390        debug $ "Bucket.mapM_: bucket was empty"
391        return ()
392    | otherwise = doMap $ fromKey bucketKey
393  where
394    doMap (Bucket sz hwRef keys values) = do
395        hw <- readSTRef hwRef
396        debug $ "Bucket.mapM_: hw was " ++ show hw ++ ", sz was " ++ show sz
397        go hw 0
398      where
399        go !hw !i | i >= hw = return ()
400                  | otherwise = do
401            k <- readArray keys i
402            v <- readArray values i
403            _ <- f (k,v)
404            go hw $ i+1
405
406
407------------------------------------------------------------------------------
408{-# INLINE foldM #-}
409foldM :: (a -> (k,v) -> ST s a) -> a -> Bucket s k v -> ST s a
410foldM f !seed0 bucketKey
411    | keyIsEmpty bucketKey = return seed0
412    | otherwise = doMap $ fromKey bucketKey
413  where
414    doMap (Bucket _ hwRef keys values) = do
415        hw <- readSTRef hwRef
416        go hw seed0 0
417      where
418        go !hw !seed !i | i >= hw = return seed
419                        | otherwise = do
420            k <- readArray keys i
421            v <- readArray values i
422            seed' <- f seed (k,v)
423            go hw seed' (i+1)
424
425
426------------------------------------------------------------------------------
427-- move i into j
428move :: Int -> Int -> MutableArray s a -> ST s ()
429move i j arr | i == j    = do
430    debug $ "move " ++ show i ++ " into " ++ show j
431    return ()
432             | otherwise = do
433    debug $ "move " ++ show i ++ " into " ++ show j
434    readArray arr i >>= writeArray arr j
435
436
437
438{-# INLINE debug #-}
439debug :: String -> ST s ()
440
441#ifdef DEBUG
442debug s = unsafeIOToST $ do
443              putStrLn s
444              hFlush stdout
445#else
446#ifdef TESTSUITE
447debug !s = do
448    let !_ = length s
449    return $! ()
450#else
451debug _ = return ()
452#endif
453#endif
454
455