-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# LANGUAGE QuasiQuotes #-}

-- | TH utilities used in "Indigo.Common.Expr"
module Indigo.Common.Expr.TH
       ( deriveExprBuildable
       ) where

import Data.Map qualified as M
import Data.Set qualified as S
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Syntax qualified as TH
import Prelude hiding (Const)
import Text.Casing qualified as C

pattern CName :: String -> (TH.Name, [BangType])
pattern $mCName :: forall {r}.
(Name, [BangType]) -> (String -> r) -> (Void# -> r) -> r
CName nm <- (TH.Name (OccName nm) _, _)

type UntypedConstr = (Name, [BangType])

-- | Generates an Buildable instance for a Expr GADT. /Note:/ This will not generate
-- additional constraints to the generated instance if those are required.
-- Inspired by 'deriveGADTNFData' from Util.TH.
deriveExprBuildable :: Name -> Q [Dec]
deriveExprBuildable :: Name -> Q [Dec]
deriveExprBuildable Name
name = do
  TyConI (DataD Cxt
_ Name
dataName [TyVarBndr ()]
vars Maybe Kind
_ [Con]
cons [DerivClause]
_) <- Name -> Q Info
reify Name
name

  let getNameFromVar :: TyVarBndr flag -> Name
getNameFromVar (PlainTV Name
n flag
_) = Name
n
      getNameFromVar (KindedTV Name
n flag
_ Kind
_) = Name
n
      convertTyVars :: Kind -> Kind
convertTyVars Kind
orig = (Element [TyVarBndr ()] -> Kind -> Kind)
-> Kind -> [TyVarBndr ()] -> Kind
forall t b. Container t => (Element t -> b -> b) -> b -> t -> b
foldr (\Element [TyVarBndr ()]
a Kind
b -> Kind -> Kind -> Kind
AppT Kind
b (Kind -> Kind) -> (Name -> Kind) -> Name -> Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Kind
VarT (Name -> Kind) -> Name -> Kind
forall a b. (a -> b) -> a -> b
$ TyVarBndr () -> Name
forall {flag}. TyVarBndr flag -> Name
getNameFromVar TyVarBndr ()
Element [TyVarBndr ()]
a) Kind
orig [TyVarBndr ()]
vars

      -- Unfolds multiple constructors of form "A, B, C :: A -> Stuff"
      -- into a list of tuples of constructor names and their data
      unfoldConstructor :: Con -> [(Name, [BangType])]
unfoldConstructor (GadtC [Name]
cs [BangType]
bangs Kind
_) = (Name -> (Name, [BangType])) -> [Name] -> [(Name, [BangType])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (,[BangType]
bangs) [Name]
cs
      unfoldConstructor (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
c) = Con -> [(Name, [BangType])]
unfoldConstructor Con
c
      unfoldConstructor Con
_ = String -> [(Name, [BangType])]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Non GADT constructors are not supported."

      (Exp
bLeft, Exp
bRight, Exp
comma) = (Lit -> Exp
LitE (String -> Lit
StringL String
"("), Lit -> Exp
LitE (String -> Lit
StringL String
")"), Lit -> Exp
LitE (String -> Lit
StringL String
", "))

      mappendAll :: [Exp] -> Exp
      mappendAll :: [Exp] -> Exp
mappendAll [] = Text -> Exp
forall a. HasCallStack => Text -> a
error Text
"impossible empty list"
      mappendAll (Exp
hd : [Exp]
rest) = (Exp -> Element [Exp] -> Exp) -> Exp -> [Exp] -> Exp
forall t b. Container t => (b -> Element t -> b) -> b -> t -> b
foldl (\Exp
res Element [Exp]
term -> Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
res) (Name -> Exp
VarE (String -> Name
mkName String
"<>")) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
Element [Exp]
term)) Exp
hd [Exp]
rest

      omitUnaryConstr :: (Exp -> Exp) -> UntypedConstr -> Q Clause
      omitUnaryConstr :: (Exp -> Exp) -> (Name, [BangType]) -> Q Clause
