Ticket #5895: Repro.hs

File Repro.hs, 6.1 KB (added by tibbe, 2 years ago)
Line 
1{-# LANGUAGE BangPatterns, CPP, MagicHash, UnboxedTuples #-}
2module Repro where
3
4import Control.Monad.ST
5import Data.Bits hiding (unsafeShiftL)
6import Data.Word
7import GHC.Exts
8import GHC.ST (ST(..))
9import Prelude hiding (filter, length, pred, read)
10
11------------------------------------------------------------------------
12-- Test harness (look at the core for this function)
13
14test :: HashMap Int Int -> HashMap Int Int
15test m = filter (< 42) m
16
17------------------------------------------------------------------------
18-- The function in question
19
20-- | /O(n)/ Filter this map by retaining only elements which values
21-- satisfy a predicate.
22filter :: (v -> Bool) -> HashMap k v -> HashMap k v
23filter p = filterWithKey (\_ v -> p v)
24{-# INLINE filter #-}
25
26-- | /O(n)/ Filter this map by retaining only elements satisfying a
27-- predicate.
28filterWithKey :: (k -> v -> Bool) -> HashMap k v -> HashMap k v
29filterWithKey pred = go
30  where
31    go Empty = Empty
32    go t@(Leaf _ (L k v))
33        | pred k v  = t
34        | otherwise = Empty
35    go (BitmapIndexed b ary) = filterA ary b
36    go (Full ary) = filterA ary fullNodeMask
37    go (Collision h ary) = filterC ary h
38
39    filterA ary0 b0 =
40        let !n = length ary0
41        in runST $ do
42            mary <- new_ n
43            step ary0 mary b0 0 0 1 n
44      where
45        step !ary !mary !b i !j !bi n
46            | i >= n = case j of
47                0 -> return Empty
48                1 -> read mary 0
49                _ -> do
50                    mary2 <- new_ j
51                    copyM mary 0 mary2 0 j
52                    ary2 <- unsafeFreeze mary2
53                    return $! if j == maxChildren
54                              then Full ary2
55                              else BitmapIndexed b ary2
56            | bi .&. b == 0 = step ary mary b i j (bi `unsafeShiftL` 1) n
57            | otherwise = case go (index ary i) of
58                Empty -> step ary mary (b .&. complement bi) (i+1) j
59                         (bi `unsafeShiftL` 1) n
60                t     -> do write mary j t
61                            step ary mary b (i+1) (j+1) (bi `unsafeShiftL` 1) n
62    {-# INLINE filterA #-}
63
64    filterC ary0 h =
65        let !n = length ary0
66        in runST $ do
67            mary <- new_ n
68            step ary0 mary 0 0 n
69      where
70        step ary mary i j n
71            | i >= n    = case j of
72                0 -> return Empty
73                1 -> do l <- read mary 0
74                        return $! Leaf h l
75                _ | i == j -> do ary2 <- unsafeFreeze mary
76                                 return $! Collision h ary2
77                  | otherwise -> do mary2 <- new_ j
78                                    copyM mary 0 mary2 0 j
79                                    ary2 <- unsafeFreeze mary2
80                                    return $! Collision h ary2
81            | pred k v  = write mary j el >> step ary mary (i+1) (j+1) n
82            | otherwise = step ary mary (i+1) j n
83          where el@(L k v) = index ary i
84{-# INLINE filterWithKey #-}
85
86------------------------------------------------------------------------
87-- Auxiliary definitions (to make this example standalone)
88
89-- | A map from keys to values.  A map cannot contain duplicate keys;
90-- each key can map to at most one value.
91data HashMap k v
92    = Empty
93    | BitmapIndexed {-# UNPACK #-} !Bitmap {-# UNPACK #-} !(Array (HashMap k v))
94    | Leaf {-# UNPACK #-} !Hash {-# UNPACK #-} !(Leaf k v)
95    | Full {-# UNPACK #-} !(Array (HashMap k v))
96    | Collision {-# UNPACK #-} !Hash {-# UNPACK #-} !(Array (Leaf k v))
97
98data Leaf k v = L !k v
99
100type Hash   = Word
101type Bitmap = Word
102type Shift  = Int
103
104bitsPerSubkey :: Int
105bitsPerSubkey = 4
106
107maxChildren :: Int
108maxChildren = fromIntegral $ 1 `unsafeShiftL` bitsPerSubkey
109
110-- | A bitmask with the 'bitsPerSubkey' least significant bits set.
111fullNodeMask :: Bitmap
112fullNodeMask = complement (complement 0 `unsafeShiftL`
113                           fromIntegral (1 `unsafeShiftL` bitsPerSubkey))
114{-# INLINE fullNodeMask #-}
115
116unsafeShiftL :: Word -> Int -> Word
117unsafeShiftL (W# x#) (I# i#) = W# (x# `uncheckedShiftL#` i#)
118{-# INLINE unsafeShiftL #-}
119
120------------------------------------------------------------------------
121-- Auxiliary array definition
122
123data Array a = Array {
124      unArray :: !(Array# a)
125    }
126
127-- | Smart constructor
128array :: Array# a -> Int -> Array a
129array ary _n = Array ary
130{-# INLINE array #-}
131
132length :: Array a -> Int
133length ary = I# (sizeofArray# (unArray ary))
134{-# INLINE length #-}
135
136data MArray s a = MArray {
137      unMArray :: !(MutableArray# s a)
138    }
139
140-- | Smart constructor
141marray :: MutableArray# s a -> Int -> MArray s a
142marray mary _n = MArray mary
143{-# INLINE marray #-}
144
145lengthM :: MArray s a -> Int
146lengthM mary = I# (sizeofMutableArray# (unMArray mary))
147{-# INLINE lengthM #-}
148
149-- | Unsafely copy the elements of an array. Array bounds are not checked.
150copyM :: MArray s e -> Int -> MArray s e -> Int -> Int -> ST s ()
151copyM !src !(I# sidx#) !dst !(I# didx#) (I# n#) =
152    ST $ \ s# ->
153    case copyMutableArray# (unMArray src) sidx# (unMArray dst) didx# n# s# of
154        s2 -> (# s2, () #)
155
156index :: Array a -> Int -> a
157index ary _i@(I# i#) =
158        case indexArray# (unArray ary) i# of (# b #) -> b
159{-# INLINE index #-}
160
161-- | Create a new mutable array of specified size, in the specified
162-- state thread, with each element containing the specified initial
163-- value.
164new :: Int -> a -> ST s (MArray s a)
165new n@(I# n#) b =
166    ST $ \s ->
167        case newArray# n# b s of
168            (# s', ary #) -> (# s', marray ary n #)
169{-# INLINE new #-}
170
171new_ :: Int -> ST s (MArray s a)
172new_ n = new n undefinedElem
173
174read :: MArray s a -> Int -> ST s a
175read ary (I# i#) = ST $ \ s ->
176        readArray# (unMArray ary) i# s
177{-# INLINE read #-}
178
179undefinedElem :: a
180undefinedElem = error "Data.HashMap.Array: Undefined element"
181{-# NOINLINE undefinedElem #-}
182
183unsafeFreeze :: MArray s a -> ST s (Array a)
184unsafeFreeze mary
185    = ST $ \s -> case unsafeFreezeArray# (unMArray mary) s of
186                   (# s', ary #) -> (# s', array ary (lengthM mary) #)
187{-# INLINE unsafeFreeze #-}
188
189write :: MArray s a -> Int -> a -> ST s ()
190write ary _i@(I# i#) b = ST $ \ s ->
191        case writeArray# (unMArray ary) i# b s of
192            s' -> (# s' , () #)
193{-# INLINE write #-}