NameSearch.hs

module Hoogle.DataBase.NameSearch
    (NameSearch, createNameSearch
    ,TextScore, searchNameSearch
    ,completionsNameSearch
    ) where

import Data.Binary.Defer
import Data.Binary.Defer.Array
import Data.Binary.Defer.Index
import qualified Data.Map as Map
import Data.Range
import General.Code
import Hoogle.Item.All
import Hoogle.TextBase.All


---------------------------------------------------------------------
-- DATA TYPES

{-
The idea is that NameItem's are sorted by name, so exact/start matching
is done by binary searching this list.

The rest of the results are taken by unioning all the suggestions in the
second element, and searching in order. All the results will end up
sorted by name (since they have identical names)

The original code was based around a Trie, gave fast performance, but
didn't merge common strings and consumed about 10x the disk space.
-}


data NameSearch = NameSearch (Array NameItem) [(Char, IntList)]

data NameItem = NameItem {key :: String
                         ,rest :: Defer [(String, [Link Entry])]}


instance Show NameSearch where
    show (NameSearch a b) =
            concat (zipWith (\a b -> show a ++ " " ++ show b) [0..] (elems a)) ++
            unlines [c : " = " ++ show d | (c,d) <- b]

instance Show NameItem where
    show (NameItem a b) = unlines $ a : map f (fromDefer b)
        where f (a,b) = unwords $ " " : a : ['#' : show (linkKey x) | x <- b]

instance BinaryDefer NameSearch where
    put (NameSearch a b) = put2 a b
    get = get2 NameSearch

instance BinaryDefer NameItem where
    put (NameItem a b) = put2 a b
    get = get2 NameItem


---------------------------------------------------------------------
-- CREATION

createNameSearch :: [Link Entry] -> NameSearch
createNameSearch xs = NameSearch (array $ Map.elems items) (Map.toList shortcuts)
    where
        items = buildItems xs
        shortcuts = buildShortcuts items


buildShortcuts :: Map.Map String NameItem -> Map.Map Char IntList
buildShortcuts = Map.map (toIntList . sort) . foldl' add Map.empty . zip [0..] . Map.keys
    where
        add mp (i,s) = foldl' g mp $ nub s
            where g mp x = Map.insertWith (++) x [i] mp


buildItems :: [Link Entry] -> Map.Map String NameItem
buildItems = Map.map norm . foldl' add Map.empty
    where
        add mp e = Map.insertWith f ltext (NameItem ltext $ Defer [(text, [e])]) mp
            where
                text = entryName $ fromLink e
                ltext = map toLower text

                f _ (NameItem a b) = NameItem a $ Defer $ g $ fromDefer b
                g [] = [(text, [e])]
                g ((x1,x2):xs) | x1 == text = (x1, e : x2) : xs
                               | otherwise = (x1,x2) : g xs

        norm (NameItem a b) = NameItem a $ Defer $ f $ fromDefer b
            where f x = sortFst [(a, sortOn linkKey b) | (a,b) <- x]


---------------------------------------------------------------------
-- SEARCHING

-- lower is better
data TextScore = TSExact | TSStart | TSNone
                 deriving (Eq,Ord)

instance Show TextScore where
    show TSExact = "exact"
    show TSStart = "start"
    show TSNone = "_"


{-
Step 1: Binary search for find the exact match
Step 2: Follow from that item finding ones which start
Step 3: Use the hint set to merge into a list of results
-}

searchNameSearch :: NameSearch -> String -> [(Link Entry,EntryView,TextScore)]
searchNameSearch (NameSearch items shortcuts) str = step1 ++ step2 ++ step3
    where
        lstr = map toLower str
        nstr = length str
        rangePrefix = FocusOn $ rangeStartCount 0 nstr

        (exact,prefix) = startPos items lstr
        (prefixes,lastpre) = followPrefixes items lstr prefix


        step1 = if isJust exact then f TSExact yes ++ f TSStart no else []
            where
                (yes,no) = partition ((==) str . fst) $ fromDefer $ rest $ items ! fromJust exact
                f scr xs = [(x, rangePrefix, scr) | x <- concatMap snd xs]

        step2 = [(x, rangePrefix, TSStart) | x <- prefixes]

        seen i = fromMaybe prefix exact <= i && i <= lastpre
        step3 = [(e,view,TSNone) | i <- xs, let x = items ! i
                , Just p <- [testMatch lstr $ key x]
                , let view = FocusOn $ rangeStartCount p nstr
                , e <- concatMap snd $ fromDefer $ rest x]
            where xs = filter (not . seen) $ intersectOrds $
                       map (maybe [] fromIntList . flip lookup shortcuts) $ nub lstr


-- Return the index of the string as the first component
-- Return the first possible index of the prefix as the second
startPos :: Array NameItem -> String -> (Maybe Int, Int)
startPos xs x = f 0 (arraySize xs - 1)
    where
        f low high | high - low < 3 = g low high
                   | otherwise =
            case compare x (key $ xs ! mid) of
                    EQ -> (Just mid, mid+1)
                    GT -> f (mid+1) high
                    LT -> f low (mid-1)
            where
                mid = (high + low) `div` 2

        g low high | low > high = (Nothing, low)
        g low high = if k == x then (Just low, low+1)
                     else if x `isPrefixOf` k then (Nothing, low)
                     else g (low+1) high
            where k = key $ xs ! low


-- Return all the items you can match following the prefix
-- Plus the last item that was a valid prefix index
followPrefixes :: Array NameItem -> String -> Int -> ([Link Entry], Int)
followPrefixes xs x i = f i
    where
        n = arraySize xs
        f i | i < n && x `isPrefixOf` key xsi = (concatMap snd (fromDefer $ rest xsi) ++ res, end)
            | otherwise = ([],i-1)
            where xsi = xs ! i
                  (res,end) = f (i+1)


testMatch :: String -> String -> Maybe Int
testMatch find within = listToMaybe [i | (i,x) <- zip [0..] $ tails within, find `isPrefixOf` x]


intersectOrd :: [Int] -> [Int] -> [Int]
intersectOrd (x:xs) (y:ys) = case compare x y of
    EQ -> x : intersectOrd xs ys
    LT -> intersectOrd xs (y:ys)
    GT -> intersectOrd (x:xs) ys
intersectOrd _ _ = []


intersectOrds :: [[Int]] -> [Int]
intersectOrds = fold1 intersectOrd


---------------------------------------------------------------------
-- COMPLETIONS

completionsNameSearch :: NameSearch -> String -> [String]
completionsNameSearch (NameSearch items _) str =
        concatMap (map fst . fromDefer . rest) $
        takeWhile ((lstr `isPrefixOf`) . key) $
        map ((!) items) [start .. arraySize items - 1]
    where
        lstr = map toLower str
        nstr = length str

        (exact,prefix) = startPos items lstr
        start = fromMaybe prefix exact


---------------------------------------------------------------------
-- IntList TYPE

type IntList = [IntRange]
data IntRange = IntRange !Int !Int

instance Show IntRange where
    show (IntRange a b) = show a ++ ".." ++ show b

instance BinaryDefer IntRange where
    put (IntRange a b) = put2 a b
    get = get2 IntRange


toIntList :: [Int] -> IntList
toIntList [] = []
toIntList (x:xs) = f x xs
    where
        f i [] = [IntRange x i]
        f i (y:ys) | y == i+1 = f y ys
                   | otherwise = IntRange x i : toIntList (y:ys)


fromIntList :: IntList -> [Int]
fromIntList = concatMap (\(IntRange a b) -> [a..b])