syb based implementation of InlineDoBind

InlineDoBind.hs

{-# LANGUAGE FlexibleContexts, TemplateHaskell, TupleSections #-}
module InlineDoBind where

import Data.Data (Data, gmapM)
import Data.Generics.Aliases (extM)
import Language.Haskell.TH (Exp(DoE,AppE,VarE,InfixE),Stmt(BindS),Pat(VarP))
import Language.Haskell.TH.Syntax (Quasi,qNewName)

import Control.Monad.Trans.Writer.Strict (runWriterT,WriterT)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Writer.Class (MonadWriter, tell, pass, listen)

pre :: m a -> a
pre _ = error "pre must be used only inside mkInlineDoBind"

mkInlineDoBind :: (Data a, Quasi m) => m a -> m a
mkInlineDoBind decsQ =
  do
    decs <- decsQ
    (newDecs, leftovers) <- runWriterT (handle decs)
    if not (null leftovers)
      then fail "must use pre only inside do"
      else pure newDecs
  where
    handle :: (Data a, Quasi m) => a -> WriterT [Stmt] m a
    handle = extM (gmapM handle) handleExp
    handleExp :: Quasi m => Exp -> WriterT [Stmt] m Exp
    handleExp (DoE sts) = DoE <$> fmap concat (mapM handleDoSt sts)
    handleExp (AppE func e) | func == VarE 'pre = do
      e1 <- handle e
      varName <- lift (qNewName "inlineBound")
      tell [BindS (VarP varName) e1]
      pure (VarE varName)
    handleExp (InfixE (Just func) op (Just e)) | func == (VarE 'pre) && op == (VarE '($)) = handleExp (AppE func e)
    handleExp e = gmapM handle e
    handleDoSt st = do
      (st1, binds) <- pass $ fmap (,const []) $ listen (handle st)
      pure (binds ++ [st1])

InlineDoBindDemo.hs

{-# LANGUAGE TemplateHaskell, PostfixOperators #-}
{-# OPTIONS_GHC -ddump-splices #-}
module InlineDoBindDemo where

import InlineDoBind

mkInlineDoBind [d|
  foo1 f g i =
    do
      i (pre f) (pre g)

  foo2 f a b c d =
    do
      f (pre a) b (pre c) d

  foo3 async getURL url1 url2 wait = do
    let (a1, a2) = (pre$ async (getURL url1), pre$ async (getURL url2))
    let (page1, page2) = (pre$ wait a1, pre$ wait a2)
    pure ()

  foo4 f g x = do
    f (g (pre x))

  foo5 f g x = do
    f $ do
        g (pre x)

  foo6 process getAction getArgument getConfig = do
    process (pre $ (pre getAction) (pre getArgument)) (pre getConfig)

  foo7 f x h y z = do
    f (pre x)
    let g = h (pre y)
    i <- (pre z)
    return i

  data R m n = R { m :: m, n :: n }
  
  foo8 x f y = do
    return R { m = pre x, n = f (pre y) }
  |]