-- | Parallel versions of 'filter' and 'simpleFilter'
module Text.Fuzzy.Parallel
(   filter,
    simpleFilter,
    Scored(..),
    -- reexports
    Fuzzy,
) where

import           Control.Monad.ST            (runST)
import           Control.Parallel.Strategies (Eval, Strategy, evalTraversable,
                                              parTraversable, rseq, using)
import           Data.Monoid.Textual         (TextualMonoid)
import           Data.Vector                 (Vector, (!))
import qualified Data.Vector                 as V
-- need to use a stable sort
import           Data.Bifunctor              (second)
import           Data.Char                   (toLower)
import           Data.Maybe                  (fromMaybe)
import qualified Data.Monoid.Textual         as T
import           Prelude                     hiding (filter)
import           Text.Fuzzy                  (Fuzzy (..))

data Scored a = Scored {Scored a -> Int
score_ :: !Int, Scored a -> a
original:: !a}
  deriving (a -> Scored b -> Scored a
(a -> b) -> Scored a -> Scored b
(forall a b. (a -> b) -> Scored a -> Scored b)
-> (forall a b. a -> Scored b -> Scored a) -> Functor Scored
forall a b. a -> Scored b -> Scored a
forall a b. (a -> b) -> Scored a -> Scored b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> Scored b -> Scored a
$c<$ :: forall a b. a -> Scored b -> Scored a
fmap :: (a -> b) -> Scored a -> Scored b
$cfmap :: forall a b. (a -> b) -> Scored a -> Scored b
Functor,Int -> Scored a -> ShowS
[Scored a] -> ShowS
Scored a -> String
(Int -> Scored a -> ShowS)
-> (Scored a -> String) -> ([Scored a] -> ShowS) -> Show (Scored a)
forall a. Show a => Int -> Scored a -> ShowS
forall a. Show a => [Scored a] -> ShowS
forall a. Show a => Scored a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Scored a] -> ShowS
$cshowList :: forall a. Show a => [Scored a] -> ShowS
show :: Scored a -> String
$cshow :: forall a. Show a => Scored a -> String
showsPrec :: Int -> Scored a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Scored a -> ShowS
Show)