omitUnaryConstr Exp -> Exp
fun (Name
conName, [BangType]
_) = do
        Name
var <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"a"
        Clause -> Q Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$
          [Pat] -> Body -> [Dec] -> Clause
Clause
            [Name -> [Pat] -> Pat
ConP Name
conName [Name -> Pat
VarP Name
var]]
            (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
fun (Name -> Exp
VarE Name
var))
            []

      -- Specific to Expr
      toFunName :: (Name, [BangType]) -> String
toFunName (CName String
"Right'") = String
"right"
      toFunName (CName String
"Left'")  = String
"left"
      toFunName (CName String
"Concat'") = String
"concatAll"
      toFunName (CName String
"Int'")  = String
"toInt"
      toFunName (TH.Name (OccName String
nm) NameFlavour
_, [BangType]
_) = Identifier String -> String
C.toCamel (Identifier String -> String) -> Identifier String -> String
forall a b. (a -> b) -> a -> b
$ String -> Identifier String
C.fromHumps String
nm

      operators :: Map String String
      operators :: Map String String
operators = [(String, String)] -> Map String String
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [ (String
"Add", String
"+"), (String
"Sub", String
"-"), (String
"Mul", String
"*"), (String
"Div", String
"/"), (String
"Mod", String
"%")
                             , (String
"Lsl", String
"<<<"), (String
"Lsr", String
">>>"), (String
"Eq'", String
"=="), (String
"Neq", String
"/="), (String
"Le", String
"=<")
                             , (String
"Lt", String
"<"), (String
"Ge", String
">="), (String
"Gt", String
">="), (String
"Or", String
"||"), (String
"Xor", String
"^"), (String
"And", String
"&&")
                             , (String
"Cons", String
".:"), (String
"Concant", String
"<>")
                             ]

      braces :: Set String
      braces :: Set String
braces = [String] -> Set String
forall a. Ord a => [a] -> Set a
S.fromList [String
"Add", String
"Sub", String
"Or", String
"Xor", String
"Lsl", String
"Lsr"]

      makeClause :: UntypedConstr -> Q Clause
      makeClause :: (Name, [BangType]) -> Q Clause
