Core.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExistentialQuantification #-}
-- | This modules provides the function `searchSome` for searching the candidates provided by `Vector` `Text`. The information about the location of matches
-- is stored in a length-tagged unboxed vector `S.Vector`. Such vectors have an `Unbox` instances which allows us to store the collection of such mathces in an
-- unboxed `U.Vector`. This significantly reduces the memory usage and pressure on garbage collector. As a result the matchers used by this function are tagged
-- with the number @n@ of needles need to be matched and are provided by `MatcherSized`. An unsized interface is provided by `Matcher` which is existenially
-- quantified over the number of needles. Functions for constructing matching and matchers have both a sized and unsized version.

module Talash.Core ( -- * Types
                     MatcherSized (..) , Matcher (..) , MatchState (..) , MatchPart (..) , MatchFull (..) , SearchSettings (..) , Indices
                     -- * Matchers and matching
                     , makeMatcher
                     -- ** Fuzzy style
                     , fuzzyMatcherSized , fuzzyMatcher , fuzzyMatchSized , fuzzyMatch
                     -- ** Orderless Style
                     , orderlessMatcherSized , orderlessMatcher , orderlessMatchSized , orderlessMatch
                     -- * Search
                     , fuzzySettings , orderlessSettings  ,  searchSome , parts , partsOrderless , minify) where

import Control.Monad.ST (ST, runST)
import qualified Data.Text as T
import Data.Text.AhoCorasick.Automaton
import Data.Text.Utf16
import qualified Data.Vector as V
import qualified Data.Vector.Algorithms.Intro as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M
import qualified Data.Vector.Unboxed.Sized as S
import GHC.TypeNats
import Prelude (fromIntegral)
import Intro
import Lens.Micro (_1 , _2 , (^.))

