{-# LANGUAGE TupleSections #-}

module Text.Regex.Anagram.Util
  where

import           Control.Applicative (Alternative, empty)
import           Data.Foldable (foldlM)
import           Data.Function (on)
import qualified Data.IntMap.Strict as M
import qualified Data.IntSet as S
import           Data.List (group, groupBy, sortOn)
import qualified Data.Vector as V
import qualified Data.Vector.Fusion.Bundle as VB
import qualified Data.Vector.Fusion.Bundle.Size as VBS
import qualified Data.Vector.Fusion.Stream.Monadic as VS
import qualified Data.Vector.Generic as VG

import Text.Regex.Anagram.Types

guard' :: Alternative m => Bool -> a -> m a
guard' :: Bool -> a -> m a
guard' Bool
True = a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
guard' Bool
False = m a -> a -> m a
forall a b. a -> b -> a
const m a
forall (f :: * -> *) a. Alternative f => f a
empty

foldMapM :: (Monad m, Monoid b) => (a -> m b) -> [a] -> m b
-- foldMapM f = fmap fold . mapM f
foldMapM :: (a -> m b) -> [a] -> m b
foldMapM a -> m b
f = (b -> a -> m b) -> b -> [a] -> m b
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (\b
b a
a -> (b
b b -> b -> b
forall a. Semigroup a => a -> a -> a
<>) (b -> b) -> m b -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
f a
a) b
forall a. Monoid a => a
mempty

concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
-- concatMapM = foldMapM
concatMapM :: (a -> m [b]) -> [a] -> m [b]
concatMapM a -> m [b]
f = ([[b]] -> [b]) -> m [[b]] -> m [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[b]] -> [b]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (m [[b]] -> m [b]) -> ([a] -> m [[b]]) -> [a] -> m [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> m [b]) -> [a] -> m [[b]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM a -> m [b]
f

{-# INLINE withRLE #-}
withRLE :: (f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE :: (f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE f (RL a) -> g (RL b)
f = g (RL b) -> RLEof g b
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE (g (RL b) -> RLEof g b)
-> (RLEof f a -> g (RL b)) -> RLEof f a -> RLEof g b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (RL a) -> g (RL b)
f (f (RL a) -> g (RL b))
-> (RLEof f a -> f (RL a)) -> RLEof f a -> g (RL b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RLEof f a -> f (RL a)
forall (f :: * -> *) a. RLEof f a -> f (RL a)
unRLE

rleLength :: RLE a -> Int
rleLength :: RLE a -> Int
rleLength = (Int -> RL a -> Int) -> Int -> [RL a] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int
l (RL a
_ Int
r) -> Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
r) Int
0 ([RL a] -> Int) -> (RLE a -> [RL a]) -> RLE a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RLE a -> [RL a]
forall (f :: * -> *) a. RLEof f a -> f (RL a)
unRLE

rle :: Eq a => [a] -> RLE a
rle :: [a] -> RLE a
rle = [RL a] -> RLE a
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE ([RL a] -> RLE a) -> ([a] -> [RL a]) -> [a] -> RLE a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> RL a) -> [[a]] -> [RL a]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
x:[a]
l) -> a -> Int -> RL a
forall a. a -> Int -> RL a
RL a
x (Int -> Int
forall a. Enum a => a -> a
succ (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l)) ([[a]] -> [RL a]) -> ([a] -> [[a]]) -> [a] -> [RL a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group

rleV :: Eq a => V.Vector a -> RLEV a
rleV :: Vector a -> RLEV a
rleV = Vector (RL a) -> RLEV a
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE (Vector (RL a) -> RLEV a)
-> (Vector a -> Vector (RL a)) -> Vector a -> RLEV a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bundle Vector (RL a) -> Vector (RL a)
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
VG.unstream (Bundle Vector (RL a) -> Vector (RL a))
-> (Vector a -> Bundle Vector (RL a)) -> Vector a -> Vector (RL a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (m :: * -> *). Monad m => Stream m a -> Stream m (RL a))
-> (Size -> Size) -> Bundle Vector a -> Bundle Vector (RL a)
forall a b (v :: * -> *).
(forall (m :: * -> *). Monad m => Stream m a -> Stream m b)
-> (Size -> Size) -> Bundle v a -> Bundle v b
VB.inplace forall (m :: * -> *). Monad m => Stream m a -> Stream m (RL a)
forall (m :: * -> *) a.
(Monad m, Eq a) =>
Stream m a -> Stream m (RL a)
rles Size -> Size
VBS.toMax (Bundle Vector a -> Bundle Vector (RL a))
-> (Vector a -> Bundle Vector a)
-> Vector a
-> Bundle Vector (RL a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Bundle Vector a
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
VG.stream where
  rles :: Stream m a -> Stream m (RL a)
rles (VS.Stream s -> m (Step s a)
step s
st) = ((Maybe (RL a), s) -> m (Step (Maybe (RL a), s) (RL a)))
-> (Maybe (RL a), s) -> Stream m (RL a)
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
VS.Stream (Maybe (RL a), s) -> m (Step (Maybe (RL a), s) (RL a))
step' (Maybe (RL a)
forall a. Maybe a
Nothing, s
st) where
    step' :: (Maybe (RL a), s) -> m (Step (Maybe (RL a), s) (RL a))
step' (Maybe (RL a)
m, s
s) = do
      Step s a
t <- s -> m (Step s a)
step s
s
      case Step s a
t of
        VS.Yield a
x s
s' -> case Maybe (RL a)
m of
          Maybe (RL a)
Nothing -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (RL a), s) (RL a)
 -> m (Step (Maybe (RL a), s) (RL a)))
-> Step (Maybe (RL a), s) (RL a)
-> m (Step (Maybe (RL a), s) (RL a))
forall a b. (a -> b) -> a -> b
$ (Maybe (RL a), s) -> Step (Maybe (RL a), s) (RL a)
forall s a. s -> Step s a
VS.Skip (RL a -> Maybe (RL a)
forall a. a -> Maybe a
Just (a -> Int -> RL a
forall a. a -> Int -> RL a
RL a
x Int
1), s
s')
          Just r :: RL a
r@(RL a
y Int
n)
            | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (RL a), s) (RL a)
 -> m (Step (Maybe (RL a), s) (RL a)))