-- | Returns the rendered output and the
-- matching score for a pattern and a text.
-- Two examples are given below:
--
-- >>> match "fnt" "infinite" "" "" id True
-- Just ("infinite",3)
--
-- >>> match "hsk" ("Haskell",1995) "<" ">" fst False
-- Just ("<h>a<s><k>ell",5)
--
{-# INLINABLE match #-}

match :: (T.TextualMonoid s)
      => s        -- ^ Pattern in lowercase except for first character
      -> t        -- ^ The value containing the text to search in.
      -> s        -- ^ The text to add before each match.
      -> s        -- ^ The text to add after each match.
      -> (t -> s) -- ^ The function to extract the text from the container.
      -> Maybe (Fuzzy t s) -- ^ The original value, rendered string and score.
match :: s -> t -> s -> s -> (t -> s) -> Maybe (Fuzzy t s)
match s
pattern t
t s
pre s
post t -> s
extract =
    if s -> Bool
forall s. TextualMonoid s => s -> Bool
null s
pat then Fuzzy t s -> Maybe (Fuzzy t s)
forall a. a -> Maybe a
Just (t -> s -> Int -> Fuzzy t s
forall t s. t -> s -> Int -> Fuzzy t s
Fuzzy t
t s
result Int
totalScore) else Maybe (Fuzzy t s)
forall a. Maybe a
Nothing
  where
    null :: (T.TextualMonoid s) => s -> Bool
    null :: s -> Bool
null = Bool -> Bool
not (Bool -> Bool) -> (s -> Bool) -> s -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> s -> Bool
forall t. TextualMonoid t => (Char -> Bool) -> t -> Bool
T.any (Bool -> Char -> Bool
forall a b. a -> b -> a
const Bool
True)

    s :: s
s = t -> s
extract t
t
    (Int
totalScore, Int
_currScore, s
result, s
pat, Bool
_) =
      ((Int, Int, s, s, Bool) -> s -> (Int, Int, s, s, Bool))
-> ((Int, Int, s, s, Bool) -> Char -> (Int, Int, s, s, Bool))
-> (Int, Int, s, s, Bool)
-> s
-> (Int, Int, s, s, Bool)
forall t a.
TextualMonoid t =>
(a -> t -> a) -> (a -> Char -> a) -> a -> t -> a
T.foldl'
        (Int, Int, s, s, Bool) -> s -> (Int, Int, s, s, Bool)
forall a. HasCallStack => a
undefined
        (\(Int
tot, Int
cur, s
res, s
pat, Bool
isFirst) Char
c ->
            case s -> Maybe (Char, s)
forall t. TextualMonoid t => t -> Maybe (Char, t)
T.splitCharacterPrefix s
pat of
              Maybe (Char, s)
Nothing -> (Int
tot, Int
0, s
res s -> s -> s
forall a. Semigroup a => a -> a -> a
<> Char -> s
forall t. TextualMonoid t => Char -> t
T.singleton Char
c, s
pat, Bool
isFirst)
              Just (Char
x, s
xs) ->
                -- the case of the first character has to match
                -- otherwise use lower case since the pattern is assumed lower
                let !c' :: Char
c' = if Bool
isFirst then Char
c else Char -> Char
toLower Char
c in
                if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
c' then
                  let cur' :: Int
cur' = Int
cur Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 in
                  (Int
tot Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
cur', Int
cur', s
res s -> s -> s
forall a. Semigroup a => a -> a -> a
<> s
pre s -> s -> s
forall a. Semigroup a => a -> a -> a
<> Char -> s
forall t. TextualMonoid t => Char -> t
T.singleton Char
c s -> s -> s
forall a. Semigroup a => a -> a -> a
<> s
post, s
xs, Bool
False)
                else (Int
tot, Int
0, s
res s -> s -> s
forall a. Semigroup a => a -> a -> a
<> Char -> s
forall t. TextualMonoid t => Char -> t
T.singleton Char
c, s
pat, Bool
isFirst)
        ) ( Int
0
          , Int
1 -- matching at the start gives a bonus (cur = 1)
          , s
forall a. Monoid a => a
mempty, s
pattern, Bool
True) s
s

-- | The function to filter a list of values by fuzzy search on the text extracted from them.
filter :: (TextualMonoid s)
       => Int      -- ^ Chunk size. 1000 works well.
       -> Int      -- ^ Max. number of results wanted
       -> s        -- ^ Pattern.
       -> [t]      -- ^ The list of values containing the text to search in.
       -> s        -- ^ The text to add before each match.
       -> s        -- ^ The text to add after each match.
       -> (t -> s) -- ^ The function to extract the text from the container.
       -> [Scored t] -- ^ The list of results, sorted, highest score first.
filter :: Int -> Int -> s -> [t] -> s -> s -> (t -> s) -> [Scored t]
filter Int
chunkSize Int
maxRes s
pattern [t]
ts s
pre s
post t -> s
extract = (forall s. ST s [Scored t]) -> [Scored t]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Scored t]) -> [Scored t])
-> (forall s. ST s [Scored t]) -> [Scored t]
forall a b. (a -> b) -> a -> b
$ do
  let v :: Vector (Fuzzy t s)
v = (Maybe (Fuzzy t s) -> Maybe (Fuzzy t s))
-> Vector (Maybe (Fuzzy t s)) -> Vector (Fuzzy t s)
forall a b. (a -> Maybe b) -> Vector a -> Vector b
V.mapMaybe Maybe (Fuzzy t s) -> Maybe (Fuzzy t s)
forall a. a -> a
id
             ((t -> Maybe (Fuzzy t s)) -> Vector t -> Vector (Maybe (Fuzzy t s))
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\t
t -> s -> t -> s -> s -> (t -> s) -> Maybe (Fuzzy t s)
forall s t.
TextualMonoid s =>
s -> t -> s -> s -> (t -> s) -> Maybe (Fuzzy t s)
match s
pattern' t
t s
pre s
post t -> s
extract) ([t] -> Vector t
forall a. [a] -> Vector a
V.fromList [t]
ts)
             Vector (Maybe (Fuzzy t s))
-> Strategy (Vector (Maybe (Fuzzy t s)))
-> Vector (Maybe (Fuzzy t s))
forall a. a -> Strategy a -> a
`using`
             Int
-> Strategy (Maybe (Fuzzy t s))
-> Strategy (Vector (Maybe (Fuzzy t s)))
forall a. Int -> Strategy a -> Vector a -> Eval (Vector a)
parVectorChunk Int
chunkSize (Strategy (Fuzzy t s) -> Strategy (Maybe (Fuzzy t s))
forall (t :: * -> *) a.
Traversable t =>
Strategy a -> Strategy (t a)
evalTraversable Strategy (Fuzzy t s)
forall s t. TextualMonoid s => Fuzzy t s -> Eval (Fuzzy t s)
forceScore))
      perfectScore :: Int
