{-|
Module      : THSH.QQ
Description : Template Haskell qausi-quote machinary, heavily using code from PyF package.
Copyright   : (c) Miao ZhiCheng, 2024
License     : MIT
Maintainer  : zhicheng.miao@gmail.com
Stability   : experimental
Portability : POSIX
-}


{-# LANGUAGE TemplateHaskellQuotes #-}

module THSH.QQ
  ( thsh
  ) where

-- ghc modules
import           GHC                        (SrcLoc, mkSrcLoc, mkSrcSpan)
import qualified Language.Haskell.TH        as TH
import           Language.Haskell.TH.Quote  (QuasiQuoter (..))
import           Language.Haskell.TH.Syntax (Code, Exp (..), Q (..))
-- base module
import           Data.List                  (intercalate)
import           Data.Maybe                 (catMaybes, listToMaybe)
import           Data.String                (fromString)
-- transformers module
import           Control.Monad.Trans.Reader (runReader)
-- parsec module
import qualified Text.Parsec                as Ps
import qualified Text.Parsec.Error          as PsError
import qualified Text.Parsec.Pos            as PsPos
--
import           THSH.Funclet               (AnyFunclet (..))
import qualified THSH.Internal.PyFInternals as PyF
import           THSH.Internal.THUtils      (reportErrorAt)
import           THSH.Script                (Script (..), genFuncletPipeCode)


-- | The quasi quoter for Template Haskell shell scripts.
thsh :: QuasiQuoter
thsh :: QuasiQuoter
thsh = QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp = String -> Q Exp
toExp,
      quotePat :: String -> Q Pat
quotePat = String -> String -> Q Pat
forall a. String -> a
mkErr String
"pattern",
      quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall a. String -> a
mkErr String
"type",
      quoteDec :: String -> Q [Dec]
quoteDec = String -> String -> Q [Dec]
forall a. String -> a
mkErr String
"declaration"
    }
  where
    mkErr :: String -> a
    mkErr :: forall a. String -> a
mkErr String
name = String -> a
forall a. HasCallStack => String -> a
error (String
"thsh : This QuasiQuoter can not be used as a " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"!")

    toExp :: String -> Q Exp
    toExp :: String -> Q Exp
toExp String
s = do
      loc <- Q Loc
TH.location
      exts <- TH.extsEnabled
      let context = Maybe (Char, Char) -> [Extension] -> ParsingContext
PyF.ParsingContext ((Char, Char) -> Maybe (Char, Char)
forall a. a -> Maybe a
Just (Char
'«', Char
'»')) [Extension]
exts

      -- Setup the parser so it matchs the real original position in the source
      -- code.
      let filename = Loc -> String
TH.loc_filename Loc
loc
      let initPos = SourcePos -> Column -> SourcePos
Ps.setSourceColumn (SourcePos -> Column -> SourcePos
Ps.setSourceLine (String -> SourcePos
PsPos.initialPos String
filename) ((Column, Column) -> Column
forall a b. (a, b) -> a
fst ((Column, Column) -> Column) -> (Column, Column) -> Column
forall a b. (a -> b) -> a -> b
$ Loc -> (Column, Column)
TH.loc_start Loc
loc)) ((Column, Column) -> Column
forall a b. (a, b) -> b
snd ((Column, Column) -> Column) -> (Column, Column) -> Column
forall a b. (a -> b) -> a -> b
$ Loc -> (Column, Column)
TH.loc_start Loc
loc)
      case runReader (Ps.runParserT (Ps.setPosition initPos >> PyF.parseGenericFormatString) () filename s) context of
        -- returns a dummy exp, so TH continues its life. This TH code won't be
        -- executed anyway, there is an error
        Left ParseError
err -> ParseError -> Q ()
reportParserErrorAt ParseError
err Q () -> Q Exp -> Q Exp
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [|()|]
        Right [Item]
items -> do
          -- stop at the first item that contains an error
          (Q (Maybe (SrcSpan, String)) -> Q (Maybe (SrcSpan, String)))
