@@ -1,38 +1,118 @@
1+{-# LANGUAGE BangPatterns #-}
2+
3+------------------------------------------------------------------------------
14-- |
2−-- Module : Data.Attoparsec.Text.FastSet
3−-- Copyright : Bryan O'Sullivan 2015
5+-- Module : Data.Attoparsec.FastSet
6+-- Copyright : Felipe Lessa 2010, Bryan O'Sullivan 2007-2015
47-- License : BSD3
58--
6−-- Maintainer : bos@serpentine.com
9+-- Maintainer : felipe.lessa@gmail.com
710-- Stability : experimental
811-- Portability : unknown
912--
10−-- Fast set membership tests for 'Char' values.
11−
13+-- Fast set membership tests for 'Char' values. We test for
14+-- membership using a hashtable implemented with Robin Hood
15+-- collision resolution. The set representation is unboxed,
16+-- and the characters and hashes interleaved, for efficiency.
17+--
18+--
19+-----------------------------------------------------------------------------
1220module Data.Attoparsec.Text.FastSet
1321 (
1422 -- * Data type
1523 FastSet
1624 -- * Construction
1725 , fromList
26+ , set
1827 -- * Lookup
1928 , member
2029 -- * Handy interface
2130 , charClass
2231 ) where
2332
24−import qualified Data.IntSet as I
25−import Data.Char (ord)
33+import Data.Bits ((.|.), (.&.), shiftR)
34+import Data.Function (on)
35+import Data.List (sort, sortBy)
36+import qualified Data.Array.Base as AB
37+import qualified Data.Array.Unboxed as A
38+import qualified Data.Text as T
39+
40+data FastSet = FastSet {
41+ table :: {-# UNPACK #-} !(A.UArray Int Int)
42+ , mask :: {-# UNPACK #-} !Int
43+ }
44+
45+data Entry = Entry {
46+ key :: {-# UNPACK #-} !Char
47+ , initialIndex :: {-# UNPACK #-} !Int
48+ , index :: {-# UNPACK #-} !Int
49+ }
2650
27−newtype FastSet = FastSet I.IntSet
51+offset :: Entry -> Int
52+offset e = index e - initialIndex e
53+
54+resolveCollisions :: [Entry] -> [Entry]
55+resolveCollisions [] = []
56+resolveCollisions [e] = [e]
57+resolveCollisions (a:b:entries) = a' : resolveCollisions (b' : entries)
58+ where (a', b')
59+ | index a < index b = (a, b)
60+ | offset a < offset b = (b { index=index a }, a { index=index a + 1 })
61+ | otherwise = (a, b { index=index a + 1 })
62+
63+pad :: Int -> [Entry] -> [Entry]
64+pad = go 0
65+ where -- ensure that we pad enough so that lookups beyond the
66+ -- last hash in the table fall within the array
67+ go !_ !m [] = replicate (max 1 m + 1) empty
68+ go k m (e:entries) = map (const empty) [k..i - 1] ++ e :
69+ go (i + 1) (m + i - k - 1) entries
70+ where i = index e
71+ empty = Entry '\0' maxBound 0
72+
73+nextPowerOf2 :: Int -> Int
74+nextPowerOf2 0 = 1
75+nextPowerOf2 x = go (x - 1) 1
76+ where go y 32 = y + 1
77+ go y k = go (y .|. (y `shiftR` k)) $ k * 2
78+
79+fastHash :: Char -> Int
80+fastHash = fromEnum
2881
2982fromList :: String -> FastSet
30−fromList = FastSet . I.fromList . map ord
83+fromList s = FastSet (AB.listArray (0, length interleaved - 1) interleaved)
84+ mask'
85+ where s' = ordNub (sort s)
86+ l = length s'
87+ mask' = nextPowerOf2 ((5 * l) `div` 4) - 1
88+ entries = pad mask' .
89+ resolveCollisions .
90+ sortBy (compare `on` initialIndex) .
91+ zipWith (\c i -> Entry c i i) s' .
92+ map ((.&. mask') . fastHash) $ s'
93+ interleaved = concatMap (\e -> [fromEnum $ key e, initialIndex e])
94+ entries
95+
96+ordNub :: Eq a => [a] -> [a]
97+ordNub [] = []
98+ordNub (y:ys) = go y ys
99+ where go x (z:zs)
100+ | x == z = go x zs
101+ | otherwise = x : go z zs
102+ go x [] = [x]
103+
104+set :: T.Text -> FastSet
105+set = fromList . T.unpack
31106
32107-- | Check the set for membership.
33108member :: Char -> FastSet -> Bool
34−member c (FastSet s) = I.member (ord c) s
35−{-# INLINE member #-}
109+member c a = go (2 * i)
110+ where i = fastHash c .&. mask a
111+ lookupAt j b = (i' <= i) && (c == c' || b)
112+ where c' = toEnum $ AB.unsafeAt (table a) j
113+ i' = AB.unsafeAt (table a) $ j + 1
114+ go j = lookupAt j . lookupAt (j + 2) . lookupAt (j + 4) .
115+ lookupAt (j + 6) . go $ j + 8
36116
37117charClass :: String -> FastSet
38118charClass = fromList . go