1 files changed, 16 insertions, 12 deletions
diff --git a/app/TreeGrepper/Match.hs b/app/TreeGrepper/Match.hs
index f882a3c..7b8cde8 100644
--- a/app/TreeGrepper/Match.hs
+++ b/app/TreeGrepper/Match.hs
@@ -10,13 +10,17 @@ where
 import Data.Aeson (FromJSON)
 import Data.Function (on)
 import Data.List (sortBy)
+import Data.List.NonEmpty (NonEmpty ((:|)))
+import Data.List.NonEmpty qualified as N
 import Data.Ord (comparing)
+import Data.Text (Text)
+import Data.Text qualified as T
 import GHC.Generics (Generic)
 
 data Match = Match
   { kind :: String,
     name :: String,
-    text :: String,
+    text :: Text,
     start :: Position,
     end :: Position
   }
@@ -39,11 +43,11 @@ instance FromJSON Position
 merge :: [Match] -> [Match]
 merge matches =
   map mergeGroup
-    . groupBy (\a b -> a.end.row + 1 == b.start.row)
+    . chainsBy (\a b -> a.end.row + 1 == b.start.row)
     $ sortBy (comparing (.start)) matches
 
-mergeGroup :: [Match] -> Match
-mergeGroup (m : ms) =
+mergeGroup :: NonEmpty Match -> Match
+mergeGroup (m :| ms) =
   m
     { text = text,
       start = start,
@@ -51,16 +55,16 @@ mergeGroup (m : ms) =
     }
   where
     mss = m : ms
-    text = unlines $ map (.text) mss
+    text = T.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
+chainsBy :: (a -> a -> Bool) -> [a] -> [NonEmpty a]
+chainsBy p = reverse . map N.reverse . go []
   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
+    go rs [] = rs
+    go [] (x : xs) = go [N.singleton x] xs
+    go (ass@((a :| as) : rs)) (x : xs)
+      | p a x = go ((x :| a : as) : rs) xs
+      | otherwise = go (N.singleton x : ass) xs