{-# 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