-> [Q (Maybe (SrcSpan, String))] -> Q [Maybe (SrcSpan, String)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Q (Maybe (SrcSpan, String)) -> Q (Maybe (SrcSpan, String))
forall a. a -> a
id ((Item -> Q (Maybe (SrcSpan, String)))
-> [Item] -> [Q (Maybe (SrcSpan, String))]
forall a b. (a -> b) -> [a] -> [b]
map Item -> Q (Maybe (SrcSpan, String))
PyF.checkOneItem [Item]
items) Q [Maybe (SrcSpan, String)]
-> ([Maybe (SrcSpan, String)] -> Q (Maybe (SrcSpan, String)))
-> Q (Maybe (SrcSpan, String))
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe (SrcSpan, String) -> Q (Maybe (SrcSpan, String))
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (SrcSpan, String) -> Q (Maybe (SrcSpan, String)))
-> ([Maybe (SrcSpan, String)] -> Maybe (SrcSpan, String))
-> [Maybe (SrcSpan, String)]
-> Q (Maybe (SrcSpan, String))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(SrcSpan, String)] -> Maybe (SrcSpan, String)
forall a. [a] -> Maybe a
listToMaybe ([(SrcSpan, String)] -> Maybe (SrcSpan, String))
-> ([Maybe (SrcSpan, String)] -> [(SrcSpan, String)])
-> [Maybe (SrcSpan, String)]
-> Maybe (SrcSpan, String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (SrcSpan, String)] -> [(SrcSpan, String)]
forall a. [Maybe a] -> [a]
catMaybes Q (Maybe (SrcSpan, String))
-> (Maybe (SrcSpan, String) -> Q Exp) -> Q Exp
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Maybe (SrcSpan, String)
Nothing -> Code Q Script -> Q Exp
forall a (m :: * -> *). Quote m => Code m a -> m Exp
TH.unTypeCode ([Item] -> Code Q Script
mkScript [Item]
items)
            Just (SrcSpan
srcSpan, String
msg) -> SrcSpan -> String -> Q ()
reportErrorAt SrcSpan
srcSpan String
msg Q () -> Q Exp -> Q Exp
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [|()|]

{- ========== mkScript ========== -}

mkScript :: [PyF.Item] -> Code Q Script
mkScript :: [Item] -> Code Q Script
mkScript [Item]
items = [|| String -> [AnyFunclet] -> Script
MkScript $$(Q Exp -> Code Q String
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce Q Exp
source) $$(Q Exp -> Code Q [AnyFunclet]
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce Q Exp
funclets) ||]
  where items' :: [(Bool, Q Exp)]
items'   = (Item -> (Bool, Q Exp)) -> [Item] -> [(Bool, Q Exp)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Item -> (Bool, Q Exp)
matchItem [Item]
items
        funclets :: Q Exp
funclets = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
appendQ (Exp -> [Exp] -> Exp) -> Q Exp -> Q ([Exp] -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [| [] |] Q ([Exp] -> Exp) -> Q [Exp] -> Q Exp
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                   ((Bool, Q Exp) -> Q Exp) -> [(Bool, Q Exp)] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Bool, Q Exp) -> Q Exp
forall a b. (a, b) -> b
snd (((Bool, Q Exp) -> Bool) -> [(Bool, Q Exp)] -> [(Bool, Q Exp)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
True) (Bool -> Bool) -> ((Bool, Q Exp) -> Bool) -> (Bool, Q Exp) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool, Q Exp) -> Bool
forall a b. (a, b) -> a
fst) [(Bool, Q Exp)]
items')
        source :: Q Exp
source   = (Column, Q Exp) -> Q Exp
forall a b. (a, b) -> b
snd ((Column, Q Exp) -> Q Exp) -> (Column, Q Exp) -> Q Exp
forall a b. (a -> b) -> a -> b
$ ((Column, Q Exp) -> (Bool, Q Exp) -> (Column, Q Exp))
-> (Column, Q Exp) -> [(Bool, Q Exp)] -> (Column, Q Exp)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(Column
c, Q Exp
rs) (Bool
isFunclet, Q Exp
frag) -> if Bool
isFunclet
                                 then (Column
c Column -> Column -> Column
forall a. Num a => a -> a -> a
+ Column
1, Exp -> Exp -> Exp
appendQ (Exp -> Exp -> Exp) -> Q Exp -> Q (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Exp
rs Q (Exp -> Exp) -> Q Exp -> Q Exp
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [| genFuncletPipeCode c |])
                                 else (Column
c, Exp -> Exp -> Exp
appendQ (Exp -> Exp -> Exp) -> Q Exp -> Q (Exp -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Exp
rs Q (Exp -> Exp) -> Q Exp -> Q Exp
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Exp
frag)
                               ) (Column
0 :: Int, [| [] |]) [(Bool, Q Exp)]
items'

-- | call `<>` between two 'Exp'
appendQ :: Exp -> Exp -> Exp
appendQ :: Exp -> Exp -> Exp
appendQ Exp
s0 Exp
s1 = Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
s0) (Name -> Exp
VarE '(<>)) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
s1)