-- | The MatcherSized type consists of a state machine for matching a fixed number of needles. The number of matches needed is encoded in the Nat parameterzing
--   the type. Here the purpose is to improve the memory consumption by utlizing the `Unbox` instance for sized tagged unboxed vectors from
--   (vector-sized)[https://hackage.haskell.org/package/vector-sized] package. This significantly reduces the memory consumption. At least in the present
--   implementation there is no benefit for correctness and dealing with the length tag is occasionally annoying.
data MatcherSized (n :: Nat) a = MatcherSized {
                              caseSensitivity :: CaseSensitivity ,
                              -- | An AhoCorasick state machine from the alfred-margaret package which does the actual string matching
                              machina :: {-# UNPACK #-} !(AcMachine a) ,
                              -- | The sizes of the /basic/ needles in code unit indices. The Left Int case is for when the length of all the
                              -- needles is 1 with Int the number of needles.
                              sizes :: !(Either Int (S.Vector n Int))}

-- | The existential version of MatcherSized
data Matcher a      = forall n. KnownNat n => Matcher (MatcherSized n a)

-- | The matching process essentially takes the form of a fold with possible early termination over the matches produced. See the runLower from the
--   alfred-margaret. Here MatchState is the return type of this fold and essentially it records the positions of the matches. Here like in alfred-margaret
--   position is the code unit index of the first code unit beyond the match. We can't use the CodeUnitIndex here because it doesn't have an unbox instance.
data MatchState (n :: Nat) a = MatchState {
                                 -- | This is used to record the present extent of the match. What extent means is different to different matching styles.
                                 endLocation :: {-# UNPACK #-} !Int ,
                                 -- | The vector recording the position of the matches.
                                 partialMatch :: {-# UNPACK #-} !(S.Vector n Int) ,
                                 -- | Any auxiliary information needed to describe the state of the match.
                                 aux :: !a} deriving Show

data MatchPart = MatchPart {matchBegin :: {-# UNPACK #-} !Int , matchEnd :: {-# UNPACK #-} !Int} deriving Show

-- | The full match consisting of a score for the match and vector consisting of the positions of the match. The score is intended as for bucketing and as a
--   result shouldn't be two large and must be non-negative . For the fuzzy style in this module @n@ contiguous matches contribute @n-1@ to the score. The
--   scores thus range from @0@ to @n-1@ where @n@ is the length of the string to be matched. For orderless style this score is always @0@.
data MatchFull (n :: Nat) = MatchFull {scored :: {-# UNPACK #-} !Int , indices :: {-# UNPACK #-} !(S.Vector n Int)} deriving Show

-- | The configuration for a search style with n needles and matcher of type a
data SearchSettings a (n :: Nat) = SearchSettings {
                                     -- | Given the matcher and the candidate text, find a match or return Nothing if there is none.
                                     match :: a -> Text -> Maybe (MatchFull n) ,
                                     -- | The maximum score for a given matcher. It determines the number of buckets.
                                     fullscore :: a -> Int ,
                                     -- | Maximum number of matches with full score to produce.
                                     maxFullMatches :: Int ,
                                     -- | The ordering to sort the matches within a given bucket. It is run with two candidates and their corresponding matches.
                                     orderAs :: Text -> S.Vector n Int -> Text -> S.Vector n Int -> Ordering}

-- | Type synonym for the index of a candidate in the backing vector along with the positions of the matches for it.
type Indices (n :: Nat) = (Int , S.Vector n Int)

-- | Unsafe, use with care. eIndex i return 1 for Left and for Right the element at @i@-th position in the vector. The vector must have at least @i+1@ elements.
-- This uses unsafeIndex so no bound checks are performed.
{-# INLINE eIndex #-}
eIndex :: KnownNat n => Int -> Either a (S.Vector n Int) -> Int
eIndex i = either (const 1) (`S.unsafeIndex` i)

{-# INLINEABLE updateMatch #-}
updateMatch :: KnownNat n => Int -> Either Int (S.Vector n Int) -> MatchState n a -> Int -> Int -> a -> MatchState n a
updateMatch !c l (MatchState !f !m _) !b !e !a = MatchState e (S.withVectorUnsafe (U.modify doWrites) m) a
  where
    doWrites s = void $ foldlM (\d i -> M.unsafeWrite s i d $> d - eIndex i l) c [e , e-1 .. b]

-- | The score for a fuzzy match.
{-# INLINE matchScore #-}
matchScore :: KnownNat n => Either Int (S.Vector n Int) -> S.Vector n Int -> Int
matchScore u v
  | S.length v == 0       = 0
  | v' <- S.fromSized v   = U.ifoldl' (\ !s !i !cc -> if cc - U.unsafeIndex v' i == eIndex i u then s+1 else s ) 0 . U.tail $ v'

{-# INLINEABLE matchStepFuzzy #-}
matchStepFuzzy :: KnownNat n => Either Int (S.Vector n Int) -> MatchState n () -> Match MatchPart -> Next (MatchState n ())
matchStepFuzzy l s@(MatchState !f !m _) (Match !i (MatchPart !b !e))
  | e - b == either id S.length l - 1                             = Done $ updateMatch (codeUnitIndex i) l s b e ()
  | f + 1 == b                                                    = Step $ updateMatch (codeUnitIndex i) l s b e ()
  | f >= b && e > f && monotonic                                  = Step $ updateMatch (codeUnitIndex i) l s b e ()
  | otherwise                                                     = Step   s
  where
    monotonic = S.unsafeIndex m f <= codeUnitIndex i - either (const $ e-f) (U.sum . U.slice f (e-f) . S.fromSized) l

{-# INLINEABLE matchStepOrderless #-}
matchStepOrderless :: KnownNat n => Either Int (S.Vector n Int) -> MatchState n (Int , Int) -> Match Int -> Next (MatchState n (Int , Int))
matchStepOrderless !lv s@(MatchState r !m (!lm , !li)) (Match (CodeUnitIndex !c) !i)
  | S.unsafeIndex m i == 0 && c - eIndex i lv >= li       = go   $ MatchState (r+1) (S.withVectorUnsafe (U.modify (\mv -> M.unsafeWrite mv i c)) m) (i , c)
  | S.unsafeIndex m i == 0 && eIndex lm lv < eIndex i lv  = Step $ MatchState r (S.withVectorUnsafe (U.modify (\mv -> M.unsafeWrite mv i c *> M.write mv lm 0)) m) (i , c)
  | otherwise                                             = Step s
  where
    go = if r == S.length m - 1 then Done else Step

kConsecutive :: Int ->  Text -> [Text]
kConsecutive k t = map (T.take k) . take (1 + T.length t - k) . T.tails $ t

{-# INLINE run #-}
run :: CaseSensitivity -> a -> (a -> Match v -> Next a) -> AcMachine v -> Text -> a
run IgnoreCase     = runLower
run CaseSensitive  = runText

-- | A general function to construct a Matcher. Returns Nothing if the string is empty or if the number of needles turns out to be non-positive
makeMatcher :: forall a. CaseSensitivity -> (Text -> Int) -- ^ The function to determine the number of needles from the query string.
                                                          -- The proxy argument is instantiated at the resulting value.
                    -> (forall n. KnownNat n => Proxy n -> CaseSensitivity -> Text -> MatcherSized n a) -- ^ The functions for constructing the matcher
                    -> Text -- ^ The query string
                    -> Maybe (Matcher a) -- ^ Nothing if the string is empty or if the number of needles turns out to be non-positive
makeMatcher c lenf matf t
  | T.null t || lenf t <= 0                                    = Nothing
  | SomeNat p <- someNatVal . fromIntegral. lenf $ t    = Just . Matcher . matf p c $ t

{-# INLINE withSensitivity #-}
withSensitivity :: CaseSensitivity -> Text -> Text
withSensitivity IgnoreCase    = lowerUtf16
withSensitivity CaseSensitive = id

-- | Constructs the matcher for fuzzy matching. The needles are all possible contigous subtrings of the string being matched. The Nat @n@ must be instantiated at the
--  length @n@ of the query string. They are n choose 2 such substrings, so to the complexity of matching is \(O(m + n^2)\) where @m@ is the length of candidate string.
--  This is a rough (and probably wrong) estimate as the updating the matchstate for each found match is not a constant time operation. Not sure if Aho Corasick is
--  the optimal way for this kind of matching but in practice it seems fast enough.
fuzzyMatcherSized :: KnownNat n => p n -> CaseSensitivity -> Text -> MatcherSized n MatchPart
fuzzyMatcherSized _ c t = MatcherSized {caseSensitivity = c , machina = build . concatMap go $ [T.length t , T.length t - 1 .. 1]
                             , sizes = if S.sum sz == S.length sz then Left (S.length sz) else Right sz }
  where
    sz      = fromMaybe (S.replicate 1) . S.fromList . map (length . unpackUtf16 . withSensitivity c . T.singleton)  . T.unpack  $ t
    go !k   = zipWith (\t' l -> (unpackUtf16 . withSensitivity c $ t' , MatchPart l (l + k -1))) (kConsecutive k t) [0 ..]

-- | Unsized version of fuzzyMatcherSized
fuzzyMatcher :: CaseSensitivity -> Text -> Maybe (Matcher MatchPart)
fuzzyMatcher c  = makeMatcher c T.length fuzzyMatcherSized

-- | Constructs the matcher for orderless matching, the needles are the words from the query string and the proxy argument should be instantiated at the
--  number of words.
orderlessMatcherSized :: KnownNat n => p n -> CaseSensitivity -> Text -> MatcherSized n Int
orderlessMatcherSized _ c t = MatcherSized {caseSensitivity = c , machina = build . zip (unpackUtf16 <$> wrds) $ [0 ..]
                               , sizes = Right . fromMaybe (S.replicate 1) . S.fromList . map (codeUnitIndex . lengthUtf16) $ wrds }
  where
    wrds = withSensitivity c <$> T.words t

-- | Unsized version of orderlessMatcherSized
orderlessMatcher :: CaseSensitivity -> Text -> Maybe (Matcher Int)
orderlessMatcher c = makeMatcher c (length . T.words) orderlessMatcherSized

{-# INLINEABLE fuzzyMatchSized#-}
fuzzyMatchSized :: KnownNat n => MatcherSized n MatchPart -> Text -> Maybe (MatchFull n)
fuzzyMatchSized (MatcherSized c m l) = full . run c (MatchState (-1) (S.replicate 0) ()) (matchStepFuzzy l) m
  where
    full (MatchState !e !u _) = if e + 1 == S.length u then Just $ MatchFull (matchScore l u) u else Nothing

fuzzyMatch :: Matcher MatchPart -> Text -> Maybe [Text]
fuzzyMatch (Matcher m) t = parts (S.fromSized <$> sizes m) t . S.fromSized . indices <$> fuzzyMatchSized m t

{-# INLINEABLE orderlessMatchSized#-}
orderlessMatchSized :: KnownNat n => MatcherSized n Int -> Text -> Maybe (MatchFull n)
orderlessMatchSized (MatcherSized c m l) = full . run c (MatchState 0 (S.replicate 0) (0,0)) (matchStepOrderless l) m
  where
    ln = either id S.length l
    full u = if endLocation u == ln then Just $ MatchFull 0 (partialMatch u) else Nothing

orderlessMatch :: Matcher Int -> Text -> Maybe [Text]
orderlessMatch (Matcher m) t = partsOrderless (S.fromSized <$> sizes m) t . S.fromSized . indices <$> orderlessMatchSized m t

-- | The parts of a string resulting from a match using the fuzzy matcher.
parts :: Either Int (U.Vector Int) -- ^ The information about the lengths of different needles.
  -> Text -- ^ The candidate string that has been matched
  -> U.Vector Int -- ^ The vector recording the positions of the needle in the matched string.
  -> [Text] -- ^ The candidate string split up according to  the match
parts v t = done . foldl' cut ([] , lengthUtf16 t) . minify v
  where
    done (ms , cp) = unsafeSliceUtf16 0 cp t : ms
    cut  (!ms , !cp) !cc  = (unsafeSliceUtf16 cc (cp - cc) t : ms , cc )

-- | The parts of a string resulting from a match using the orderless matcher. See parts for an explanation of arguments.
partsOrderless :: Either Int (U.Vector Int) -> Text -> U.Vector Int -> [Text]
partsOrderless v t u = parts (map (`U.backpermute` fst up) v) t (snd up)
  where
    up = U.unzip . U.modify (V.sortBy (comparing snd)) . U.imap (,) $ u

{-# INLINE eIndexU #-}
eIndexU :: Int -> Either a (U.Vector Int) -> Int
eIndexU i = either (const 1) (`U.unsafeIndex` i)

-- | Shorten a match by collapsing the contiguous sub-matches together.
minify :: Either Int (U.Vector Int) -> U.Vector Int -> [CodeUnitIndex]
minify v s
  | U.length s == 1       = map CodeUnitIndex [a , a - eIndexU 0 v]
  | otherwise             = map CodeUnitIndex . (U.last s :) . snd . U.ifoldl' go (a , [a - eIndexU 0 v]) . U.tail $ s
  where
    a = U.unsafeHead s
    go (!l , s) !i !c = if c - l <= eIndexU (i+1) v then (c , s) else (c , c - eIndexU (i+1) v : l : s )
-- | The default ordering used in this module to sort matches within a given bucket. Prefers the matche for which the last part is closest to the end. To tie
-- break prefers the shorter matched string.
defOrdering :: Text -> S.Vector n Int -> Text -> S.Vector n Int -> Ordering
defOrdering t1 s1 t2 s2
  | el == EQ   = compare (T.length t1) (T.length t2)
  | otherwise  = el
  where
    el = compare (T.length t1 - U.maximum (S.fromSized s1)) (T.length t2 - U.maximum (S.fromSized s2))

-- | Search functions suitable for fuzzy matching. The candidate @c@ will match query @s@ if @c@ contains all the characters in @s@ in order. In general there
--   can be several ways of matching. This tries to find a match with minimum number of parts of. It does not find the minimum number of parts, if that requires
--   reducing the extent of the partial match during search. E.g. matching @"as"@ against @"talash"@ the split will be @["tal","as","h"]@ and not
--   @["t","a","la","s","h"]@. While matching @"talash best match testing hat"@ against @"tea"@ will not result in @["talash best match ","te","sting h","a","t"]@ since
--   @"te"@ occurs only after we have match all three letters and we can't know if we will find the @"a"@ without going through the string.
fuzzySettings :: KnownNat n => Int -> SearchSettings (MatcherSized n MatchPart) n
fuzzySettings !m = SearchSettings { match = fuzzyMatchSized , fullscore = \t -> either id S.length (sizes t) - 1 , maxFullMatches = m , orderAs = defOrdering}

-- | Search functions that match the words in i.e. space separated substring in any order. @"talash best"@ will match @"be as"@ with the split
--   @["tal","as","h","be","st"]@ but @"talash best"@ will not match @"bet"@.
orderlessSettings :: KnownNat n => Int -> SearchSettings (MatcherSized n Int) n
orderlessSettings n = SearchSettings {match = orderlessMatchSized , fullscore = const 0, maxFullMatches = n , orderAs = defOrdering}

-- | Given a matcher, search for matches in a vector of text. This function only searches for matches among the strings at indices which are in 3rd argument.
searchSome :: forall a n. KnownNat n => SearchSettings a n -- ^ The configuration for finding matches
  -> a -- ^ The matcher
  -> V.Vector Text -- ^ The vector of candidates
  -> U.Vector Int -- ^ The subset of indices of candidates to search from
  -> (U.Vector Int , U.Vector (Indices n)) -- ^ The new set of filtered indices in the vector and the vector containing the indices for each match found.
searchSome config !t v i = runST $ uncurry (searchSomeST config t v i)
                                   =<< ((,) <$> M.unsafeNew (U.length i)
                                            <*> V.replicateM (fullscore config t + 1) (M.unsafeNew (min (U.length i) $ maxFullMatches config)))

-- This functions is somewhat ridiculous and probably over optimized. Its only purpose is to be used in searchSome. searchSome can more simply be written as
-- a mapMaybe followed by a sort but this doesn't allow for a early termination of matching once we have found enough matches. A streaming solution will probably
-- be simpler for generating the candidates but will still need to be read into a vector and sorted so for now this works and is fast enough.
searchSomeST :: forall s a n. KnownNat n => SearchSettings a n -- ^
  -> a -> V.Vector Text -> U.Vector Int -> M.STVector s Int -- ^ The mutable vector into which we write the new filtered indices
  -> V.Vector (M.STVector s (Indices n)) -- ^ The vector @mv@ containing the buckets i.e. mutable vectors one for each possible score gathering the matches
  ->  ST s (U.Vector Int , U.Vector (Indices n)) -- ^ The final vector of filtered indices and the vector of
searchSomeST config !t v i mi mv = end =<< U.foldM' go (0, U.replicate (fullscore config t + 1) 0) i
  where
    -- If we don't have enough matches with the highest score, @cc@ goes through other score in the decreasing order filling up the bucket of highest order.
    -- Stopping when either the bucket is full or else we have run through all the scores. Before filling each bucket is sorted according the ordering supplied  
    -- by the configuration. Once the filling is done this bucket is frozen see `sortTake` and returned along with the frozen vector of filtered indcies see `end`.
    end (l , u)  = (,) <$> (U.unsafeFreeze . M.unsafeTake l $ mi) <*> sortTake u
    sortTake u = (\l -> U.unsafeFreeze . M.unsafeTake l $ msv) =<< (doSort (-1) (U.unsafeLast u) *> U.ifoldM' cc (U.unsafeLast u) (U.reverse . U.unsafeInit $ u))
    msv = V.unsafeLast mv
    cc r ix n
      | r >= maxFullMatches config     = pure r
      | r + n <= maxFullMatches config = doSort ix n *> doFill r (fullscore config t - 1 - ix) n $> r + n
      | otherwise                      = doSort ix n *> doFill r (fullscore config t - 1 - ix) (maxFullMatches config - r) $> maxFullMatches config
    doFill r ix n = traverse_ (\k -> M.unsafeWrite msv (r+k) =<< M.unsafeRead (V.unsafeIndex mv ix) k) [0 .. n-1]
    doSort   ix   = V.sortByBounds cmfn (V.unsafeIndex mv (fullscore config t - 1 - ix)) 0
    cmfn (!i1,!s1) (!i2,!s2) =  orderAs config (V.unsafeIndex v i1) s1 (V.unsafeIndex v i2) s2
    bstMch = match config t . V.unsafeIndex v
    bign   = (>= maxFullMatches config)
    -- The idx is the number of total number of matches that have been written up to now. nf is the vector carrying the same information for each bucket.
    -- elm is current index to check for a match. So we check if v ! elem is a match, if it is we write the match into the appropriate bucket and elem into
    -- the vector of filtered indices  and increment idx and the appropriate index of nf. If we have already reached the needed number of matches of full score
    -- we instead just skip matching and write idx into the vector of filtered indices.
    go (!idx , !nf) !elm
      | bign (U.last nf)                                                  = M.unsafeWrite mi idx elm $> (idx + 1 , nf)
      | Just !mch <- bstMch elm , bign (U.unsafeIndex nf (scored mch))    = M.unsafeWrite mi idx elm $> (idx + 1 , nf)
      | Just (MatchFull !s !ixs) <- bstMch elm                            = M.unsafeWrite mi idx elm *> wrt s ixs $> inc s
      | otherwise                                                         = pure (idx , nf)
      where
        wrt s ixs = M.unsafeWrite (V.unsafeIndex mv s) (U.unsafeIndex nf s) (elm , ixs)
        inc s = (idx + 1 , U.modify (\mu -> M.unsafeModify mu (+1) s) nf)