@@ -28,20 +28,23 @@ module Data.Attoparsec.Text.FastSet
2828 , charClass
2929 ) where
3030
31−import Data.Bits
32−import Data.List (sort, sortBy, nub)
31+import Data.Bits ((.|.), (.&.), shiftR)
3332import Data.Function (on)
33+import Data.List (sort, sortBy, nub)
3434import qualified Data.Array.Base as AB
3535import qualified Data.Array.Unboxed as A
3636import qualified Data.Text as T
3737
38−data FastSet = FastSet {table :: {-# UNPACK #-} !(A.UArray Int Int),
39− mask :: {-# UNPACK #-} !Int}
40− deriving (Eq, Ord, Show)
38+data FastSet = FastSet {
39+ table :: {-# UNPACK #-} !(A.UArray Int Int)
40+ , mask :: {-# UNPACK #-} !Int
41+ } deriving (Eq, Ord, Show)
4142
42−data Entry = Entry {key :: {-# UNPACK #-} !Char,
43− initialIndex :: {-# UNPACK #-} !Int,
44− index :: {-# UNPACK #-} !Int} deriving (Show)
43+data Entry = Entry {
44+ key :: {-# UNPACK #-} !Char
45+ , initialIndex :: {-# UNPACK #-} !Int
46+ , index :: {-# UNPACK #-} !Int
47+ } deriving (Show)
4548
4649offset :: Entry -> Int
4750offset e = index e - initialIndex e
@@ -49,54 +52,58 @@ offset e = index e - initialIndex e
4952resolveCollisions :: [Entry] -> [Entry]
5053resolveCollisions [] = []
5154resolveCollisions [e] = [e]
52−resolveCollisions (a:b:entries) = a' : resolveCollisions (b':entries)
53− where (a', b')
54− | index a < index b = (a, b)
55− | offset a < offset b = (b{index=index a}, a{index=index a + 1})
56− | otherwise = (a, b{index=index a + 1})
55+resolveCollisions (a:b:entries) = a' : resolveCollisions (b' : entries)
56+ where (a', b')
57+ | index a < index b = (a, b)
58+ | offset a < offset b = (b { index=index a }, a { index=index a + 1 })
59+ | otherwise = (a, b { index=index a + 1 })
5760
58−pad :: Int -> [Entry]-> [Entry]
61+pad :: Int -> [Entry] -> [Entry]
5962pad = go 0
60− where go _ m [] = replicate (max 1 m) empty
61− go k m (e:entries) = map (const empty) [k..i - 1] ++ e : go (i + 1) (m + i - k - 1) entries
62− where i = index e
63− empty = Entry '\0' maxBound 0
63+ where go _ m [] = replicate (max 1 m) empty
64+ go k m (e:entries) = map (const empty) [k..i - 1] ++ e :
65+ go (i + 1) (m + i - k - 1) entries
66+ where i = index e
67+ empty = Entry '\0' maxBound 0
6468
6569nextPowerOf2 :: Int -> Int
66−nextPowerOf2 0 = 1
67−nextPowerOf2 x = go (x - 1) 1
68− where go y 32 = y + 1
69− go y k = go (y .|. (y `shiftR` k)) $ k * 2
70+nextPowerOf2 0 = 1
71+nextPowerOf2 x = go (x - 1) 1
72+ where go y 32 = y + 1
73+ go y k = go (y .|. (y `shiftR` k)) $ k * 2
7074
7175fastHash :: Char -> Int
7276fastHash = fromEnum
7377
7478fromList :: String -> FastSet
75−fromList s = FastSet (AB.listArray (0, length interleaved - 1) interleaved) mask'
76− where s' = nub $ sort s
77− l = length s'
78− mask' = nextPowerOf2 ((5 * l) `div` 4) - 1
79− indeces = map ((.&. mask') . fastHash) s'
80− entries = pad mask' .
81− resolveCollisions .
82− sortBy (compare `on` initialIndex) $
83− zipWith (\c i -> Entry c i i) s' indeces
84− interleaved = concatMap (\e -> [fromEnum $ key e, initialIndex e]) entries
79+fromList s = FastSet (AB.listArray (0, length interleaved - 1) interleaved)
80+ mask'
81+ where s' = nub $ sort s
82+ l = length s'
83+ mask' = nextPowerOf2 ((5 * l) `div` 4) - 1
84+ indices = map ((.&. mask') . fastHash) s'
85+ entries = pad mask' .
86+ resolveCollisions .
87+ sortBy (compare `on` initialIndex) .
88+ zipWith (\c i -> Entry c i i) s' $ indices
89+ interleaved = concatMap (\e -> [fromEnum $ key e, initialIndex e])
90+ entries
8591
8692set :: T.Text -> FastSet
8793set = fromList . T.unpack
8894
8995-- | Check the set for membership.
9096member :: Char -> FastSet -> Bool
91−member c a = go (2 * i)
92− where i = fastHash c .&. mask a
93− lookupAt j b = (i' <= i) && (c == c' || b)
94− where c' = toEnum $ AB.unsafeAt (table a) j
95− i' = AB.unsafeAt (table a) $ j + 1
96− go j = lookupAt j $ lookupAt (j + 2) $ lookupAt (j + 4) $ lookupAt (j + 6) $ go (j + 8)
97+member c a = go (2 * i)
98+ where i = fastHash c .&. mask a
99+ lookupAt j b = (i' <= i) && (c == c' || b)
100+ where c' = toEnum $ AB.unsafeAt (table a) j
101+ i' = AB.unsafeAt (table a) $ j + 1
102+ go j = lookupAt j . lookupAt (j + 2) . lookupAt (j + 4) .
103+ lookupAt (j + 6) . go $ j + 8
97104
98105charClass :: String -> FastSet
99106charClass = fromList . go
100− where go (a:'-':b:xs) = [a..b] ++ go xs
101− go (x:xs) = x : go xs
102− go _ = ""
107+ where go (a:'-':b:xs) = [a..b] ++ go xs
108+ go (x:xs) = x : go xs
109+ go _ = ""