makeClause c :: (Name, [BangType])
c@(CName String
"C") = (Exp -> Exp) -> (Name, [BangType]) -> Q Clause
omitUnaryConstr (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty") (Exp -> Exp) -> (Exp -> Exp) -> Exp -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp -> Exp
AppE (Name -> Exp
ConE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"PrintAsValue")) (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"V") = (Exp -> Exp) -> (Name, [BangType]) -> Q Clause
omitUnaryConstr (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty")) (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"Div") = (Name, [BangType]) -> Q Clause
forall {m :: * -> *} {a}. Quote m => (Name, [a]) -> m Clause
constructBinaryOperatorWithProxy (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"Mod") = (Name, [BangType]) -> Q Clause
forall {m :: * -> *} {a}. Quote m => (Name, [a]) -> m Clause
constructBinaryOperatorWithProxy (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"Contract") = (Name, [BangType]) -> Q Clause
forall {m :: * -> *} {b}. Quote m => (Name, b) -> m Clause
constructUnaryWithProxy (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"ObjMan") = (Exp -> Exp) -> (Name, [BangType]) -> Q Clause
omitUnaryConstr (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty")) (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"Construct") = (Name, [BangType]) -> Q Clause
forall {m :: * -> *} {b}. Quote m => (Name, b) -> m Clause
construct (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"ConstructWithoutNamed") = (Name, [BangType]) -> Q Clause
forall {m :: * -> *} {b}. Quote m => (Name, b) -> m Clause
construct (Name, [BangType])
c
      makeClause c :: (Name, [BangType])
c@(CName String
"StInsertNew") =
        String -> (Name, [BangType]) -> Maybe (String, Bool) -> Q Clause
generalClauseImpl String
"pretty" (Name, [BangType])
c Maybe (String, Bool)
forall a. Maybe a
Nothing
      makeClause c :: (Name, [BangType])
c@(TH.Name (OccName String
nm) NameFlavour
_, [BangType]
_) =
        String -> (Name, [BangType]) -> Maybe (String, Bool) -> Q Clause
generalClauseImpl String
"pretty" (Name, [BangType])
c (Maybe (String, Bool) -> Q Clause)
-> Maybe (String, Bool) -> Q Clause
forall a b. (a -> b) -> a -> b
$ (, String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member String
nm Set String
braces) (String -> (String, Bool)) -> Maybe String -> Maybe (String, Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Map String String -> Maybe String
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
nm Map String String
operators

      -- Construct build instance for operators that have several proxy values
      constructBinaryOperatorWithProxy :: (Name, [a]) -> m Clause
constructBinaryOperatorWithProxy (Name
conName, [a]
bangs) = do
        Name
x <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"a"
        Name
y <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"a"
        let nm :: String
nm = Name -> String
nameBase Name
conName
            op :: String
op = String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe (Text -> String
forall a. HasCallStack => Text -> a
error Text
"Unknown operator") (Maybe String -> String) -> Maybe String -> String
forall a b. (a -> b) -> a -> b
$ String -> Map String String -> Maybe String
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
nm Map String String
operators
            var_x :: Exp
var_x = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty") (Name -> Exp
VarE Name
x)
            var_y :: Exp
var_y = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty") (Name -> Exp
VarE Name
y)
            wilds :: [Pat]
wilds = (a -> Pat) -> [a] -> [Pat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (\a
_ -> Pat
WildP) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
2 [a]
bangs)
            varList :: [Pat]
varList = ((Name -> Pat) -> [Name] -> [Pat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map Name -> Pat
VarP [Name
x, Name
y]) [Pat] -> [Pat] -> [Pat]
forall a. [a] -> [a] -> [a]
++ [Pat]
wilds

            pretties :: [Exp]
pretties = [Exp
var_x, Lit -> Exp
LitE (String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ String
" " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
op String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" "), Exp
var_y]

        Clause -> m Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> m Clause) -> Clause -> m Clause
forall a b. (a -> b) -> a -> b
$
          [Pat] -> Body -> [Dec] -> Clause
Clause
            [Name -> [Pat] -> Pat
ConP Name
conName [Pat]
varList]
            (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
mappendAll [Exp]
pretties)
            []

      -- Construct build instance for constructors with a single argument preceeded
      -- by a single proxy
      constructUnaryWithProxy :: (Name, b) -> m Clause
constructUnaryWithProxy (Name
conName, b
_bangs) = do
        Name
x <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"a"
        let var_x :: Exp
var_x = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"pretty") (Name -> Exp
VarE Name
x)
        Clause -> m Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> m Clause) -> Clause -> m Clause
forall a b. (a -> b) -> a -> b
$
          [Pat] -> Body -> [Dec] -> Clause
Clause
            [Name -> [Pat] -> Pat
ConP Name
conName [Pat
WildP, Name -> Pat
VarP Name
x]]
            (Exp -> Body
NormalB  Exp
var_x)
            []


      construct :: (Name, b) -> m Clause
construct (Name
conName, b
_) = do
        Name
proxy <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"proxy"
        Name
r <- String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"rec"
        let showTypeRep :: Exp -> Exp
showTypeRep = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"show") (Exp -> Exp) -> (Exp -> Exp) -> Exp -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"typeRep")
        Exp
mappendRec <- [| \x -> mconcat (intersperse "," (recordToList (rmap (\ex -> Const (pretty ex)) x))) |]
        Clause -> m Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> m Clause) -> Clause -> m Clause
forall a b. (a -> b) -> a -> b
$
          [Pat] -> Body -> [Dec] -> Clause
