editing issues preserves comment style

1 files changed, 54 insertions, 0 deletions

diff --git a/app/TreeGrepper/Comment.hs b/app/TreeGrepper/Comment.hs
index 1a6aed2..0ca9543 100644
--- a/app/TreeGrepper/Comment.hs
+++ b/app/TreeGrepper/Comment.hs
@@ -1,13 +1,19 @@
 module TreeGrepper.Comment
   ( Comment (..),
     getComments,
+    CommentStyle (..),
+    uncomment,
+    comment,
   )
 where
 
 import Control.Exception (throw)
 import Data.Aeson qualified as A
+import Data.Binary (Binary)
 import Data.ByteString.Lazy.Char8 qualified as B
 import Data.Function ((&))
+import Data.List (find)
+import Data.Maybe (fromMaybe)
 import Data.Text qualified as T
 import Exception qualified as E
 import GHC.Generics (Generic)
@@ -15,6 +21,7 @@ import Process (proc, sh)
 import System.FilePath (takeExtension)
 import System.Process.Typed (setWorkingDir)
 import TreeGrepper.FileType (FileType (..))
+import TreeGrepper.FileType qualified as G
 import TreeGrepper.Match (Match (..), Position (..))
 import TreeGrepper.Match qualified as G
 import TreeGrepper.Result (Result (..))
@@ -33,6 +40,53 @@ data Comment = Comment
   }
   deriving (Show, Generic)
 
+data CommentStyle
+  = LineStyle T.Text
+  | BlockStyle T.Text T.Text
+  deriving (Eq, Show, Generic, Binary)
+
+comment :: CommentStyle -> T.Text -> T.Text
+comment (LineStyle linePrefix) = T.unlines . map ((linePrefix <> " ") <>) . T.lines
+comment (BlockStyle blockStart blockEnd) = (blockStart <>) . (<> blockEnd)
+
+uncomment :: FileType -> T.Text -> (CommentStyle, T.Text)
+uncomment fileType rawText =
+  maybe
+    ( ( LineStyle info.lineStart,
+        stripLineComments (G.info fileType).lineStart text
+      )
+    )
+    ( \(blockInfo, blockStart) ->
+        ( BlockStyle blockStart blockInfo.blockEnd,
+          stripBlockComment blockStart blockInfo.blockEnd text
+        )
+    )
+    $ do
+      blockInfo <- info.block
+      (,) blockInfo <$> find (`T.isPrefixOf` text) blockInfo.blockStart
+  where
+    info = G.info fileType
+    text = stripLines rawText
+    stripLines = T.intercalate "\n" . map T.strip . T.lines
+
+stripLineComments :: T.Text -> T.Text -> T.Text
+stripLineComments lineStart text =
+  onLines
+    ( \line ->
+        fromMaybe line . fmap T.stripStart $
+          T.stripPrefix lineStart line
+    )
+    text
+  where
+    onLines f = T.intercalate "\n" . map f . T.lines
+
+stripBlockComment :: T.Text -> T.Text -> T.Text -> T.Text
+stripBlockComment blockStart blockEnd text =
+  T.strip
+    . (fromMaybe text . T.stripSuffix blockEnd)
+    . (fromMaybe text . T.stripPrefix blockStart)
+    $ text
+
 fromMatch :: Result -> Match -> Comment
 fromMatch Result {..} Match {..} = Comment {..}