diff --git a/fuzzyfind.cabal b/fuzzyfind.cabal --- a/fuzzyfind.cabal +++ b/fuzzyfind.cabal @@ -1,6 +1,6 @@ cabal-version: 2.2 name: fuzzyfind -version: 0.1.0 +version: 2.1.0 license: MIT license-file: LICENSE author: Rúnar Bjarnason @@ -10,7 +10,7 @@ bug-reports: http://github.com/runarorama/fuzzyfind/issues copyright: Copyright (C) 2021 Unison Computing build-type: Simple -tested-with: GHC == 8.8.4 +tested-with: GHC == 8.8.4, GHC == 8.10.4 synopsis: Fuzzy text matching description: A package that provides an API for fuzzy text search in Haskell, using a modified version of the Smith-Waterman algorithm. The search is intended to behave similarly to the excellent fzf tool by Junegunn Choi. category: Text @@ -23,9 +23,21 @@ type: git location: https://github.com/runarorama/fuzzyfind -library +library exposed-modules: Text.FuzzyFind other-extensions: DeriveGeneric, OverloadedLists, ScopedTypeVariables, ViewPatterns - build-depends: base ^>=4.13.0.0, array ^>=0.5.4.0, containers ^>=0.6.2.1 + build-depends: base ==4.*, massiv ==0.6.*, containers ==0.6.*, text ==1.2.* hs-source-dirs: src default-language: Haskell2010 + +executable bench + main-is: Bench.hs + ghc-options: -w -threaded -rtsopts -with-rtsopts=-N -v0 + hs-source-dirs: tests + other-modules: + default-language: Haskell2010 + build-depends: + base ==4.*, + criterion ==1.5.*, + deepseq ==1.4.*, + fuzzyfind diff --git a/src/Text/FuzzyFind.hs b/src/Text/FuzzyFind.hs --- a/src/Text/FuzzyFind.hs +++ b/src/Text/FuzzyFind.hs @@ -1,3 +1,6 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE OverloadedLists #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -18,14 +21,21 @@ module Text.FuzzyFind where import Control.Monad (join) -import Data.Array +import Data.Massiv.Array ( Array, (!), + Ix2(..), + (...), + forM, + forM_ ) -import qualified Data.Array as Array +import qualified Data.Massiv.Array as A +import qualified Data.Massiv.Array.Unsafe as A +import qualified Data.Massiv.Array.Mutable as M import Data.Char (isAlphaNum, isLower, isUpper, toLower) import Data.Foldable (maximumBy, toList, foldl') import Data.Function (on) +import Data.Maybe (fromMaybe) import Data.Sequence ( Seq (..), ViewL (..), @@ -34,11 +44,13 @@ viewr, (<|) ) -import Data.List (sortOn) import qualified Data.Sequence as Seq import GHC.Generics (Generic) +import Data.Text (Text) +import qualified Data.Text as Text +import Control.Monad.ST (runST) --- | Calling @bestMatch query string@ will return 'Nothing' if @query@ is not a +-- | @bestMatch query string@ will return 'Nothing' if @query@ is not a -- subsequence of @string@. Otherwise, it will return the "best" way to line up -- the characters in @query@ with the characters in @string@. Lower-case -- characters in the @query@ are assumed to be case-insensitive, and upper-case @@ -56,7 +68,7 @@ -- datatype), but the lowest score is @0@. -- -- A substring from the query will generate a 'Match', and any characters from --- the -- input that don't result in a 'Match' will generate a 'Gap'. +-- the input that don't result in a 'Match' will generate a 'Gap'. -- Concatenating all the 'Match' and 'Gap' results should yield the original -- input string. -- @@ -69,11 +81,12 @@ -- 1. Contiguous characters from the query string. For example, @bestMatch "pp"@ -- will find the last two ps in "pickled pepper". -- 2. Characters at the beginnings of words. For example, @bestMatch "pp"@ --- will find the first two Ps in "Peter Piper". --- 3. Characters at CamelCase humps. For example, @bestMatch "bm" "BatMan"@ +-- will find the first two Ps in \"Peter Piper\". +-- 3. Characters at CamelCase humps. For example, @bestMatch "bm" \"BatMan\"@ +-- will score higher than @bestMatch "bm" \"Batman\".@ -- 4. The algorithm strongly prefers the first character of the query pattern -- to be at the beginning of a word or CamelHump. For example, --- @bestMatch "mn" "Bat Man"@ will score higher than @bestMatch "atn" "Batman"@. +-- @bestMatch "mn" \"Bat Man\"@ will score higher than @bestMatch "atn" \"Batman\"@. -- -- All else being equal, matches that occur later in the input string are preferred. bestMatch :: String -- ^ The query pattern. @@ -88,7 +101,7 @@ defaultConsecutiveBonus -- | Finds input strings that match all the given input patterns. For each input --- that matches, it returns one 'Alignment'. The output is sorted by 'score', +-- that matches, it returns one 'Alignment'. The output is not sorted. -- ascending. -- -- For example: @@ -103,14 +116,22 @@ -- red macadamia -- * ******* -- @ -fuzzyFind :: [String] -- ^ The query patterns. - -> [String] -- ^ The input strings. - -> [Alignment] -fuzzyFind query strings = - sortOn score - $ strings - >>= (\s -> toList - $ foldl' (\a q -> (<>) <$> a <*> bestMatch q s) (Just mempty) query +fuzzyFind + :: [String] -- ^ The query patterns. + -> [String] -- ^ The input strings. + -> [Alignment] +fuzzyFind = (fmap fst .) . fuzzyFindOn id + +-- | A version of 'fuzzyFind' that searches on the given text field of the data. +fuzzyFindOn :: (a -> String) -> [String] -> [a] -> [(Alignment, a)] +fuzzyFindOn f query d = + d + >>= (\s -> + toList + $ (, s) + <$> foldl' (\a q -> (<>) <$> a <*> bestMatch q (f s)) + (Just mempty) + query ) instance Semigroup Alignment where @@ -124,7 +145,7 @@ -- | An 'Alignment' is a 'Score' together with a 'Result'. Better results have -- higher scores. data Alignment - = Alignment { score :: Score, result :: Result } + = Alignment { score :: !Score, result :: !Result } deriving (Eq, Ord, Show, Generic) -- | The base score given to a matching character @@ -151,124 +172,155 @@ defaultFirstCharBonusMultiplier = 2 -- | We prefer consecutive runs of matched characters in the pattern, so we --- impose a penalty for any gaps, proportional to the size of the gap. -defaultGapPenalty :: Int -> Int -defaultGapPenalty 1 = 3 -defaultGapPenalty n = max 0 (3 + n) +-- impose a penalty for any gaps, which is added to the size of the gap. +defaultGapPenalty :: Int +defaultGapPenalty = 3 -- | We give a bonus to consecutive matching characters. --- A number about the same as the `boundaryBonus` will strongly prefer +-- A number about the same as the boundary bonus will prefer -- runs of consecutive characters vs finding acronyms. defaultConsecutiveBonus :: Int -defaultConsecutiveBonus = defaultGapPenalty 8 +defaultConsecutiveBonus = 11 -- | Renders an 'Alignment' as a pair of lines with "*" on the lower line -- indicating the location of pattern matches. -highlight :: Alignment -> String -highlight (Alignment s (Result segments)) = - foldMap prettySegment segments <> "\n" <> foldMap showGaps segments - where - prettySegment (Gap xs) = toList xs - prettySegment (Match xs) = toList xs - showGaps (Gap xs) = replicate (length xs) ' ' - showGaps (Match xs) = replicate (length xs) '*' +-- highlight' :: Alignment -> Text +-- highlight' (Alignment s (Result segments)) = +-- foldMap prettySegment segments <> "\n" <> foldMap showGaps segments +-- where +-- prettySegment (Gap xs) = xs +-- prettySegment (Match xs) = xs +-- showGaps (Gap xs) = Text.pack $ replicate (Text.length xs) ' ' +-- showGaps (Match xs) = Text.pack $ replicate (Text.length xs) '*' +-- highlight :: Alignment -> String +-- highlight = Text.unpack . highlight' + -- | A highly configurable version of 'bestMatch'. bestMatch' :: Int -- ^ Base score for a matching character. See 'defaultMatchScore'. -> Int -- ^ Base score for a mismatched character. See 'defaultMismatchScore'. - -> (Int -> Int) -- ^ Penalty for a gap of the given length. See 'defaultGapPenalty'. + -> Int -- ^ Additional penalty for a gap. See 'defaultGapPenalty'. -> Int -- ^ Bonus score for a match at the beginning of a word. See 'defaultBoundaryBonus'. -> Int -- ^ Bonus score for a match on a CamelCase hump. See 'defaultCamelCaseBonus'. -> Int -- ^ Bonus multiplier for matching the first character of the pattern. -- See 'defaultFirstCharBonusMultiplier'. - -> Int -- ^ Bonus score for each consecutive character matched. See 'defaultFirstCharBonusMultiplier'. + -> Int -- ^ Bonus score for each consecutive character matched. + -- See 'defaultFirstCharBonusMultiplier'. -> String -- ^ The query pattern. -> String -- ^ The input string. -> Maybe Alignment bestMatch' matchScore mismatchScore gapPenalty boundaryBonus camelCaseBonus firstCharBonusMultiplier consecutiveBonus query str - = Alignment (totalScore m nx) . reverseResult <$> traceback + = Alignment (totalScore m nx) . Result <$> traceback where - totalScore i j = if i > m then 0 else hs ! (i, j) + bonuses ! (i, j) - table = unlines - [ unwords - $ (if y > 0 then show $ b' ! y else " ") - : [ show (totalScore x y) | x <- [0 .. m] ] - | y <- [0 .. n] - ] + totalScore i j = + if i > m then 0 else (A.index' hs (i :. j)) + (A.index' bonuses (i :. j)) + -- table = unlines + -- [ unwords + -- $ (if y > 0 then show $ str' ! y else " ") + -- : [ show (totalScore x y) | x <- [0 .. m] ] + -- | y <- [0 .. n] + -- ] similarity a b = if a == b || a == toLower b then matchScore else mismatchScore - traceback = (gaps (drop nx str) <>) <$> go (m, nx) - go (0, j) = pure $ gaps (take j str) - go (i, 0) = Nothing - go (i, j) = if similarity (a' ! i) (b' ! j) > 0 - then (match (b' ! j) <>) <$> go (i - 1, j - 1) - else (gap (b' ! j) <>) <$> go (i, j - 1) - nx = localMax m n - localMax m n = maximumBy - (\b d -> compare (totalScore m b) (totalScore m d)) - [ j | j <- [1 .. n] ] - m = length query - n = length str - a' = Array.listArray (1, m) query - b' = Array.listArray (1, n) str - hs = Array.listArray bounds [ h i j | (i, j) <- Array.range bounds ] - bonuses = Array.listArray bounds [ bonus i j | (i, j) <- Array.range bounds ] - bounds = ((0, 0), (m, n)) + traceback :: Maybe (Seq ResultSegment) + traceback = (<> gaps (drop nx str)) <$> go [] [] (-1) m nx + go r m currOp 0 j = (gaps (take j str) <>) <$> case m of + [] -> Just r + _ -> case currOp of + 1 -> Just (r :|> Match (reverse m)) + 0 -> Just (r :|> Gap (reverse m)) + -1 -> Nothing + go _ _ _ _ 0 = Nothing + go r m currOp i j = + if similarity (A.index' query' (i - 1)) (A.index' str' (j - 1)) > 0 + then case currOp of + 0 -> + go (r :|> Gap (reverse m)) [A.index' str' (j - 1)] 1 (i - 1) (j - 1) + _ -> go r (A.index' str' (j - 1) : m) 1 (i - 1) (j - 1) + else case currOp of + 1 -> go (r :|> Match (reverse m)) [A.index' str' (j - 1)] 0 i (j - 1) + _ -> go r (A.index' str' (j - 1) : m) 0 i (j - 1) + nx = localMax m n 1 0 0 + localMax m n j r s = if j > n + then r + else + let s' = totalScore m j + in localMax m n (j + 1) (if s' > s then j else r) s' + query' = A.fromList A.Seq query :: Array A.U A.Ix1 Char + str' = A.fromList A.Seq str :: Array A.U A.Ix1 Char + m = A.unSz $ A.size query' + n = A.unSz $ A.size str' + hs :: Array A.U Ix2 Int + hs = M.createArrayST_ (A.Sz (m + 1 :. n + 1)) $ \marr -> do + A.forM_ ((0 :. 0) ... (m :. n)) $ \(i :. j) -> if (i == 0 || j == 0) + then M.writeM marr (i :. j) 0 + else do + scoreMatch <- do + hprev <- M.readM marr ((i - 1) :. (j - 1)) + pure + $ hprev + + similarity (A.index' query' (i - 1)) (A.index' str' (j - 1)) + + A.index' bonuses (i :. j) + scoreGap <- do + (arr :: Array A.U A.Ix1 Int) <- forM (1 ... j) $ \l -> + (\x -> x - (l + gapPenalty)) <$> M.readM marr (i :. (j - l)) + pure . fromMaybe 0 $ A.maximumM arr + M.writeM marr (i :. j) (scoreMatch `max` scoreGap `max` 0) + bonuses = A.makeArray A.Seq (A.Sz (m + 1 :. n + 1)) f :: Array A.U Ix2 Int + where f (i :. j) = bonus i j + bonus :: Int -> Int -> Int bonus 0 j = 0 bonus i 0 = 0 - bonus i j = if similarity (a' ! i) (b' ! j) > 0 - then multiplier * (boundary + camel + consecutive) - else 0 + bonus i j = + if similarity (A.index' query' (i - 1)) (A.index' str' (j - 1)) > 0 + then multiplier * (boundary + camel + consecutive) + else 0 where boundary = - if j < 2 || isAlphaNum (b' ! j) && not (isAlphaNum (b' ! (j - 1))) + if j < 2 || isAlphaNum (A.index' str' (j - 1)) && not + (isAlphaNum (A.index' str' (j - 2))) then boundaryBonus else 0 - camel = if j > 1 && isLower (b' ! (j - 1)) && isUpper (b' ! j) - then camelCaseBonus - else 0 + camel = + if j > 1 && isLower (A.index' str' (j - 2)) && isUpper + (A.index' str' (j - 1)) + then + camelCaseBonus + else + 0 multiplier = if i == 1 then firstCharBonusMultiplier else 1 consecutive = let - similar = i > 0 && j > 0 && similarity (a' ! i) (b' ! j) > 0 + similar = + i + > 0 + && j + > 0 + && similarity (A.index' query' (i - 1)) (A.index' str' (j - 1)) + > 0 afterMatch = - i > 1 && j > 1 && similarity (a' ! (i - 1)) (b' ! (j - 1)) > 0 + i + > 1 + && j + > 1 + && similarity (A.index' query' (i - 2)) (A.index' str' (j - 2)) + > 0 beforeMatch = - i < m && j < n && similarity (a' ! (i + 1)) (b' ! (j + 1)) > 0 + i < m && j < n && similarity (A.index' query' i) (A.index' str' j) > 0 in if similar && (afterMatch || beforeMatch) then consecutiveBonus else 0 - h 0 _ = 0 - h _ 0 = 0 - h i j = scoreMatch `max` scoreGap `max` 0 - where - scoreMatch = - hs ! (i - 1, j - 1) + similarity (a' ! i) (b' ! j) + bonuses ! (i, j) - scoreGap = maximum [ hs ! (i, j - l) - gapPenalty l | l <- [1 .. j] ] -data ResultSegment = Gap (Seq Char) | Match (Seq Char) +gaps :: String -> Seq ResultSegment +gaps s = [Gap s] + +data ResultSegment = Gap !String | Match !String deriving (Eq, Ord, Show, Generic) -- | Concatenating all the 'ResultSegment's should yield the original input string. newtype Result = Result { segments :: Seq ResultSegment } deriving (Eq, Ord, Show, Generic) -match :: Char -> Result -match a = Result [Match [a]] - -gap :: Char -> Result -gap a = Result [Gap [a]] - -gaps :: String -> Result -gaps s = Result [Gap . Seq.fromList $ reverse s] - -reverseResult :: Result -> Result -reverseResult (Result xs) = Result . Seq.reverse $ reverseSegment <$> xs - -reverseSegment :: ResultSegment -> ResultSegment -reverseSegment (Gap xs) = Gap (Seq.reverse xs) -reverseSegment (Match xs) = Match (Seq.reverse xs) - instance Monoid Result where mempty = Result [] @@ -291,25 +343,24 @@ drop' :: Int -> Result -> Result drop' n m | n < 1 = m drop' n (Result (viewl -> Gap g :< t)) = - Result [Gap (Seq.drop n g)] <> drop' (n - Seq.length g) (Result t) + Result [Gap (drop n g)] <> drop' (n - length g) (Result t) drop' n (Result (viewl -> Match g :< t)) = - Result [Match (Seq.drop n g)] <> drop' (n - Seq.length g) (Result t) + Result [Match (drop n g)] <> drop' (n - length g) (Result t) merge :: Result -> Result -> Result merge (Result Seq.Empty) ys = ys merge xs (Result Seq.Empty) = xs merge (Result xs) (Result ys ) = case (viewl xs, viewl ys) of (Gap g :< t, Gap g' :< t') - | Seq.length g <= Seq.length g' -> Result [Gap g] - <> merge (Result t) (drop' (Seq.length g) (Result ys)) + | length g <= length g' -> Result [Gap g] + <> merge (Result t) (drop' (length g) (Result ys)) | otherwise -> Result [Gap g'] - <> merge (drop' (Seq.length g') (Result xs)) (Result t') + <> merge (drop' (length g') (Result xs)) (Result t') (Match m :< t, Match m' :< t') - | Seq.length m >= Seq.length m' -> Result [Match m] - <> merge (Result t) (drop' (Seq.length m) (Result ys)) + | length m >= length m' -> Result [Match m] + <> merge (Result t) (drop' (length m) (Result ys)) | otherwise -> Result [Match m'] - <> merge (drop' (Seq.length m') (Result xs)) (Result t') + <> merge (drop' (length m') (Result xs)) (Result t') (Gap g :< t, Match m' :< t') -> - Result [Match m'] <> merge (drop' (Seq.length m') (Result xs)) (Result t') + Result [Match m'] <> merge (drop' (length m') (Result xs)) (Result t') (Match m :< t, Gap g' :< t') -> - Result [Match m] <> merge (Result t) (drop' (Seq.length m) (Result ys)) - + Result [Match m] <> merge (Result t) (drop' (length m) (Result ys)) diff --git a/tests/Bench.hs b/tests/Bench.hs new file mode 100644 --- /dev/null +++ b/tests/Bench.hs @@ -0,0 +1,28 @@ +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# OPTIONS_GHC -fno-warn-orphans -fno-warn-type-defaults #-} + +module Main where + +import Criterion.Main +import Text.FuzzyFind +import Control.DeepSeq + +deriving instance NFData Alignment +deriving instance NFData Result +deriving instance NFData ResultSegment + +main :: IO () +main = defaultMain + [ env lipsumFile $ \lipsum -> + bgroup "fzf" + $ (\n -> bench (show (4 ^ n)) . nf (fuzzyFind ["lipsum"]) $ take + (4 ^ n) + lipsum + ) + <$> [0 .. 6] + ] + +lipsumFile :: IO [String] +lipsumFile = ((take 50 <$>) . lines) <$> readFile "tests/lipsum.txt" +