fradrive/src/Utils/TH.hs

-- SPDX-FileCopyrightText: 2022-24 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Steffen Jost <jost@tcs.ifi.lmu.de>,Steffen Jost <s.jost@fraport.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Utils.TH where
-- Common Utility Functions that require TemplateHaskell

-- import Data.Char

import ClassyPrelude.Yesod
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
-- import Control.Monad
-- import Control.Monad.Trans.Class
-- import Control.Monad.Trans.Maybe
-- import Control.Monad.Trans.Except

import Data.List ((!!), foldl)

import Control.Lens
import Control.Monad.Fail

import Utils.I18n

import qualified Data.Char as Char
import Data.Universe (Universe, Finite)
import qualified Data.Map as Map
import qualified Data.Text as Text

import Utils.PathPiece

--------------------
-- Non-TH Helpers --
--------------------

-- | Invert a permutation specified as a list of all numbers 1..n, each number occurring precisely once, i.e. sort l == [1..(length l)]
--   Useful for @permuteFun
inversePermutation :: [Int] -> [Int]
inversePermutation l = map snd . sortOn fst $ zip l [1..]


------------
-- Tuples --
------------

-- Alternatively uses lenses: "^. _3" projects the 3rd component of an n-tuple for any n >=3, requires import Control.Lens
projNI :: Int -> Int -> ExpQ -- generic projection gives I-th element of N-tuple, i.e. snd3 = $(projNI 3 2)  --ghci -fth
-- Usage like so: $(projN n m) :: (t1,..,tn) -> tm (for m<=n)
projNI n i = do
  x <- newName "x"
  let rhs = varE x
  let pat = tupP $ replicate (pred i) wildP ++ varP x : replicate (n-i) wildP
  lam1E pat rhs


-- | Generic projections N-tuples that are actually left-associative pairs
--   i.e. @$(leftAssociativePairProjection c n m :: (..(t1 `c` t2) `c` .. `c` tn) -> tm@ (for m<=n)
leftAssociativePairProjection :: Name -> Int -> Int -> ExpQ
leftAssociativePairProjection constructor n i = do
    x <- newName "x"
    lamE [pat x n] (varE x)
  where
    pat x 1       = varP x
    pat x w
      | w==i      = conP constructor [wildP, varP x]
      | otherwise = conP constructor [pat x (pred w), wildP]

-- | Generic projections n-tuples that are actually left-associative pairs with differing constructors
--   i.e. @$(leftAssociativePairProjection [c1,c2,..,cn] m :: (..(t1 `c1` t2) `c2` .. `cn` t(n+1) -> tm@ (for m<=n+1)
leftAssociativeProjection :: [Name] -> Int -> ExpQ
leftAssociativeProjection constructors@(length -> n) (pred -> i)
  | n < i  = error $ "Util.TH.leftAssociativeProjection not given enough constructors: " <> show constructors
  | otherwise = do
    x <- newName "x"
    lamE [pat x n] (varE x)
  where
    pat x 0       = varP x
    pat x w@(pred -> v)
      | w==i      = conP (constructors !! v) [wildP, varP x]
      | otherwise = conP (constructors !! v) [pat x v, wildP]

-- Extract constructor names from a type definition of left-associative pair-constructors (i.e. Esqueleto-Joins in a table-expression type)
extractConstructorNames :: Name -> Q [Name]
extractConstructorNames td = do
  TyConI (TySynD _ [] ty) <- reify td -- executed at compile time, so failure is acceptable
  reverse . concat <$> mapM getDataConstructors (go ty)
  where
    go :: Type -> [Name]
    go (AppT (AppT (ConT name) rest) _) = name : go rest
    go _                                = []

    -- At this point we have the Type-Constructors, but we actually need the Data-Constructors:
    getDataConstructors :: Name -> Q [Name]
    getDataConstructors conName = do
      info <- reify conName
      case info of
        TyConI (DataD    _ _ _ _ constr _) -> return $ concatMap getConNames constr
        TyConI (NewtypeD _ _ _ _ constr _) -> return $           getConNames constr
        _ -> return []

    getConNames :: Con -> [Name]
    getConNames (NormalC  name _) = [name]
    getConNames (RecC     name _) = [name]
    getConNames (InfixC _ name _) = [name]
    getConNames (ForallC _ _ con) = getConNames con
    getConNames _ = []

{-
Example:

Suppose
  type LmsTableExpr = (                   E.SqlExpr (Entity QualificationUser)
                        `E.InnerJoin`     E.SqlExpr (Entity User)
                        `E.InnerJoin`     E.SqlExpr (Entity LmsUser)
                        `E.LeftOuterJoin` E.SqlExpr (Maybe (Entity QualificationUserBlock))
                      )
then
  info <- reify ''LmsTableExpr
with
  info = TyConI (TySynD Handler.Utils.LMS.LmsTableExpr []
    (AppT
      (AppT
        (ConT Database.Esqueleto.Internal.Internal.LeftOuterJoin)
        (AppT
          (AppT
            (ConT Database.Esqueleto.Internal.Internal.InnerJoin)
            (AppT
              (AppT
                (ConT Database.Esqueleto.Internal.Internal.InnerJoin)
                (AppT
                  (ConT Database.Esqueleto.Internal.Internal.SqlExpr)
                  (AppT
                    (ConT Database.Persist.Class.PersistEntity.Entity)
                    (ConT Model.QualificationUser)
              ) ) )
              (AppT (ConT Database.Esqueleto.Internal.Internal.SqlExpr)
                (AppT
                  (ConT Database.Persist.Class.PersistEntity.Entity)
                  (ConT Model.User)
        ) ) ) ) )
        (AppT
          (ConT Database.Esqueleto.Internal.Internal.SqlExpr)
          (AppT
            (ConT Database.Persist.Class.PersistEntity.Entity)
            (ConT Model.LmsUser)
    ) ) ) )
    (AppT
      (ConT Database.Esqueleto.Internal.Internal.SqlExpr)
        (AppT
          (ConT GHC.Maybe.Maybe)
          (AppT
            (ConT Database.Persist.Class.PersistEntity.Entity)
            (ConT Model.QualificationUserBlock)
  ) ) ) ) )
-}


---------------
-- Functions --
---------------

-- | Generic permutation of function, i.e. $(permuteFun [2,1]) == flip
--   Note that the function is applied to the permuted arguments, so usually the inverted permutation is required (see @inversePermutation)
permuteFun :: [Int] -> ExpQ
permuteFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
        ln  = length perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..ln] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

altFun :: [Int] -> ExpQ -- generic permutation/repetition of function arguments, i.e. $(permuteFun [2,1]) == flip
altFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
        mx  = maximum $ impureNonNull perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..mx] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