-> Step (Maybe (RL a), s) (RL a)
-> m (Step (Maybe (RL a), s) (RL a))
forall a b. (a -> b) -> a -> b
$ (Maybe (RL a), s) -> Step (Maybe (RL a), s) (RL a)
forall s a. s -> Step s a
VS.Skip (RL a -> Maybe (RL a)
forall a. a -> Maybe a
Just (a -> Int -> RL a
forall a. a -> Int -> RL a
RL a
x (Int -> RL a) -> Int -> RL a
forall a b. (a -> b) -> a -> b
$ Int -> Int
forall a. Enum a => a -> a
succ Int
n), s
s')
            | Bool
otherwise -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (RL a), s) (RL a)
 -> m (Step (Maybe (RL a), s) (RL a)))
-> Step (Maybe (RL a), s) (RL a)
-> m (Step (Maybe (RL a), s) (RL a))
forall a b. (a -> b) -> a -> b
$ RL a -> (Maybe (RL a), s) -> Step (Maybe (RL a), s) (RL a)
forall a s. a -> s -> Step s a
VS.Yield RL a
r (RL a -> Maybe (RL a)
forall a. a -> Maybe a
Just (a -> Int -> RL a
forall a. a -> Int -> RL a
RL a
x Int
1), s
s')
        VS.Skip s
s' -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (RL a), s) (RL a)
 -> m (Step (Maybe (RL a), s) (RL a)))
-> Step (Maybe (RL a), s) (RL a)
-> m (Step (Maybe (RL a), s) (RL a))
forall a b. (a -> b) -> a -> b
$ (Maybe (RL a), s) -> Step (Maybe (RL a), s) (RL a)
forall s a. s -> Step s a
VS.Skip (Maybe (RL a)
m, s
s')
        Step s a
VS.Done -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (RL a), s) (RL a)
 -> m (Step (Maybe (RL a), s) (RL a)))
