{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE TemplateHaskell #-}

module Test.Sandwich.TH.HasMainFunction (
  fileHasMainFunction
  , ShouldWarnOnParseError(..)
  ) where

import Control.Monad
import Data.String.Interpolate
import Language.Haskell.Exts hiding (name)
import Language.Haskell.TH (Q, runIO, reportWarning)

-- import Debug.Trace


data ShouldWarnOnParseError = WarnOnParseError | NoWarnOnParseError
  deriving (ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
(ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool)
-> (ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool)
-> Eq ShouldWarnOnParseError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
== :: ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
$c/= :: ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
/= :: ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
Eq)

-- | Use haskell-src-exts to determine if a give Haskell file has an exported main function
-- Parse with all extensions enabled, which will hopefully parse anything
fileHasMainFunction :: FilePath -> ShouldWarnOnParseError -> Q Bool
fileHasMainFunction :: String -> ShouldWarnOnParseError -> Q Bool
fileHasMainFunction String
path ShouldWarnOnParseError
shouldWarnOnParseError = IO (ParseResult (Module SrcSpanInfo))
-> Q (ParseResult (Module SrcSpanInfo))
forall a. IO a -> Q a
runIO ([Extension] -> String -> IO (ParseResult (Module SrcSpanInfo))
parseFileWithExts [Extension
x | x :: Extension
x@(EnableExtension KnownExtension
_) <- [Extension]
knownExtensions] String
path) Q (ParseResult (Module SrcSpanInfo))
-> (ParseResult (Module SrcSpanInfo) -> Q Bool) -> Q Bool
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  x :: ParseResult (Module SrcSpanInfo)
x@(ParseFailed {}) -> do
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ShouldWarnOnParseError
shouldWarnOnParseError ShouldWarnOnParseError -> ShouldWarnOnParseError -> Bool
forall a. Eq a => a -> a -> Bool
== ShouldWarnOnParseError
WarnOnParseError) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
      String -> Q ()
reportWarning [i|Failed to parse #{path}: #{x}|]
    Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  ParseOk (Module SrcSpanInfo
_ (Just ModuleHead SrcSpanInfo
moduleHead) [ModulePragma SrcSpanInfo]
_ [ImportDecl SrcSpanInfo]
_ [Decl SrcSpanInfo]
decls) -> do
    -- traceM [i|Sucessfully parsed #{path}: #{moduleHead}|]
    case ModuleHead SrcSpanInfo
moduleHead of
      ModuleHead SrcSpanInfo
_ ModuleName SrcSpanInfo
_ Maybe (WarningText SrcSpanInfo)
_ (Just (ExportSpecList SrcSpanInfo
_ ((ExportSpec SrcSpanInfo -> Bool)
-> [ExportSpec SrcSpanInfo] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ExportSpec SrcSpanInfo -> Bool
forall l. ExportSpec l -> Bool
isMainFunction -> Bool
True))) -> do
        -- traceM [i|FOUND MAIN FUNCTION FOR #{path}|]
        Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
      ModuleHead SrcSpanInfo
_ ModuleName SrcSpanInfo
_ Maybe (WarningText SrcSpanInfo)
_ Maybe (ExportSpecList SrcSpanInfo)
Nothing -> do
        -- traceM [i|LOOKING FOR DECLS: #{decls}|]
        Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> Q Bool) -> Bool -> Q Bool
forall a b. (a -> b) -> a -> b
$ (Decl SrcSpanInfo -> Bool) -> [Decl SrcSpanInfo] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Decl SrcSpanInfo -> Bool
forall l. Decl l -> Bool
isMainDecl [Decl SrcSpanInfo]
decls
      ModuleHead SrcSpanInfo
_ -> Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  ParseOk Module SrcSpanInfo
_ -> do
    String -> Q ()
reportWarning [i|Successfully parsed #{path} but no module head found|]
    Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

isMainFunction :: ExportSpec l -> Bool
isMainFunction :: forall l. ExportSpec l -> Bool
isMainFunction (EVar l
_ QName l
name) = QName l -> Bool
forall l. QName l -> Bool
isMainFunctionQName QName l
name
isMainFunction ExportSpec l
_ = Bool
False

isMainFunctionQName :: QName l -> Bool
isMainFunctionQName :: forall l. QName l -> Bool
isMainFunctionQName (Qual l
_ ModuleName l
_ Name l
name) = Name l -> Bool
forall l. Name l -> Bool
isMainFunctionName Name l
name
isMainFunctionQName (UnQual l
_ Name l
name) = Name l -> Bool
forall l. Name l -> Bool
isMainFunctionName Name l
name
isMainFunctionQName QName l
_ = Bool
False

isMainFunctionName :: Name l -> Bool
isMainFunctionName :: forall l. Name l -> Bool
isMainFunctionName (Ident l
_ String
"main") = Bool
True
isMainFunctionName (Symbol l
_ String
"main") = Bool
True
isMainFunctionName Name l
_ = Bool
False

isMainDecl :: Decl l -> Bool
isMainDecl :: forall l. Decl l -> Bool
isMainDecl (PatBind l
_ (PVar l
_ (Ident l
_ String
"main")) Rhs l
_ Maybe (Binds l)
_) = Bool
True
-- isMainDecl decl = trace [i|Looking at decl: #{decl}|] False
isMainDecl Decl l
_ = Bool
False