-- | Match an item to a funclet expression (True) or a formatted String expression (False)
matchItem :: PyF.Item -> (Bool, Q Exp)
matchItem :: Item -> (Bool, Q Exp)
matchItem (PyF.Raw String
x)                   = (Bool
False, [| x |])
matchItem (PyF.Replacement (HsExpr GhcPs
_, Exp
expr) Maybe FormatMode
y) =
  let isFunclet :: Bool
isFunclet = case Exp
expr of
                    AppE (VarE Name
a) Exp
_ -> if | Name -> String
TH.nameBase Name
a String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"sh" -> Bool
True
                                          | Name -> String
TH.nameBase Name
a String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"fn" -> Bool
True
                                          | Bool
otherwise             -> Bool
False
                    Exp
_               -> Bool
False
  in (Bool
isFunclet, if Bool
isFunclet then [| [MkAnyFunclet $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
expr)] |]
                 else [| $(Maybe FormatMode -> Q Exp
PyF.getFormatExpr Maybe FormatMode
y) $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
expr) |])

{- ========== reportParserErrorAt ========== -}

reportParserErrorAt :: Ps.ParseError -> Q ()
reportParserErrorAt :: ParseError -> Q ()
reportParserErrorAt ParseError
err = SrcSpan -> String -> Q ()
reportErrorAt SrcSpan
srcSpan String
msg
  where
    msg :: String
msg = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n" ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ ParseError -> [String]
formatErrorMessages ParseError
err

    srcSpan :: SrcSpan
srcSpan = SrcLoc -> SrcLoc -> SrcSpan
mkSrcSpan SrcLoc
loc SrcLoc
loc'

    loc :: SrcLoc
loc = SourcePos -> SrcLoc
srcLocFromParserError (ParseError -> SourcePos
Ps.errorPos ParseError
err)
    loc' :: SrcLoc
loc' = SourcePos -> SrcLoc
srcLocFromParserError (SourcePos -> Column -> SourcePos
Ps.incSourceColumn (ParseError -> SourcePos
Ps.errorPos ParseError
err) Column
1)

formatErrorMessages :: Ps.ParseError -> [String]
formatErrorMessages :: ParseError -> [String]
formatErrorMessages ParseError
err
  -- If there is an explicit error message from parsec, use only that
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Message] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Message]
messages = (Message -> String) -> [Message] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Message -> String
PsError.messageString [Message]
messages
  -- Otherwise, uses parsec formatting
  | Bool
otherwise = [String
-> String -> String -> String -> String -> [Message] -> String
PsError.showErrorMessages String
"or" String
"unknown parse error" String
"expecting" String
"unexpected" String
"end of input" (ParseError -> [Message]
PsError.errorMessages ParseError
err)]
  where
    ([Message]
_sysUnExpect, [Message]
msgs1) = (Message -> Bool) -> [Message] -> ([Message], [Message])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (String -> Message
PsError.SysUnExpect String
"" Message -> Message -> Bool
forall a. Eq a => a -> a -> Bool
==) (ParseError -> [Message]
PsError.errorMessages ParseError
err)
    ([Message]
_unExpect, [Message]
msgs2) = (Message -> Bool) -> [Message] -> ([Message], [Message])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (String -> Message
PsError.UnExpect String
"" Message -> Message -> Bool
forall a. Eq a => a -> a -> Bool
==) [Message]
msgs1
    ([Message]
_expect, [Message]
messages) = (Message -> Bool) -> [Message] -> ([Message], [Message])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (String -> Message
PsError.Expect String
"" Message -> Message -> Bool
forall a. Eq a => a -> a -> Bool
==) [Message]
msgs2

srcLocFromParserError :: Ps.SourcePos -> SrcLoc
srcLocFromParserError :: SourcePos -> SrcLoc
srcLocFromParserError SourcePos
sourceLoc = SrcLoc
srcLoc
  where
    line :: Column
line = SourcePos -> Column
Ps.sourceLine SourcePos
sourceLoc
    column :: Column
column = SourcePos -> Column
Ps.sourceColumn SourcePos
sourceLoc
    name :: String
name = SourcePos -> String
Ps.sourceName SourcePos
sourceLoc
    srcLoc :: SrcLoc
srcLoc = FastString -> Column -> Column -> SrcLoc
mkSrcLoc (String -> FastString
forall a. IsString a => String -> a
fromString String
name) Column
line Column
column