Clause
            [Name -> [Pat] -> Pat
ConP Name
conName [Name -> Pat
VarP Name
proxy, Name -> Pat
VarP Name
r]]
            (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$
              [Exp] -> Exp
mappendAll [Exp -> Exp
showTypeRep (Name -> Exp
VarE Name
proxy), Exp
bLeft, (Exp -> Exp -> Exp
AppE Exp
mappendRec (Name -> Exp
VarE Name
r)), Exp
bRight])
            []

      -- Constructs a clause "build (ConName a1 a2) = "CON_NAME" <> "(" <> pretty a1 <> pretty a2 <> ")"
      -- The first argument is a workaround: it returns name of printing function either "pretty" or "show"
      -- by index argument.
      -- This is mainly for 'StInsertNew' constructor to print error with show
      generalClauseImpl :: String -> (Name, [BangType]) -> Maybe (String, Bool) -> Q Clause
      generalClauseImpl :: String -> (Name, [BangType]) -> Maybe (String, Bool) -> Q Clause
generalClauseImpl String
funName c :: (Name, [BangType])
c@(Name
conName, [BangType]
bangs) Maybe (String, Bool)
isInfix = do
        [Name]
varNames <- (BangType -> Q Name) -> [BangType] -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\BangType
_ -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"a") [BangType]
bangs

        -- useful constants
        let funStr :: Exp
funStr = Lit -> Exp
LitE (String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ (Name, [BangType]) -> String
toFunName (Name, [BangType])
c)
        let pretties :: [Exp]
pretties =
              case (Name -> Exp) -> [Name] -> [Exp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map (\Name
e -> Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
funName) (Name -> Exp
VarE Name
e)) [Name]
varNames of
                -- Nullary operator: "function_name"
                []    -> [Exp
funStr]
                -- Infix binary operator without braces: pretty x <> " operator " <> pretty y
                [Exp
x, Exp
y] | Just (String
op, Bool
False) <- Maybe (String, Bool)
isInfix ->
                  [Exp
x, Lit -> Exp
LitE (String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ String
" " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
op String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" "), Exp
y]
                -- Infix binary operator with braces: "(" <> pretty x <> " operator " <> pretty y <> ")"
                [Exp
x, Exp
y] | Just (String
op, Bool
True) <- Maybe (String, Bool)
isInfix ->
                  [Exp
bLeft, Exp
x, Lit -> Exp
LitE (String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ String
" " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
op String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" "), Exp
y, Exp
bRight]
                -- Infix binary operator with braces: "function_name" <> "(" <> pretty x1 <> "," ... <> pretty xn <> ")"
                [Exp]
xs -> Exp
funStr Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: Exp
bLeft Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: (Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
intersperse Exp
comma [Exp]
xs [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ [Exp
bRight])
        Clause -> Q Clause
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$
          [Pat] -> Body -> [Dec] -> Clause
Clause
            [Name -> [Pat] -> Pat
ConP Name
conName ((Name -> Pat) -> [Name] -> [Pat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map Name -> Pat
VarP [Name]
varNames)]
            (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
mappendAll [Exp]
pretties)
            []

  [Clause]
clauses <- ((Name, [BangType]) -> Q Clause)
-> [(Name, [BangType])] -> Q [Clause]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name, [BangType]) -> Q Clause
makeClause ([(Name, [BangType])] -> Q [Clause])
-> [(Name, [BangType])] -> Q [Clause]
forall a b. (a -> b) -> a -> b
$ [Con]
cons [Con] -> (Con -> [(Name, [BangType])]) -> [(Name, [BangType])]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Con -> [(Name, [BangType])]
unfoldConstructor
  [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return
    [ Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
InstanceD
        Maybe Overlap
forall a. Maybe a
Nothing
        []
        (Kind -> Kind -> Kind
AppT (Name -> Kind
ConT (Name -> Kind) -> Name -> Kind
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"Buildable") (Kind -> Kind
convertTyVars (Kind -> Kind) -> Kind -> Kind
forall a b. (a -> b) -> a -> b
$ Name -> Kind
ConT Name
dataName))
        [Name -> [Clause] -> Dec
FunD (String -> Name
mkName String
"build") [Clause]
clauses]
    ]