Underscore.hs

module Underscore (doUnderscore) where

import Language.Haskell.TH
import Control.Monad (unless)
import Data.Generics (everything, everywhereM, mkM, mkQ)
import Data.Monoid (Any(Any))

-- | The placeholder for arguments is the wildcard @_@
placeholder :: Exp
placeholder = UnboundVarE (mkName "_")

-- | Predicate for checking if the 'placeholder' is contained in an 'Exp'
hasPlaceholder :: Exp -> Bool
hasPlaceholder = everything (||) (False `mkQ` (==) placeholder)

doUnderscore :: ExpQ -> ExpQ
doUnderscore expQ =
  do exprs <- traverse onlyNoBindS =<< onlyDoE =<< expQ
     startTransformation (reverse exprs)

-- | Ensure that only no-bind statements are used in the do-notation
onlyNoBindS :: Stmt -> ExpQ
onlyNoBindS (NoBindS s) = pure s
onlyNoBindS _           = fail nonNoBind

-- | Ensure that the top-most expression is do-notation
onlyDoE :: Exp -> Q [Stmt]
onlyDoE (DoE stmts) = pure stmts
onlyDoE _           = fail noTopDo

startTransformation ::
  [Exp] {- ^ reversed list of expressions from top-level do-notation -} ->
  ExpQ  {- ^ fina let-expression                                     -}
startTransformation [] = fail emptyDoNotation
startTransformation (x:xs)
  | hasPlaceholder x = fail placeholderInFinalExpr
  | otherwise        = stepTransformation (length xs) [] x xs

stepTransformation ::
  Int   {- ^ counter for generating unique names -} ->
  [Dec] {- ^ accumulated let bindings            -} ->
  Exp   {- ^ current result                      -} ->
  [Exp] {- ^ remaining expressions to transform  -} ->
  ExpQ  {- ^ final let-expression                -}
stepTransformation _ []   cur []       = pure cur -- skip let when not needed
stepTransformation _ decs cur []       = pure (LetE decs cur)
stepTransformation i decs cur (fun:xs) =
  do var <- newName ("x" ++ show i)

     let decs' = ValD (VarP var) (NormalB cur) [] : decs

         replaceUnderscore e
           | e == placeholder = (Any True, VarE var)
           | otherwise        = pure e

         (Any found, fun') = everywhereM (mkM replaceUnderscore) fun

     unless found (reportWarning (missingPlaceholder i))

     stepTransformation (i-1) decs' fun' xs

------------------------------------------------------------------------
-- Warning and error messages
------------------------------------------------------------------------

wrapMsg :: String -> String
wrapMsg msg = "doUnderscore: " ++ msg

missingPlaceholder :: Int -> String
missingPlaceholder i = wrapMsg ("Placeholder not used in " ++ ordinal i ++ " expression")

placeholderInFinalExpr, emptyDoNotation, noTopDo, nonNoBind :: String
placeholderInFinalExpr = wrapMsg "placeholder used in final expression"
emptyDoNotation        = wrapMsg "non-empty do-notation expected"
noTopDo                = wrapMsg "top-level do-notation expected"
nonNoBind              = wrapMsg "only no-bind statements expected"

-- | Pretty rendering of ordinals for error messages.
ordinal :: Int -> String
ordinal 1 = "1st"
ordinal 2 = "2nd"
ordinal 3 = "3rd"
ordinal n = show n ++ "th"