TH.hs

-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# LANGUAGE QuasiQuotes #-}

module Indigo.Internal.Expr.TH
       ( deriveExprBuildable
       ) where

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

pattern CName :: String -> (TH.Name, [BangType])
pattern 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 = do
  TyConI (DataD _ dataName vars _ cons _) <- reify name

  let getNameFromVar (PlainTV n) = n
      getNameFromVar (KindedTV n _) = n
      convertTyVars orig = foldr (\a b -> AppT b . VarT $ getNameFromVar a) orig vars

      -- Unfolds multiple constructors of form "A, B, C :: A -> Stuff"
      -- into a list of tuples of constructor names and their data
      unfoldConstructor (GadtC cs bangs _) = map (,bangs) cs
      unfoldConstructor (ForallC _ _ c) = unfoldConstructor c
      unfoldConstructor _ = fail "Non GADT constructors are not supported."

      (bLeft, bRight, comma) = (LitE (StringL "("), LitE (StringL ")"), LitE (StringL ", "))

      mappendAll :: [Exp] -> Exp
      mappendAll [] = error "impossible empty list"
      mappendAll (hd : rest) = foldl (\res term -> InfixE (Just res) (VarE (mkName "<>")) (Just term)) hd rest

      omitUnaryConstr :: (Exp -> Exp) -> UntypedConstr -> Q Clause
      omitUnaryConstr fun (conName, _) = do
        var <- newName "a"
        return $
          Clause
            [ConP conName [VarP var]]
            (NormalB $ fun (VarE var))
            []

      -- Specific to Expr
      toFunName (CName "Right'") = "right"
      toFunName (CName "Left'")  = "left"
      toFunName (CName "Concat'") = "concatAll"
      toFunName (CName "Int'")  = "toInt"
      toFunName (TH.Name (OccName nm) _, _) = C.toCamel $ C.fromHumps nm

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

      braces :: Set String
      braces = S.fromList ["Add", "Sub", "Or", "Xor", "Lsl", "Lsr"]

      makeClause :: UntypedConstr -> Q Clause
      makeClause c@(CName "C") = omitUnaryConstr (AppE (VarE $ mkName "pretty") . AppE (ConE $ mkName "PrintAsValue")) c
      makeClause c@(CName "V") = omitUnaryConstr (AppE (VarE $ mkName "pretty")) c
      makeClause c@(CName "ObjMan") = omitUnaryConstr (AppE (VarE $ mkName "pretty"))  c
      makeClause c@(CName "Construct") = construct c
      makeClause c@(CName "ConstructWithoutNamed") = construct c
      makeClause c@(CName "StInsertNew") =
        generalClauseImpl "pretty" c Nothing
      makeClause c@(TH.Name (OccName nm) _, _) =
        generalClauseImpl "pretty" c $ (, S.member nm braces) <$> M.lookup nm operators

      construct (conName, _) = do
        proxy <- newName "proxy"
        r <- newName "rec"
        let showTypeRep = AppE (VarE $ mkName "show") . AppE (VarE $ mkName "typeRep")
        mappendRec <- [| \x -> mconcat (intersperse "," (recordToList (rmap (\ex -> Const (pretty ex)) x))) |]
        return $
          Clause
            [ConP conName [VarP proxy, VarP r]]
            (NormalB $
              mappendAll [showTypeRep (VarE proxy), bLeft, (AppE mappendRec (VarE r)), 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 funName c@(conName, bangs) isInfix = do
        varNames <- traverse (\_ -> newName "a") bangs
        -- useful constants
        let funStr = LitE (StringL $ toFunName c)
        let pretties =
              case map (\e -> AppE (VarE $ mkName funName) (VarE e)) varNames of
                -- Nullary operator: "function_name"
                []    -> [funStr]
                -- Infix binary operator without braces: pretty x <> " operator " <> pretty y
                [x, y] | Just (op, False) <- isInfix ->
                  [x, LitE (StringL $ " " <> op <> " "), y]
                -- Infix binary operator with braces: "(" <> pretty x <> " operator " <> pretty y <> ")"
                [x, y] | Just (op, True) <- isInfix ->
                  [bLeft, x, LitE (StringL $ " " <> op <> " "), y, bRight]
                -- Infix binary operator with braces: "function_name" <> "(" <> pretty x1 <> "," ... <> pretty xn <> ")"
                xs -> funStr : bLeft : (intersperse comma xs ++ [bRight])
        return $
          Clause
            [ConP conName $ map VarP varNames]
            (NormalB $ mappendAll pretties)
            []

  clauses <- traverse makeClause $ cons >>= unfoldConstructor
  return
    [ InstanceD
        Nothing
        []
        (AppT (ConT $ mkName "Buildable") (convertTyVars $ ConT dataName))
        [FunD (mkName "build") clauses]
    ]