-- |
curryN :: Int -> ExpQ
curryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = map varP (fn:xs)
  let tup = tupE (map varE xs)
  let rhs = appE (varE fn) tup
  lamE pat rhs

uncurryN :: Int -> ExpQ
uncurryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = [VarP fn, TupP (map VarP xs)]
  let rhs = foldl AppE (VarE fn) (map VarE xs)
  return $ LamE pat rhs


afterN :: Int -> ExpQ -- apply a function after another of arity N, i.e. $(afterN 1) = (.)
afterN n = do
  f <- newName "f"
  g <- newName "g"
  --let rhs = [|$(curryN n) (g . ($(uncurryN n) f))|]
  lamE [varP g, varP f] [|$(curryN n) $(varE g) . $(uncurryN n) $(varE f)|]


-- Special Show-Instances for Themes
deriveShowWith :: (String -> String) -> Name -> Q [Dec]
deriveShowWith = deriveSimpleWith ''Show 'show

deriveSimpleWith ::  Name -> Name -> (String -> String) -> Name -> Q [Dec]
deriveSimpleWith cls fun strOp ty = do
    (TyConI tyCon)  <- reify ty
    (tyConName, cs) <- case tyCon of
        DataD [] nm [] _ cs _ -> return (nm, cs)
        _ -> fail "deriveShowTheme: tyCon must be a plain datatype enumeration"
    let instanceT = conT cls `appT` conT tyConName
    return <$> instanceD (return []) instanceT [genDecs cs]
  where
    genDecs :: [Con] -> Q Dec
    genDecs cs = funD fun (map genClause cs)

    genClause :: Con -> Q Clause
    genClause (NormalC name []) =
      let pats = [ConP name []]
          body = NormalB $ LitE $ StringL $ strOp $ nameBase name
      in return $ Clause pats body []
    genClause _ = fail "deriveShowTheme: constructors not allowed to have arguments"

embedRenderMessage :: Name -- ^ Foundation type
                   -> Name -- ^ Type to embed into message type
                   -> (Text -> Text) -- ^ Mangle constructor names
                   -> DecsQ