-> Step (Maybe (RL a), s) (RL a)
-> m (Step (Maybe (RL a), s) (RL a))
forall a b. (a -> b) -> a -> b
$ Step (Maybe (RL a), s) (RL a)
-> (RL a -> Step (Maybe (RL a), s) (RL a))
-> Maybe (RL a)
-> Step (Maybe (RL a), s) (RL a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Step (Maybe (RL a), s) (RL a)
forall s a. Step s a
VS.Done (\RL a
r -> RL a -> (Maybe (RL a), s) -> Step (Maybe (RL a), s) (RL a)
forall a s. a -> s -> Step s a
VS.Yield RL a
r (Maybe (RL a)
forall a. Maybe a
Nothing, s
s)) Maybe (RL a)
m

sortRLE :: Ord a => RLE a -> RLE a
sortRLE :: RLE a -> RLE a
sortRLE = ([RL a] -> [RL a]) -> RLE a -> RLE a
forall (f :: * -> *) a (g :: * -> *) b.
(f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE (([RL a] -> [RL a]) -> RLE a -> RLE a)
-> ([RL a] -> [RL a]) -> RLE a -> RLE a
forall a b. (a -> b) -> a -> b
$ ([RL a] -> RL a) -> [[RL a]] -> [RL a]
forall a b. (a -> b) -> [a] -> [b]
map (\(RL a
x Int
r:[RL a]
l) -> a -> Int -> RL a
forall a. a -> Int -> RL a
RL a
x (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ RLE a -> Int
forall a. RLE a -> Int
rleLength ([RL a] -> RLE a
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE [RL a]
l))) ([[RL a]] -> [RL a]) -> ([RL a] -> [[RL a]]) -> [RL a] -> [RL a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RL a -> RL a -> Bool) -> [RL a] -> [[RL a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (a -> a -> Bool) -> (RL a -> a) -> RL a -> RL a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` RL a -> a
forall a. RL a -> a
unRL) ([RL a] -> [[RL a]]) -> ([RL a] -> [RL a]) -> [RL a] -> [[RL a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RL a -> a) -> [RL a] -> [RL a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn RL a -> a
forall a. RL a -> a
unRL

filterRLE :: (a -> Bool) -> RLE a -> RLE a
filterRLE :: (a -> Bool) -> RLE a -> RLE a
filterRLE a -> Bool
f = ([RL a] -> [RL a]) -> RLE a -> RLE a
forall (f :: * -> *) a (g :: * -> *) b.
(f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE (([RL a] -> [RL a]) -> RLE a -> RLE a)
-> ([RL a] -> [RL a]) -> RLE a -> RLE a
forall a b. (a -> b) -> a -> b
$ (RL a -> Bool) -> [RL a] -> [RL a]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Bool
f (a -> Bool) -> (RL a -> a) -> RL a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RL a -> a
forall a. RL a -> a
unRL)

chrStr :: [Chr] -> ChrStr
chrStr :: [Int] -> ChrStr
chrStr = (Int -> Int -> Int) -> [(Int, Int)] -> ChrStr
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
M.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ([(Int, Int)] -> ChrStr)
-> ([Int] -> [(Int, Int)]) -> [Int] -> ChrStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> (Int, Int)) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (, Int
1)

chrStrRLE :: ChrStr -> RLE Chr
chrStrRLE :: ChrStr -> RLE Int
chrStrRLE = [RL Int] -> RLE Int
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE ([RL Int] -> RLE Int) -> (ChrStr -> [RL Int]) -> ChrStr -> RLE Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> RL Int) -> [(Int, Int)] -> [RL Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Int -> RL Int) -> (Int, Int) -> RL Int
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> Int -> RL Int
forall a. a -> Int -> RL a
RL) ([(Int, Int)] -> [RL Int])
-> (ChrStr -> [(Int, Int)]) -> ChrStr -> [RL Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ChrStr -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
M.toList

nullChar :: PatChar -> Bool
nullChar :: PatChar -> Bool
nullChar (PatSet ChrSet
s) = ChrSet -> Bool
S.null ChrSet
s
nullChar PatChar
_ = Bool
False

notChar :: PatChar -> PatChar
notChar :: PatChar -> PatChar
notChar (PatChr Int
c) = ChrSet -> PatChar
PatNot (Int -> ChrSet
S.singleton Int
c)
notChar (PatSet ChrSet
s) = ChrSet -> PatChar
PatNot ChrSet
s
notChar (PatNot ChrSet
s) = ChrSet -> PatChar
PatSet ChrSet
s

intersectChrStr :: PatChar -> ChrStr -> ChrStr
intersectChrStr :: PatChar -> ChrStr -> ChrStr
intersectChrStr (PatSet ChrSet
s) ChrStr
t = ChrStr -> ChrSet -> ChrStr
forall a. IntMap a -> ChrSet -> IntMap a
M.restrictKeys ChrStr
t ChrSet
s
intersectChrStr (PatNot ChrSet
n) ChrStr
t = ChrStr -> ChrSet -> ChrStr
forall a. IntMap a -> ChrSet -> IntMap a
M.withoutKeys ChrStr
t ChrSet
n
intersectChrStr (PatChr Int
c) ChrStr
t = (Int -> ChrStr) -> Maybe Int -> ChrStr
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int -> Int -> ChrStr
forall a. Int -> a -> IntMap a
M.singleton Int
c) (Maybe Int -> ChrStr) -> Maybe Int -> ChrStr
forall a b. (a -> b) -> a -> b
$ Int -> ChrStr -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
M.lookup Int
c ChrStr
t

allChrs :: PatChar -> ChrStr -> Bool
allChrs :: PatChar -> ChrStr -> Bool
allChrs PatChar
p = ChrStr -> Bool
forall a. IntMap a -> Bool
M.null (ChrStr -> Bool) -> (ChrStr -> ChrStr) -> ChrStr -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatChar -> ChrStr -> ChrStr
intersectChrStr (PatChar -> PatChar
notChar PatChar
p)

intersectChr :: ChrSet -> PatChar -> PatChar
intersectChr :: ChrSet -> PatChar -> PatChar
intersectChr ChrSet
s p :: PatChar
p@(PatChr Int
c)
  | Int -> ChrSet -> Bool
S.member Int
c ChrSet
s = PatChar
p
  | Bool
otherwise = PatChar
forall a. Monoid a => a
mempty
intersectChr ChrSet
s (PatSet ChrSet
t) = ChrSet -> PatChar
PatSet (ChrSet -> PatChar) -> ChrSet -> PatChar
forall a b. (a -> b) -> a -> b
$ ChrSet -> ChrSet -> ChrSet
S.intersection ChrSet
s ChrSet
t
intersectChr ChrSet
s (PatNot ChrSet
n) = ChrSet -> PatChar
PatSet (ChrSet -> PatChar) -> ChrSet -> PatChar
forall a b. (a -> b) -> a -> b
$ ChrSet -> ChrSet -> ChrSet
S.difference ChrSet
s ChrSet
n

differenceChr :: ChrSet -> PatChar -> PatChar
differenceChr :: ChrSet -> PatChar -> PatChar
differenceChr ChrSet
n p :: PatChar
p@(PatChr Int
c)
  | Int -> ChrSet -> Bool
S.member Int
c ChrSet
n = PatChar
forall a. Monoid a => a
mempty
  | Bool
otherwise = PatChar
p
differenceChr ChrSet
n (PatSet ChrSet
s) = ChrSet -> PatChar
PatSet (ChrSet -> PatChar) -> ChrSet -> PatChar
forall a b. (a -> b) -> a -> b
$ ChrSet -> ChrSet -> ChrSet
S.difference ChrSet
s ChrSet
n
differenceChr ChrSet
n (PatNot ChrSet
m) = ChrSet -> PatChar
PatNot (ChrSet -> PatChar) -> ChrSet -> PatChar
forall a b. (a -> b) -> a -> b
$ ChrSet -> ChrSet -> ChrSet
S.union ChrSet
m ChrSet
n

intersectChar :: PatChar -> PatChar -> PatChar
intersectChar :: PatChar -> PatChar -> PatChar
intersectChar (PatSet ChrSet
s) PatChar
p =  ChrSet -> PatChar -> PatChar
intersectChr ChrSet
s PatChar
p
intersectChar (PatNot ChrSet
n) PatChar
p = ChrSet -> PatChar -> PatChar
differenceChr ChrSet
n PatChar
p
intersectChar p :: PatChar
p@(PatChr Int
c) (PatChr Int
d)
  | Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
d = PatChar
p
  | Bool
otherwise = PatChar
forall a. Monoid a => a
mempty
intersectChar PatChar
a PatChar
b = PatChar -> PatChar -> PatChar
intersectChar PatChar
b PatChar
a