{-# LANGUAGE BangPatterns #-}

------------------------------------------------------------------------------
-- |
-- Module      :  Data.Attoparsec.FastSet
-- Copyright   :  Felipe Lessa 2010, Bryan O'Sullivan 2007-2015
-- License     :  BSD3
--
-- Maintainer  :  felipe.lessa@gmail.com
-- Stability   :  experimental
-- Portability :  unknown
--
-- Fast set membership tests for 'Char' values. We test for
-- membership using a hashtable implemented with Robin Hood
-- collision resolution. The set representation is unboxed,
-- and the characters and hashes interleaved, for efficiency.
--
--
-----------------------------------------------------------------------------
module Data.Attoparsec.Text.FastSet
    (
    -- * Data type
      FastSet
    -- * Construction
    , fromList
    , set
    -- * Lookup
    , member
    -- * Handy interface
    , charClass
    ) where

import Data.Bits ((.|.), (.&.), shiftR)
import Data.Function (on)
import Data.List (sort, sortBy)
import qualified Data.Array.Base as AB
import qualified Data.Array.Unboxed as A
import qualified Data.Text as T

data FastSet = FastSet {
    FastSet -> UArray Int Int
table :: {-# UNPACK #-} !(A.UArray Int Int)
  , FastSet -> Int
mask  :: {-# UNPACK #-} !Int
  }

data Entry = Entry {
    Entry -> Char
key          :: {-# UNPACK #-} !Char
  , Entry -> Int
initialIndex :: {-# UNPACK #-} !Int
  , Entry -> Int
index        :: {-# UNPACK #-} !Int
  }

offset :: Entry -> Int
offset :: Entry -> Int
offset Entry
e = Entry -> Int
index Entry
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Entry -> Int
initialIndex Entry
e

resolveCollisions :: [Entry] -> [Entry]
resolveCollisions :: [Entry] -> [Entry]
resolveCollisions [] = []
resolveCollisions [Entry
e] = [Entry
e]
resolveCollisions (Entry
a:Entry
b:[Entry]
entries) = Entry
a' Entry -> [Entry] -> [Entry]
forall a. a -> [a] -> [a]
: [Entry] -> [Entry]
resolveCollisions (Entry
b' Entry -> [Entry] -> [Entry]
forall a. a -> [a] -> [a]
: [Entry]
entries)
  where (Entry
a', Entry
b')
          | Entry -> Int
index Entry
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Entry -> Int
index Entry
b   = (Entry
a, Entry
b)
          | Entry -> Int
offset Entry
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Entry -> Int
offset Entry
b = (Entry
b { index :: Int
index=Entry -> Int
index Entry
a }, Entry
a { index :: Int
index=Entry -> Int
index Entry
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 })
          | Bool
otherwise           = (Entry
a, Entry
b { index :: Int
index=Entry -> Int
index Entry
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 })

pad :: Int -> [Entry] -> [Entry]
pad :: Int -> [Entry] -> [Entry]
pad = Int -> Int -> [Entry] -> [Entry]
go Int
0
  where -- ensure that we pad enough so that lookups beyond the
        -- last hash in the table fall within the array
        go :: Int -> Int -> [Entry] -> [Entry]
go !Int
_ !Int
m []          = Int -> Entry -> [Entry]
forall a. Int -> a -> [a]
replicate (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Entry
empty
        go  Int
k  Int
m (Entry
e:[Entry]
entries) = (Int -> Entry) -> [Int] -> [Entry]
forall a b. (a -> b) -> [a] -> [b]
map (Entry -> Int -> Entry
forall a b. a -> b -> a
const Entry
empty) [Int
k..Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [Entry] -> [Entry] -> [Entry]
forall a. [a] -> [a] -> [a]
++ Entry
e Entry -> [Entry] -> [Entry]
forall a. a -> [a] -> [a]
:
                               Int -> Int -> [Entry] -> [Entry]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Entry]
entries
          where i :: Int
i            = Entry -> Int
index Entry
e
        empty :: Entry
empty                = Char -> Int -> Int -> Entry
Entry Char
'\0' Int
forall a. Bounded a => a
maxBound Int
0

nextPowerOf2 :: Int -> Int
nextPowerOf2 :: Int -> Int
nextPowerOf2 Int
0  = Int
1
nextPowerOf2 Int
x  = Int -> Int -> Int
forall t. (Num t, Bits t) => t -> Int -> t
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
1
  where go :: t -> Int -> t
go t
y Int
32 = t
y t -> t -> t
forall a. Num a => a -> a -> a
+ t
1
        go t
y Int
k  = t -> Int -> t
go (t
y t -> t -> t
forall a. Bits a => a -> a -> a
.|. (t
y t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
k)) (Int -> t) -> Int -> t
forall a b. (a -> b) -> a -> b
$ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2

fastHash :: Char -> Int
fastHash :: Char -> Int
fastHash = Char -> Int
forall a. Enum a => a -> Int
fromEnum

fromList :: String -> FastSet
fromList :: String -> FastSet
fromList String
s = UArray Int Int -> Int -> FastSet
FastSet ((Int, Int) -> [Int] -> UArray Int Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
AB.listArray (Int
0, [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
interleaved Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int]
interleaved)
             Int
mask'
  where s' :: String
s'      = String -> String
forall a. Eq a => [a] -> [a]
ordNub (String -> String
forall a. Ord a => [a] -> [a]
sort String
s)
        l :: Int
l       = String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s'
        mask' :: Int
mask'   = Int -> Int
nextPowerOf2 ((Int
5 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
l) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        entries :: [Entry]
entries = Int -> [Entry] -> [Entry]
pad Int
mask' ([Entry] -> [Entry]) -> (String -> [Entry]) -> String -> [Entry]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                  [Entry] -> [Entry]
resolveCollisions ([Entry] -> [Entry]) -> (String -> [Entry]) -> String -> [Entry]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                  (Entry -> Entry -> Ordering) -> [Entry] -> [Entry]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> (Entry -> Int) -> Entry -> Entry -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Entry -> Int
initialIndex) ([Entry] -> [Entry]) -> (String -> [Entry]) -> String -> [Entry]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                  (Char -> Int -> Entry) -> String -> [Int] -> [Entry]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Char
c Int
i -> Char -> Int -> Int -> Entry
Entry Char
c Int
i Int
i) String
s' ([Int] -> [Entry]) -> (String -> [Int]) -> String -> [Entry]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                  (Char -> Int) -> String -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask') (Int -> Int) -> (Char -> Int) -> Char -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
fastHash) (String -> [Entry]) -> String -> [Entry]
forall a b. (a -> b) -> a -> b
$ String
s'
        interleaved :: [Int]
interleaved = (Entry -> [Int]) -> [Entry] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Entry
e -> [Char -> Int
forall a. Enum a => a -> Int
fromEnum (Char -> Int) -> Char -> Int
forall a b. (a -> b) -> a -> b
$ Entry -> Char
key Entry
e, Entry -> Int
initialIndex Entry
e])
                      [Entry]
entries

ordNub :: Eq a => [a] -> [a]
ordNub :: [a] -> [a]
ordNub []     = []
ordNub (a
y:[a]
ys) = a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
go a
y [a]
ys
  where go :: a -> [a] -> [a]
go a
x (a
z:[a]
zs)
          | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
z    = a -> [a] -> [a]
go a
x [a]
zs
          | Bool
otherwise = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a -> [a] -> [a]
go a
z [a]
zs
        go a
x []       = [a
x]

set :: T.Text -> FastSet
set :: Text -> FastSet
set = String -> FastSet
fromList (String -> FastSet) -> (Text -> String) -> Text -> FastSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack

-- | Check the set for membership.
member :: Char -> FastSet -> Bool
member :: Char -> FastSet -> Bool
member Char
c FastSet
a           = Int -> Bool
go (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i)
  where i :: Int
i            = Char -> Int
fastHash Char
c Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. FastSet -> Int
mask FastSet
a
        lookupAt :: Int -> Bool -> Bool
lookupAt Int
j Bool
b = (Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i) Bool -> Bool -> Bool
&& (Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
c' Bool -> Bool -> Bool
|| Bool
b)
            where c' :: Char
c' = Int -> Char
forall a. Enum a => Int -> a
toEnum (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ UArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
AB.unsafeAt (FastSet -> UArray Int Int
table FastSet
a) Int
j
                  i' :: Int
i' = UArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
AB.unsafeAt (FastSet -> UArray Int Int
table FastSet
a) (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        go :: Int -> Bool
go Int
j         = Int -> Bool -> Bool
lookupAt Int
j (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Bool -> Bool
lookupAt (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Bool -> Bool
lookupAt (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                       Int -> Bool -> Bool
lookupAt (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6) (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Bool
go (Int -> Bool) -> Int -> Bool
forall a b. (a -> b) -> a -> b
$ Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8

charClass :: String -> FastSet
charClass :: String -> FastSet
charClass = String -> FastSet
fromList (String -> FastSet) -> (String -> String) -> String -> FastSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
go
  where go :: String -> String
go (Char
a:Char
'-':Char
b:String
xs) = [Char
a..Char
b] String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
go String
xs
        go (Char
x:String
xs)       = Char
x Char -> String -> String
forall a. a -> [a] -> [a]
: String -> String
go String
xs
        go String
_            = String
""