Apply do-syntax.

3 files changed, 24 insertions, 2 deletions

diff --git a/Sweeten.hs b/Sweeten.hs
new file mode 100644
index 0000000..929f9f8
--- /dev/null
+++ b/Sweeten.hs
@@ -0,0 +1,21 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+module Sweeten where
+
+import Data.Data
+import Data.Generics.Aliases
+import Data.Generics.Schemes
+import Language.Haskell.Exts.Syntax as HS
+
+applyDoSyntax :: forall l. (Typeable l, Data l) => Exp l -> Exp l
+applyDoSyntax = everywhere (mkT (applyDoSyntax1 :: Exp l -> Exp l))
+
+applyDoSyntax1 :: Exp l -> Exp l
+applyDoSyntax1 exp@(InfixApp la x (QVarOp lb (UnQual lc (Symbol ld ">>="))) (Lambda le [pat] y)) =
+    case y of
+        Do lf ss -> Do la (Generator lb pat x : ss)
+        _        -> Do la [Generator lb pat x, Qualifier le y]
+applyDoSyntax1 exp@(InfixApp la x (QVarOp lb (UnQual lc (Symbol ld ">>"))) y) =
+    case y of
+        Do lf ss -> Do la (Qualifier lb x : ss)
+        _        -> Do la [Qualifier lb x, Qualifier lc y]
+applyDoSyntax1 exp = exp
diff --git a/monkeypatch.cabal b/monkeypatch.cabal
index 23bc405..2b06032 100644
--- a/monkeypatch.cabal
+++ b/monkeypatch.cabal
@@ -52,7 +52,7 @@ extra-source-files:  CHANGELOG.md
 
 executable monkeypatch
   main-is:             monkeypatch.hs
-  other-modules:       GrepNested
+  other-modules:       GrepNested, Sweeten
   -- other-extensions:
   build-depends:       base >=4.10.1.0  && <=4.12
                      , containers ^>=0.5.10.2
diff --git a/monkeypatch.hs b/monkeypatch.hs
index b72e346..37e52a6 100644
--- a/monkeypatch.hs
+++ b/monkeypatch.hs
@@ -49,6 +49,7 @@ import Text.PrettyPrint             (Doc, doubleQuotes, empty, text, vcat, ($$),
                                      (<+>))
 import Text.Show.Pretty
 
+import Sweeten
 import GrepNested
 
 trace :: p -> a -> a
@@ -631,7 +632,7 @@ transpile o fname incs (CTranslUnit edecls _) = do
                                 return $ foldr applyComputation (Computation Map.empty Map.empty retUnit) xs
                         case mhask of
                             Just hask -> do printHeader
-                                            mapM_ (putStrLn . ("  "++)) $ lines $ HS.prettyPrint $ comp hask
+                                            mapM_ (putStrLn . ("  "++)) $ lines $ HS.prettyPrint $ applyDoSyntax $ comp hask
                             Nothing -> forM_ (oSelectFunction o) $ \_ -> do
                                 printHeader
                                 forM_ bdy $ \d -> do