perfectScore = Fuzzy s s -> Int
forall t s. TextualMonoid s => Fuzzy t s -> Int
score (Fuzzy s s -> Int) -> Fuzzy s s -> Int
forall a b. (a -> b) -> a -> b
$ Fuzzy s s -> Maybe (Fuzzy s s) -> Fuzzy s s
forall a. a -> Maybe a -> a
fromMaybe (String -> Fuzzy s s
forall a. HasCallStack => String -> a
error (String -> Fuzzy s s) -> String -> Fuzzy s s
forall a b. (a -> b) -> a -> b
$ (s -> String) -> s -> String
forall t. TextualMonoid t => (t -> String) -> t -> String
T.toString s -> String
forall a. HasCallStack => a
undefined s
pattern) (Maybe (Fuzzy s s) -> Fuzzy s s) -> Maybe (Fuzzy s s) -> Fuzzy s s
forall a b. (a -> b) -> a -> b
$
        s -> s -> s -> s -> (s -> s) -> Maybe (Fuzzy s s)
forall s t.
TextualMonoid s =>
s -> t -> s -> s -> (t -> s) -> Maybe (Fuzzy t s)
match s
pattern' s
pattern' s
"" s
"" s -> s
forall a. a -> a
id
  [Scored t] -> ST s [Scored t]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Scored t] -> ST s [Scored t]) -> [Scored t] -> ST s [Scored t]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector (Fuzzy t s) -> [Scored t]
forall s t.
TextualMonoid s =>
Int -> Int -> Vector (Fuzzy t s) -> [Scored t]
partialSortByAscScore Int
maxRes Int
perfectScore Vector (Fuzzy t s)
v
  where
      -- Preserve case for the first character, make all others lowercase
      pattern' :: s
pattern' = case s -> Maybe (Char, s)
forall t. TextualMonoid t => t -> Maybe (Char, t)
T.splitCharacterPrefix s
pattern of
          Just (Char
c, s
rest) -> Char -> s
forall t. TextualMonoid t => Char -> t
T.singleton Char
c s -> s -> s
forall a. Semigroup a => a -> a -> a
<> (Char -> Char) -> s -> s
forall t. TextualMonoid t => (Char -> Char) -> t -> t
T.map Char -> Char
toLower s
rest
          Maybe (Char, s)
_              -> s
pattern

