2 files changed, 69 insertions, 45 deletions
diff --git a/app/Main.hs b/app/Main.hs
index abc8d4c..126caab 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -64,6 +64,8 @@ import System.Exit (ExitCode (ExitFailure), exitWith)
 import System.FilePath as F
 import System.IO (hPutStrLn, stderr)
 import System.Process.Typed qualified as P
+import TreeGrepper.Match (Match (..), Position (..))
+import TreeGrepper.Match qualified as TM
 
 data Options = Options
   { optCommand :: Command
@@ -155,25 +157,6 @@ data TreeGrepperResult = TreeGrepperResult
 
 instance A.FromJSON TreeGrepperResult
 
-data Match = Match
-  { kind :: String,
-    name :: String,
-    text :: String,
-    start :: Position,
-    end :: Position
-  }
-  deriving (Show, Generic)
-
-instance A.FromJSON Match
-
-data Position = Position
-  { row :: Int,
-    column :: Int
-  }
-  deriving (Show, Generic)
-
-instance A.FromJSON Position
-
 getIssues :: String -> IO [TreeGrepperResult]
 getIssues filename =
   let extension = F.takeExtension filename
@@ -217,32 +200,7 @@ getIssues filename =
 
 fixTreeGrepper :: TreeGrepperResult -> TreeGrepperResult
 fixTreeGrepper treeGrepperResult =
-  treeGrepperResult {matches = mergeMatches treeGrepperResult.matches}
-  where
-    mergeMatches matches =
-      Maybe.catMaybes
-        [ subs ms
-          | ms <- groupBy eq matches
-        ]
-    groupBy p xs = reverse (map reverse (groupBy' [] p xs))
-    groupBy' as p [] = as
-    groupBy' [] p (x : xs) = groupBy' [[x]] p xs
-    groupBy' (ass@((a : as) : rs)) p (x : xs)
-      | p a x = groupBy' ((x : a : as) : rs) p xs
-      | otherwise = groupBy' ([x] : ass) p xs
-    eq m n = m.end.row + 1 == n.start.row
-    subs [] = Nothing
-    subs (mss@(m : _)) =
-      Just
-        ( m
-            { start = start mss,
-              end = end mss,
-              text = unlines (map (.text) mss)
-            }
-        )
-    start ms = minimumBy (O.comparing loc) (map (.start) ms)
-    end ms = maximumBy (O.comparing loc) (map (.end) ms)
-    loc x = (x.row, x.column)
+  treeGrepperResult {matches = TM.merge treeGrepperResult.matches}
 
 getFiles :: IO [String]
 getFiles =
diff --git a/app/TreeGrepper/Match.hs b/app/TreeGrepper/Match.hs
new file mode 100644
index 0000000..f882a3c
--- /dev/null
+++ b/app/TreeGrepper/Match.hs
@@ -0,0 +1,66 @@
+{-# LANGUAGE OverloadedRecordDot #-}
+
+module TreeGrepper.Match
+  ( Match (..),
+    Position (..),
+    merge,
+  )
+where
+
+import Data.Aeson (FromJSON)
+import Data.Function (on)
+import Data.List (sortBy)
+import Data.Ord (comparing)
+import GHC.Generics (Generic)
+
+data Match = Match
+  { kind :: String,
+    name :: String,
+    text :: String,
+    start :: Position,
+    end :: Position
+  }
+  deriving (Show, Generic)
+
+instance FromJSON Match
+
+data Position = Position
+  { row :: Int,
+    column :: Int
+  }
+  deriving (Eq, Show, Generic)
+
+instance Ord Position where
+  compare = compare `on` (\p -> (p.row, p.column))
+
+instance FromJSON Position
+
+{- tree-grepper (tree-sitter) is unable to match sibling comments blocks as a whole. We thus merge matches if they are line-adjacent. -}
+merge :: [Match] -> [Match]
+merge matches =
+  map mergeGroup
+    . groupBy (\a b -> a.end.row + 1 == b.start.row)
+    $ sortBy (comparing (.start)) matches
+
+mergeGroup :: [Match] -> Match
+mergeGroup (m : ms) =
+  m
+    { text = text,
+      start = start,
+      end = end
+    }
+  where
+    mss = m : ms
+    text = unlines $ map (.text) mss
+    start = minimum $ map (.start) mss
+    end = maximum $ map (.end) mss
+
+{- A version of `Data.List.groupBy` that uses the last added group-member for comparison with new candidates for the group. `Data.List.groupBy` uses the initial member for all subsequent comparisons. -}
+groupBy :: (a -> a -> Bool) -> [a] -> [[a]]
+groupBy p xs = reverse . map reverse $ go [] p xs
+  where
+    go rs _ [] = rs
+    go [] p (x : xs) = go [[x]] p xs
+    go (ass@((a : as) : rs)) p (x : xs)
+      | p a x = go ((x : a : as) : rs) p xs
+      | otherwise = go ([x] : ass) p xs