I'm writing some performance-sensitive code which requires a lookup in a binary search tree. The tree itself will be generated once at the beginning of each run, and then never modified. Lookups will be performed quite often, easily in the millions, so I'm only concerned with lookup performance.
The keys of the tree are Chars and can cheaply be converted to Int, so I assumed that Data.Map and cousins would be the best choice: I'm using a strict IntMap here. However, I'm getting performance significantly below what I would expect from a binary search tree. I demonstrate here with a rather contrived example.
import qualified Data.IntMap.Strict as IM
divideByTen :: Int -> Int
divideByTen n = maybe 0 snd $ IM.lookupLE n divideByTenMap
divideByTenMap :: IM.IntMap Int
divideByTenMap = IM.fromList $ zip [0,10..1000] [0..]
My data will not be uniformly distributed, so performance of certain subsets will be more important than overall. Hence it makes sense to introduce some shortcuts for these commonly used inputs.
divideByTenFaster :: Int -> Int
divideByTenFaster n
| n < 10 = 0
| n < 20 = 1
| n < 30 = 2
| n < 40 = 3
| n < 50 = 4
| n < 60 = 5
| n < 70 = 6
| n < 80 = 7
| n < 90 = 8
| n < 100 = 9
| otherwise = divideByTen n
There are 100 elements of this IntMap, so a binary search should take at most 7 comparisons. Hence I would expect divideByTenFaster to be faster than divideByTen for n < 70, and slower afterwards. What I see instead is that divideByTenFaster is twice as fast as the IntMap lookup, even for n < 100. Running the following benchmarking code:
import Criterion.Main
main :: IO ()
main = do
defaultMainWith defaultConfig $
[ bench "divideByTen" $ nf divideByTen 99
, bench "divideByTenFaster" $ nf divideByTenFaster 99
]
we get
benchmarking divideByTen
time 18.50 ns (18.47 ns .. 18.54 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 18.49 ns (18.47 ns .. 18.51 ns)
std dev 79.42 ps (61.51 ps .. 110.3 ps)
benchmarking divideByTenFaster
time 7.164 ns (7.154 ns .. 7.176 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 7.161 ns (7.145 ns .. 7.176 ns)
std dev 49.60 ps (36.84 ps .. 68.32 ps)
In my real problem I'm seeing even more significant differences: the IntMap lookup is up to an order of magnitude slower. I've found myself writing long and difficult-to-maintain shortcuts, essentially manually building an unbalanced binary search tree, which significantly improves performance in common cases, and slows down the general case slightly.
I could move to an array-based solution, but I had thought this was the ideal use case for Map. Are there some optimisations I'm missing, or is another data type a better choice?
For full context the problem I'm dealing with is looking up Unicode character width, as determined by the Unicode specification. The lookup tree has around 700 to 1100 elements, depending on context. Performance of the ASCII subset is very important, and beyond that widely-used scripts should be quite fast (Latin, Han ideographs, Arabic, Devanagari, etc.), while performance for lesser-used scripts can be sacrificed a bit.
Edit:
Given the many suggestions in the comments, I've benchmarked a few different approaches to this problem. For pure lookup performance nothing beats a giant array, though that requires a non-trivial memory and construction cost, but other approaches all perform similarly, with the exception of directly implementing Data.Map.Strict's structure here (thanks to @dfeuer's comment below).
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Map.Internal as MInt
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as VU
bstSize :: Int
bstSize = 800
domainSize :: Int
domainSize = 1100000
divideByTenShortcut :: Int -> Int
divideByTenShortcut n
| n < 10 = 0
| n < 20 = 1
| n < 30 = 2
| n < 40 = 3
| n < 50 = 4
| n < 60 = 5
| n < 70 = 6
| n < 80 = 7
| n < 90 = 8
| n < 100 = 9
| otherwise = divideByTenIntMapLookup n
divideByTenIntMapLookup :: Int -> Int
divideByTenIntMapLookup n = maybe 0 snd $ IM.lookupLE n divideByTenIntMap
divideByTenIntMap :: IM.IntMap Int
divideByTenIntMap = IM.fromList $ zip [0,10..bstSize*10] [0..]
divideByTenMapLookup :: Int -> Int
divideByTenMapLookup n = maybe 0 snd $ M.lookupLE n divideByTenMap
divideByTenMap :: M.Map Int Int
divideByTenMap = M.fromList $ zip [0,10..bstSize*10] [0..]
arrayBinarySearch :: G.Vector v Int => v Int -> Int -> Int
arrayBinarySearch vs search = loop 0 (G.length vs)
where
loop !low !high
| high <= low = low
| otherwise = case compare midval search of
LT -> loop midpoint high
GT -> loop low (midpoint - 1)
EQ -> midpoint
where
midpoint = (low + high + 1) `quot` 2
midval = G.unsafeIndex vs midpoint
divideByTenArrayLookup :: Int -> Int
divideByTenArrayLookup n = let i = arrayBinarySearch divideByTenArray n in V.unsafeIndex divideByTenArraySols i
divideByTenArray :: V.Vector Int
divideByTenArray = V.fromList [0,10..bstSize*10]
divideByTenArraySols :: V.Vector Int
divideByTenArraySols = V.fromList [0..bstSize]
divideByTenArrayUnboxedLookup :: Int -> Int
divideByTenArrayUnboxedLookup n = let i = arrayBinarySearch divideByTenArrayUnboxed n in VU.unsafeIndex divideByTenArrayUnboxedSols i
divideByTenArrayUnboxed :: VU.Vector Int
divideByTenArrayUnboxed = VU.fromList [0,10..bstSize*10]
divideByTenArrayUnboxedSols :: VU.Vector Int
divideByTenArrayUnboxedSols = VU.fromList [0..bstSize]
divideByTenFullArrayLookup :: Int -> Int
divideByTenFullArrayLookup = V.unsafeIndex divideByTenFullArray
divideByTenFullArray :: V.Vector Int
divideByTenFullArray = V.fromList $ map (`quot` 10) [0..domainSize]
divideByTenFullArrayUnboxedLookup :: Int -> Int
divideByTenFullArrayUnboxedLookup = VU.unsafeIndex divideByTenFullArrayUnboxed
divideByTenFullArrayUnboxed :: VU.Vector Int
divideByTenFullArrayUnboxed = VU.fromList $ map (`quot` 10) [0..domainSize]
divideByTenCharMapLookup :: Int -> Int
divideByTenCharMapLookup n = lookupLE n divideByTenCharMap
divideByTenCharMap :: CharMap
divideByTenCharMap = fromList $ zip [0,10..bstSize*10] [0..]
data CharMap = Bin {-# UNPACK #-} !Size !Int !Int !CharMap !CharMap
| Tip
type Size = Int
lookupLE :: Int -> CharMap -> Int
lookupLE = goNothing
where
goNothing !_ Tip = 0
goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l
EQ -> x
GT -> goJust k kx x r
goJust !_ !_ x' Tip = x'
goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l
EQ -> x
GT -> goJust k kx x r
{-# INLINABLE lookupLE #-}
fromList :: [(Int, Int)] -> CharMap
fromList = repack . M.fromList
where
repack MInt.Tip = Tip
repack (MInt.Bin s k v l r) = Bin s k v (repack l) (repack r)
import Control.Exception (evaluate)
import Criterion.Main
main :: IO ()
main = do
evaluate divideByTenIntMap
evaluate divideByTenMap
evaluate divideByTenCharMap
evaluate divideByTenArray
evaluate divideByTenArrayUnboxed
evaluate divideByTenFullArray
evaluate divideByTenFullArrayUnboxed
defaultMainWith defaultConfig
[ bench "divideByTenIntMap" $ nf divideByTenIntMapLookup 99
, bench "divideByTenShortcut" $ nf divideByTenShortcut 99
, bench "divideByTenMap" $ nf divideByTenMapLookup 99
, bench "divideByTenCharMap" $ nf divideByTenCharMapLookup 99
, bench "divideByTenArray" $ nf divideByTenArrayLookup 99
, bench "divideByTenArrayUnboxed" $ nf divideByTenArrayUnboxedLookup 99
, bench "divideByTenFullArray" $ nf divideByTenFullArrayLookup 99
, bench "divideByTenFullArrayUnboxed" $ nf divideByTenFullArrayUnboxedLookup 99
]
Build profile: -w ghc-8.10.7 -O1
benchmarking divideByTenIntMap
time 21.31 ns (21.26 ns .. 21.36 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 21.31 ns (21.27 ns .. 21.36 ns)
std dev 137.3 ps (104.7 ps .. 190.5 ps)
benchmarking divideByTenShortcut
time 6.881 ns (6.868 ns .. 6.893 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 6.896 ns (6.873 ns .. 6.947 ns)
std dev 112.5 ps (40.52 ps .. 184.2 ps)
variance introduced by outliers: 23% (moderately inflated)
benchmarking divideByTenMap
time 23.31 ns (22.94 ns .. 23.68 ns)
0.999 R² (0.998 R² .. 1.000 R²)
mean 23.02 ns (22.91 ns .. 23.26 ns)
std dev 500.1 ps (245.4 ps .. 865.6 ps)
variance introduced by outliers: 33% (moderately inflated)
benchmarking divideByTenCharMap
time 15.89 ns (15.88 ns .. 15.92 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 15.90 ns (15.88 ns .. 15.92 ns)
std dev 68.89 ps (44.90 ps .. 113.1 ps)
benchmarking divideByTenArray
time 27.07 ns (27.01 ns .. 27.13 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 27.03 ns (26.92 ns .. 27.13 ns)
std dev 334.5 ps (248.0 ps .. 474.7 ps)
variance introduced by outliers: 14% (moderately inflated)
benchmarking divideByTenArrayUnboxed
time 22.89 ns (22.82 ns .. 22.99 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 22.94 ns (22.84 ns .. 23.10 ns)
std dev 418.6 ps (301.6 ps .. 555.1 ps)
variance introduced by outliers: 26% (moderately inflated)
benchmarking divideByTenFullArray
time 6.639 ns (6.626 ns .. 6.651 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 6.635 ns (6.626 ns .. 6.644 ns)
std dev 28.87 ps (24.14 ps .. 34.83 ps)
benchmarking divideByTenFullArrayUnboxed
time 6.029 ns (6.017 ns .. 6.041 ns)
1.000 R² (1.000 R² .. 1.000 R²)
mean 6.021 ns (6.013 ns .. 6.029 ns)
std dev 26.09 ps (21.68 ps .. 31.51 ps)
Benchmark doclayout-bench: FINISH
So the problem doesn't seem to be with IntMap per se. Checking the output of -ddump-simpl also shows that values seem to be being unboxed appropriately.