-- | Return all elements of the list that have a fuzzy
-- match against the pattern. Runs with default settings where
-- nothing is added around the matches, as case insensitive.
--
-- >>> simpleFilter "vm" ["vim", "emacs", "virtual machine"]
-- ["vim","virtual machine"]
{-# INLINABLE simpleFilter #-}
simpleFilter :: (TextualMonoid s)
             => Int -- ^ Chunk size. 1000 works well.
             -> Int -- ^ Max. number of results wanted
             -> s   -- ^ Pattern to look for.
             -> [s] -- ^ List of texts to check.
             -> [Scored s] -- ^ The ones that match.
simpleFilter :: Int -> Int -> s -> [s] -> [Scored s]
simpleFilter Int
chunk Int
maxRes s
pattern [s]
xs =
  Int -> Int -> s -> [s] -> s -> s -> (s -> s) -> [Scored s]
forall s t.
TextualMonoid s =>
Int -> Int -> s -> [t] -> s -> s -> (t -> s) -> [Scored t]
filter Int
chunk Int
maxRes s
pattern [s]
xs s
forall a. Monoid a => a
mempty s
forall a. Monoid a => a
mempty s -> s
forall a. a -> a
id

--------------------------------------------------------------------------------

-- | Evaluation that forces the 'score' field
forceScore :: TextualMonoid s => Fuzzy t s -> Eval(Fuzzy t s)
forceScore :: Fuzzy t s -> Eval (Fuzzy t s)
forceScore it :: Fuzzy t s
it@Fuzzy{Int
score :: Int
score :: forall t s. TextualMonoid s => Fuzzy t s -> Int
score} = do
  Int
score' <- Strategy Int
forall a. Strategy a
rseq Int
score
  Fuzzy t s -> Eval (Fuzzy t s)
forall (m :: * -> *) a. Monad m => a -> m a
return Fuzzy t s
it{score :: Int
score = Int
score'}

--------------------------------------------------------------------------------

-- | Divides a vector in chunks, applies the strategy in parallel to each chunk.
parVectorChunk :: Int -> Strategy a -> Vector a -> Eval (Vector a)
parVectorChunk :: Int -> Strategy a -> Vector a -> Eval (Vector a)
parVectorChunk Int
chunkSize Strategy a
st Vector a
v =
    [Vector a] -> Vector a
forall a. [Vector a] -> Vector a
V.concat ([Vector a] -> Vector a) -> Eval [Vector a] -> Eval (Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Vector a -> Eval (Vector a)) -> Strategy [Vector a]
forall (t :: * -> *) a.
Traversable t =>
Strategy a -> Strategy (t a)
parTraversable (Strategy a -> Vector a -> Eval (Vector a)
forall (t :: * -> *) a.
Traversable t =>
Strategy a -> Strategy (t a)
evalTraversable Strategy a
st) (Int -> Vector a -> [Vector a]
forall a. Int -> Vector a -> [Vector a]
chunkVector Int
chunkSize Vector a
v)

-- >>> chunkVector 3 (V.fromList [0..10])
-- >>> chunkVector 3 (V.fromList [0..11])
-- >>> chunkVector 3 (V.fromList [0..12])
-- [[0,1,2],[3,4,5],[6,7,8],[9,10]]
-- [[0,1,2],[3,4,5],[6,7,8],[9,10,11]]
-- [[0,1,2],[3,4,5],[6,7,8],[9,10,11],[12]]
chunkVector :: Int -> Vector a -> [Vector a]
chunkVector :: Int -> Vector a -> [Vector a]
chunkVector Int
chunkSize Vector a
v = do
    let indices :: [(Int, Int)]
indices = Int -> (Int, Int) -> [(Int, Int)]
chunkIndices Int
chunkSize (Int
0,Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
v)
    [Int -> Int -> Vector a -> Vector a
forall a. Int -> Int -> Vector a -> Vector a
V.slice Int
l (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Vector a
v | (Int
l,Int
h) <- [(Int, Int)]
indices]

-- >>> chunkIndices 3 (0,9)
-- >>> chunkIndices 3 (0,10)
-- >>> chunkIndices 3 (0,11)
-- [(0,2),(3,5),(6,8)]
-- [(0,2),(3,5),(6,8),(9,9)]
-- [(0,2),(3,5),(6,8),(9,10)]
chunkIndices :: Int -> (Int,Int) -> [(Int,Int)]
chunkIndices :: Int -> (Int, Int) -> [(Int, Int)]
chunkIndices Int
chunkSize (Int
from,Int
to) =
  ((Int, Int) -> (Int, Int)) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Int) -> (Int, Int) -> (Int, Int)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Int -> Int
forall a. Enum a => a -> a
pred) ([(Int, Int)] -> [(Int, Int)]) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
  [Int] -> [(Int, Int)]
forall a. [a] -> [(a, a)]
pairwise ([Int] -> [(Int, Int)]) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$
  [Int
from, Int
fromInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
chunkSize .. Int
toInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
to]

pairwise :: [a] -> [(a,a)]
pairwise :: [a] -> [(a, a)]
pairwise []       = []
pairwise [a
_]      = []
pairwise (a
x:a
y:[a]
xs) = (a
x,a
y) (a, a) -> [(a, a)] -> [(a, a)]
forall a. a -> [a] -> [a]
: [a] -> [(a, a)]
forall a. [a] -> [(a, a)]
pairwise (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)

-- | A stable partial sort ascending by score. O(N) best case, O(wanted*N) worst case
partialSortByAscScore :: TextualMonoid s
            => Int  -- ^ Number of items needed
            -> Int  -- ^ Value of a perfect score
            -> Vector (Fuzzy t s)
            -> [Scored t]
partialSortByAscScore :: Int -> Int -> Vector (Fuzzy t s) -> [Scored t]
partialSortByAscScore Int
wantedCount Int
perfectScore Vector (Fuzzy t s)
v = Int -> SortState Any -> [Scored t] -> [Scored t]
forall a. Int -> SortState a -> [Scored t] -> [Scored t]
loop Int
0 (Int -> Int -> Int -> SortState Any
forall a. Int -> Int -> Int -> SortState a
SortState Int
forall a. Bounded a => a
minBound Int
perfectScore Int
0) [] where
  l :: Int
l = Vector (Fuzzy t s) -> Int
forall a. Vector a -> Int
V.length Vector (Fuzzy t s)
v
  loop :: Int -> SortState a -> [Scored t] -> [Scored t]
loop Int
index st :: SortState a
st@SortState{Int
foundCount :: forall a. SortState a -> Int
scoreWanted :: forall a. SortState a -> Int
bestScoreSeen :: forall a. SortState a -> Int
foundCount :: Int
scoreWanted :: Int
bestScoreSeen :: Int
..} [Scored t]
acc
    | Int
foundCount Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
wantedCount = [Scored t] -> [Scored t]
forall a. [a] -> [a]
reverse [Scored t]
acc
    | Int
index Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l
-- ProgressCancelledException
    = if Int
bestScoreSeen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
scoreWanted
        then Int -> SortState a -> [Scored t] -> [Scored t]
loop Int
0 SortState a
st{scoreWanted :: Int
scoreWanted = Int
bestScoreSeen, bestScoreSeen :: Int
bestScoreSeen = Int
forall a. Bounded a => a
minBound} [Scored t]
acc
        else [Scored t] -> [Scored t]
forall a. [a] -> [a]
reverse [Scored t]
acc
    | Bool
otherwise =
      case Vector (Fuzzy t s)
vVector (Fuzzy t s) -> Int -> Fuzzy t s
forall a. Vector a -> Int -> a
!Int
index of
        Fuzzy t s
x | Fuzzy t s -> Int
forall t s. TextualMonoid s => Fuzzy t s -> Int
score Fuzzy t s
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
scoreWanted
          -> Int -> SortState a -> [Scored t] -> [Scored t]
loop (Int
indexInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SortState a
st{foundCount :: Int
foundCount = Int
foundCountInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1} (Fuzzy t s -> Scored t
forall s t. TextualMonoid s => Fuzzy t s -> Scored t
toScored Fuzzy t s
xScored t -> [Scored t] -> [Scored t]
forall a. a -> [a] -> [a]
:[Scored t]
acc)
          | Fuzzy t s -> Int
forall t s. TextualMonoid s => Fuzzy t s -> Int
score Fuzzy t s
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
scoreWanted Bool -> Bool -> Bool
&& Fuzzy t s -> Int
forall t s. TextualMonoid s => Fuzzy t s -> Int
score Fuzzy t s
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bestScoreSeen
          -> Int -> SortState a -> [Scored t] -> [Scored t]
loop (Int
indexInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SortState a
st{bestScoreSeen :: Int
bestScoreSeen = Fuzzy t s -> Int
forall t s. TextualMonoid s => Fuzzy t s -> Int
score Fuzzy t s
x} [Scored t]
acc
          | Bool
otherwise
          -> Int -> SortState a -> [Scored t] -> [Scored t]
loop (Int
indexInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) SortState a
st [Scored t]
acc

toScored :: TextualMonoid s => Fuzzy t s -> Scored t
toScored :: Fuzzy t s -> Scored t
toScored Fuzzy{s
t
Int
original :: forall t s. TextualMonoid s => Fuzzy t s -> t
rendered :: forall t s. TextualMonoid s => Fuzzy t s -> s
score :: Int
rendered :: s
original :: t
score :: forall t s. TextualMonoid s => Fuzzy t s -> Int
..} = Int -> t -> Scored t
forall a. Int -> a -> Scored a
Scored Int
score t
original

data SortState a = SortState
  { SortState a -> Int
bestScoreSeen :: !Int
  , SortState a -> Int
scoreWanted   :: !Int
  , SortState a -> Int
foundCount    :: !Int
  }
  deriving Int -> SortState a -> ShowS
[SortState a] -> ShowS
SortState a -> String
(Int -> SortState a -> ShowS)
-> (SortState a -> String)
-> ([SortState a] -> ShowS)
-> Show (SortState a)
forall a. Int -> SortState a -> ShowS
forall a. [SortState a] -> ShowS
forall a. SortState a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SortState a] -> ShowS
$cshowList :: forall a. [SortState a] -> ShowS
show :: SortState a -> String
$cshow :: forall a. SortState a -> String
showsPrec :: Int -> SortState a -> ShowS
$cshowsPrec :: forall a. Int -> SortState a -> ShowS
Show