-- ^ @embedRenderMessage ''Foundation ''MessageType mangle@ declares a
--   `RenderMessage Foundation MessageType` instance expecting there
--   to be one constructor for each constructor of @MessageType@ in
--   scope, taking the same arguments:
--
-- > data NewMessage = NewMessageOne | NewMessageTwo | NewMessageThree
-- > data FoundationMessage = MsgOne | MsgThree
-- > data FoundationEvenMessage = MsgTwo
-- >
-- > -- embedRenderMessage ''Foundation ''NewMessage $ drop 2 . splitCamel
-- > instance RenderMessage Foundation NewMessage where
-- >   renderMessage f ls = \case
-- >     NewMessageOne -> renderMessage f ls MsgOne
-- >     NewMessageTwo -> renderMessage f ls MsgTwo
embedRenderMessage f inner mangle = do
  DatatypeInfo{..} <- reifyDatatype inner

  f' <- newName "f"
  ls <- newName "ls"

  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let constr = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP constructorName $ map varP vars) (normalB [e|renderMessage $(varE f') $(varE ls) $(constr)|]) []

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT inner)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ lamCaseE matches) []
      ]
    ]

-- ^ Like @embedRenderMessage, but for newtype definitions
embedRenderMessageVariant :: Name -- ^ Foundation Type
                          -> Name -- ^ Name of newtype
                          -> (Text -> Text) -- ^ Mangle constructor names
                          -> DecsQ
embedRenderMessageVariant f newT mangle = do
  [ConstructorInfo{ constructorName = newtypeName, constructorFields = [ ConT newtypeInner ] }] <- datatypeCons <$> reifyDatatype newT
  DatatypeInfo{..} <- reifyDatatype newtypeInner

  f' <- newName "f"
  ls <- newName "ls"

  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let body = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP newtypeName [conP constructorName $ map varP vars]) (normalB [e|renderMessage $(varE f') $(varE ls) $(body)|]) []

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT newT)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ lamCaseE matches) []
      ]
    ]


dispatchTH :: Name -- ^ Datatype to pattern match
           -> ExpQ
-- ^ Produces a lambda-case-expression matching all constructors of the named datatype and calling a function (named after the constructor prefixed with @dispatch@) on the fields of each constructor
dispatchTH dType = do
  DatatypeInfo{..} <- reifyDatatype dType
  let
    matches = map mkMatch datatypeCons
    mkMatch ConstructorInfo{..} = do
      pats <- forM constructorFields $ \_ -> newName "x"
      let fName = mkName $ "dispatch" <> nameBase constructorName
      match (conP constructorName $ map varP pats) (normalB $ foldl (\e pat -> e `appE` varE pat) (varE fName) pats) []
  lamCaseE matches


mkI18nWidgetEnum :: String -> FilePath -> DecsQ
mkI18nWidgetEnum (splitCamel -> namebase) basename = do
  itemsAvailable <- i18nWidgetFilesAvailable' basename
  let items = Map.mapWithKey (\k _ -> typPrefix <> unPathPiece k) itemsAvailable
  sequence
    [ dataD (cxt []) dataName [] Nothing
        [ normalC (mkName conName) []
        | (_, conName) <- Map.toAscList items
        ]
        [ derivClause (Just StockStrategy)
            [ conT ''Eq
            , conT ''Ord
            , conT ''Read
            , conT ''Show
            , conT ''Enum
            , conT ''Bounded
            , conT ''Generic
            -- , conT ''Typeable
            ]
        , derivClause (Just AnyclassStrategy)
            [ conT ''Universe
            , conT ''Finite
            , conT ''NFData
            ]
        ]
    , instanceD (cxt []) (conT ''PathPiece `appT` conT dataName)
      [ funD 'toPathPiece
        [ clause [conP (mkName con) []] (normalB . litE . stringL $ repack int) []
        | (int, con) <- Map.toList items
        ]
      , funD 'fromPathPiece
        [ clause [varP $ mkName "t"]
          ( guardedB
            [ (,) <$> normalG [e|$(varE $ mkName "t") == int|] <*> [e|Just $(conE $ mkName con)|]
            | (int, con) <- Map.toList items
            ]) []
        , clause [wildP] (normalB [e|Nothing|]) []
        ]
      ]
    , sigD (mkName $ valPrefix <> "ItemMap") [t|Map Text $(conT dataName)|]
    , funD (mkName $ valPrefix <> "ItemMap")
        [ clause [] (normalB [e| Map.fromList $(listE . map (\(int, con) -> tupE [litE . stringL $ repack int, conE $ mkName con]) $ Map.toList items) |]) []
        ]
    ]
  where
    unPathPiece :: Text -> String
    unPathPiece = repack . mconcat . map (over _head Char.toUpper) . Text.splitOn "-"

    dataName = mkName $ typPrefix <> "Item"

    typPrefix = concat $ over (takingWhile Char.isLower $ _head . traverse) Char.toUpper namebase
    valPrefix = concat $ over (takingWhile Char.isUpper $ _head . traverse) Char.toLower namebase