Refactor UnicodeSyntax.hs (#317)

* Refactor UnicodeSyntax.hs
Co-authored-by: Jasper Van der Jeugt <jaspervdj@gmail.com>

1 files changed, 10 insertions, 32 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/UnicodeSyntax.hs b/lib/Language/Haskell/Stylish/Step/UnicodeSyntax.hs
index 2f0def6..ff01dee 100644
--- a/lib/Language/Haskell/Stylish/Step/UnicodeSyntax.hs
+++ b/lib/Language/Haskell/Stylish/Step/UnicodeSyntax.hs
@@ -52,37 +52,17 @@ groupPerLine :: [((Int, Int), a)] -> [(Int, [(Int, a)])]
 groupPerLine = M.toList . M.fromListWith (++) .
     map (\((r, c), x) -> (r, [(c, x)]))
 
-
---------------------------------------------------------------------------------
-typeSigs :: Module -> Lines -> [((Int, Int), String)]
-typeSigs module' ls =
-    [ (pos, "::")
+-- | Find symbol positions in the module.  Currently only searches in type
+-- signatures.
+findSymbol :: Module -> Lines -> String -> [((Int, Int), String)]
+findSymbol module' ls sym =
+    [ (pos, sym)
     | TypeSig _ funLoc typeLoc <- everything (rawModuleDecls $ moduleDecls module') :: [Sig GhcPs]
-    , (_, funEnd)              <- infoPoints funLoc
-    , (typeStart, _)           <- infoPoints [hsSigWcType typeLoc]
-    , pos                      <- maybeToList $ between funEnd typeStart "::" ls
-    ]
-
---------------------------------------------------------------------------------
-contexts :: Module -> Lines -> [((Int, Int), String)]
-contexts module' ls =
-    [ (pos, "=>")
-    | TypeSig _ _ typeLoc <- everything (rawModuleDecls $ moduleDecls module') :: [Sig GhcPs]
-    , (start, end)        <- infoPoints [hsSigWcType typeLoc]
-    , pos                 <- maybeToList $ between start end "=>" ls
+    , (funStart, _)            <- infoPoints funLoc
+    , (_, typeEnd)             <- infoPoints [hsSigWcType typeLoc]
+    , pos                      <- maybeToList $ between funStart typeEnd sym ls
     ]
 
-
---------------------------------------------------------------------------------
-typeFuns :: Module -> Lines -> [((Int, Int), String)]
-typeFuns module' ls =
-    [ (pos, "->")
-    | TypeSig _ _ typeLoc <- everything (rawModuleDecls $ moduleDecls module') :: [Sig GhcPs]
-    , (start, end)        <- infoPoints [hsSigWcType typeLoc]
-    , pos                 <- maybeToList $ between start end "->" ls
-    ]
-
-
 --------------------------------------------------------------------------------
 -- | Search for a needle in a haystack of lines. Only part the inside (startRow,
 -- startCol), (endRow, endCol) is searched. The return value is the position of
@@ -113,7 +93,5 @@ step' alp lg ls module' = applyChanges changes ls
   where
     changes = (if alp then addLanguagePragma lg "UnicodeSyntax" module' else []) ++
         replaceAll perLine
-    perLine = sort $ groupPerLine $
-        typeSigs module' ls ++
-        contexts module' ls ++
-        typeFuns module' ls
+    toReplace = [ "::", "=>", "->" ]
+    perLine = sort $ groupPerLine $ concatMap (findSymbol module' ls) toReplace