early exit

lib - Plugin.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}

module Plugin where

import Data.Data (Data)
import Data.Generics.Uniplate.Data qualified as Uniplate
import GHC
import GhcPlugins
import Plugins

deriving stock instance Data HsParsedModule

plugin = defaultPlugin {parsedResultAction = \_ _ -> pure . f}
  where
    f :: HsParsedModule -> HsParsedModule
    f = Uniplate.transformBi g

    g :: Located [ExprLStmt GhcPs] -> Located [ExprLStmt GhcPs]
    g (L loc stms) = L loc (go stms)

    go :: [ExprLStmt GhcPs] -> [ExprLStmt GhcPs]
    go [] = []
    go (x@(EarlyStm pat body fail) : xs) =
      let bind = noLoc (HsVar NoExtField (noLoc (mkRdrUnqual (mkVarOcc ">>="))))
          lcase = HsLamCase NoExtField (MG NoExtField (noLoc (map noLoc cases)) Generated)
          onSuccess = GRHS NoExtField [] (noLoc (HsDo NoExtField DoExpr (noLoc xs)))
          onFail = GRHS NoExtField [] (fail)
          cases =
            [ Match NoExtField CaseAlt [pat] (GRHSs NoExtField [noLoc onSuccess] (noLoc emptyLocalBinds)),
              Match NoExtField CaseAlt [noLoc (WildPat NoExtField)] (GRHSs NoExtField [noLoc onFail] (noLoc emptyLocalBinds))
            ]
          act = OpApp NoExtField (noLoc (HsPar NoExtField body)) bind (noLoc lcase)
       in [noLoc (BodyStmt NoExtField (noLoc act) noSyntaxExpr noSyntaxExpr)]
    go (x : xs) = x : go xs

out :: Outputable a => a -> String
out = showSDocUnsafe . ppr

pattern Or :: LHsExpr GhcPs
pattern Or <- L _ (HsVar _ (L _ (rdrNameString -> "||")))

rdrNameString :: RdrName -> String
rdrNameString = occNameString . rdrNameOcc

pattern EarlyStm :: LPat GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs -> LStmt GhcPs (LHsExpr GhcPs)
pattern EarlyStm pat body fail <- L _ (BindStmt _ pat (L _ (OpApp _ body Or fail)) _ _)

package.yaml

name: hspg

dependencies:
  - base
  - uniplate
  - ghc

executables:
  hspg-exe:
    main: Main.hs
    source-dirs: src
    dependencies:
      - hspg

library:
  source-dirs: lib

src - Main.hs

{-# OPTIONS_GHC -fplugin=Plugin #-}

prompt :: String -> IO String
prompt msg = do
  putStr msg
  getLine

main = do
  "hello" <- prompt $ "put 'hello'> " || putStrLn "fail"
  putStrLn "ok"
  pure ()