module ContextFilter (filterContext) where

filterContext ::
  Int         {- ^ context before       -} ->
  Int         {- ^ context after        -} ->
  (a -> Bool) {- ^ predicate            -} ->
  [a]         {- ^ inputs               -} ->
  [a]         {- ^ matches with context -}
filterContext :: Int -> Int -> (a -> Bool) -> [a] -> [a]
filterContext Int
before Int
after a -> Bool
p [a]
xs
  | Int
before Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"filterContext: bad before"
  | Int
after  Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"filterContext: bad after"
  | Bool
otherwise  = Selection -> [a] -> [a]
forall a. Selection -> [a] -> [a]
selectList (Int -> Selection -> Selection
dropSelection Int
before Selection
selects) [a]
xs
  where
    width :: Int
width = Int
before Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
after

    selects :: Selection
selects = Int -> [a] -> Selection
go Int
0 [a]
xs

    go :: Int -> [a] -> Selection
go Int
n []       = Int -> Selection
replicateKeep Int
n
    go Int
n (a
y : [a]
ys)
      | a -> Bool
p a
y       = Selection -> Selection
Keep (Int -> [a] -> Selection
go Int
width [a]
ys)
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = Selection -> Selection
Keep (Int -> [a] -> Selection
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
ys)
      | Bool
otherwise = Selection -> Selection
Skip (Int -> [a] -> Selection
go Int
n [a]
ys)

data Selection = End | Keep Selection | Skip Selection
  deriving (Int -> Selection -> ShowS
[Selection] -> ShowS
Selection -> [Char]
(Int -> Selection -> ShowS)
-> (Selection -> [Char])
-> ([Selection] -> ShowS)
-> Show Selection
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Selection] -> ShowS
$cshowList :: [Selection] -> ShowS
show :: Selection -> [Char]
$cshow :: Selection -> [Char]
showsPrec :: Int -> Selection -> ShowS
$cshowsPrec :: Int -> Selection -> ShowS
Show)

replicateKeep :: Int -> Selection
replicateKeep :: Int -> Selection
replicateKeep Int
0 = Selection
End
replicateKeep Int
i = Selection -> Selection
Keep (Int -> Selection
replicateKeep (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))

dropSelection :: Int -> Selection -> Selection
dropSelection :: Int -> Selection -> Selection
dropSelection Int
0 Selection
x        = Selection
x
dropSelection Int
_ Selection
End      = Selection
End
dropSelection Int
i (Keep Selection
x) = Int -> Selection -> Selection
dropSelection (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Selection
x
dropSelection Int
i (Skip Selection
x) = Int -> Selection -> Selection
dropSelection (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Selection
x

selectList :: Selection -> [a] -> [a]
selectList :: Selection -> [a] -> [a]
selectList (Keep Selection
x) (a
y : [a]
ys) = a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Selection -> [a] -> [a]
forall a. Selection -> [a] -> [a]
selectList Selection
x [a]
ys
selectList (Skip Selection
x) (a
_ : [a]
ys) =     Selection -> [a] -> [a]
forall a. Selection -> [a] -> [a]
selectList Selection
x [a]
ys
selectList Selection
_        [a]
_        = []