{-# 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' :: forall (m :: * -> *) a. Alternative m => Bool -> a -> m a
guard' Bool
True = a -> m a
forall a. 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 a. 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 :: forall (m :: * -> *) b a.
(Monad m, Monoid b) =>
(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 :: forall (m :: * -> *) a b. Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM a -> m [b]
f = ([[b]] -> [b]) -> m [[b]] -> m [b]
forall a b. (a -> b) -> m a -> 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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM a -> m [b]
f

{-# INLINE withRLE #-}
withRLE :: (f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE :: forall (f :: * -> *) a (g :: * -> *) b.
(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 :: forall a. RLE a -> Chr
rleLength = (Chr -> RL a -> Chr) -> Chr -> [RL a] -> Chr
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Chr
l (RL a
_ Chr
r) -> Chr
l Chr -> Chr -> Chr
forall a. Num a => a -> a -> a
+ Chr
r) Chr
0 ([RL a] -> Chr) -> (RLE a -> [RL a]) -> RLE a -> Chr
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 :: forall a. Eq a => [a] -> RLE a
rle = [RL a] -> RLEof [] a
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE ([RL a] -> RLEof [] a) -> ([a] -> [RL a]) -> [a] -> RLEof [] 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 -> Chr -> RL a
forall a. a -> Chr -> RL a
RL a
x (Chr -> Chr
forall a. Enum a => a -> a
succ (Chr -> Chr) -> Chr -> Chr
forall a b. (a -> b) -> a -> b
$ [a] -> Chr
forall a. [a] -> Chr
forall (t :: * -> *) a. Foldable t => t a -> Chr
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 :: forall a. Eq a => Vector a -> RLEV a
rleV = Vector (RL a) -> RLEof Vector a
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE (Vector (RL a) -> RLEof Vector a)
-> (Vector a -> Vector (RL a)) -> Vector a -> RLEof Vector 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 Stream m a -> Stream m (RL a)
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 a. a -> m 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 -> Chr -> RL a
forall a. a -> Chr -> RL a
RL a
x Chr
1), s
s')
          Just r :: RL a
r@(RL a
y Chr
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 a. a -> m 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 -> Chr -> RL a
forall a. a -> Chr -> RL a
RL a
x (Chr -> RL a) -> Chr -> RL a
forall a b. (a -> b) -> a -> b
$ Chr -> Chr
forall a. Enum a => a -> a
succ Chr
n), s
s')
            | Bool
otherwise -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall a. a -> m 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 -> Chr -> RL a
forall a. a -> Chr -> RL a
RL a
x Chr
1), s
s')
        VS.Skip s
s' -> Step (Maybe (RL a), s) (RL a) -> m (Step (Maybe (RL a), s) (RL a))
forall a. a -> m 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 a. a -> m 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 :: forall a. Ord a => RLE a -> RLE a
sortRLE = ([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] a
forall (f :: * -> *) a (g :: * -> *) b.
(f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE (([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] a)
-> ([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] 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 Chr
r:[RL a]
l) -> a -> Chr -> RL a
forall a. a -> Chr -> RL a
RL a
x (Chr
r Chr -> Chr -> Chr
forall a. Num a => a -> a -> a
+ RLEof [] a -> Chr
forall a. RLE a -> Chr
rleLength ([RL a] -> RLEof [] 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 :: forall a. (a -> Bool) -> RLE a -> RLE a
filterRLE a -> Bool
f = ([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] a
forall (f :: * -> *) a (g :: * -> *) b.
(f (RL a) -> g (RL b)) -> RLEof f a -> RLEof g b
withRLE (([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] a)
-> ([RL a] -> [RL a]) -> RLEof [] a -> RLEof [] 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 :: [Chr] -> ChrStr
chrStr = (Chr -> Chr -> Chr) -> [(Chr, Chr)] -> ChrStr
forall a. (a -> a -> a) -> [(Chr, a)] -> IntMap a
M.fromListWith Chr -> Chr -> Chr
forall a. Num a => a -> a -> a
(+) ([(Chr, Chr)] -> ChrStr)
-> ([Chr] -> [(Chr, Chr)]) -> [Chr] -> ChrStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Chr -> (Chr, Chr)) -> [Chr] -> [(Chr, Chr)]
forall a b. (a -> b) -> [a] -> [b]
map (, Chr
1)

chrStrRLE :: ChrStr -> RLE Chr
chrStrRLE :: ChrStr -> RLE Chr
chrStrRLE = [RL Chr] -> RLE Chr
forall (f :: * -> *) a. f (RL a) -> RLEof f a
RLE ([RL Chr] -> RLE Chr) -> (ChrStr -> [RL Chr]) -> ChrStr -> RLE Chr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Chr, Chr) -> RL Chr) -> [(Chr, Chr)] -> [RL Chr]
forall a b. (a -> b) -> [a] -> [b]
map ((Chr -> Chr -> RL Chr) -> (Chr, Chr) -> RL Chr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Chr -> Chr -> RL Chr
forall a. a -> Chr -> RL a
RL) ([(Chr, Chr)] -> [RL Chr])
-> (ChrStr -> [(Chr, Chr)]) -> ChrStr -> [RL Chr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ChrStr -> [(Chr, Chr)]
forall a. IntMap a -> [(Chr, 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 Chr
c) = ChrSet -> PatChar
PatNot (Chr -> ChrSet
S.singleton Chr
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 Chr
c) ChrStr
t = (Chr -> ChrStr) -> Maybe Chr -> ChrStr
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Chr -> Chr -> ChrStr
forall a. Chr -> a -> IntMap a
M.singleton Chr
c) (Maybe Chr -> ChrStr) -> Maybe Chr -> ChrStr
forall a b. (a -> b) -> a -> b
$ Chr -> ChrStr -> Maybe Chr
forall a. Chr -> IntMap a -> Maybe a
M.lookup Chr
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 Chr
c)
  | Chr -> ChrSet -> Bool
S.member Chr
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 Chr
c)
  | Chr -> ChrSet -> Bool
S.member Chr
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 Chr
c) (PatChr Chr
d)
  | Chr
c Chr -> Chr -> Bool
forall a. Eq a => a -> a -> Bool
== Chr
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