GenService.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE CPP #-}

module Thrift.Compiler.GenService
  ( genServiceDecls
  , genServiceImports
  , genServiceExports
  ) where

import Control.Monad
#if __GLASGOW_HASKELL__ <= 804
import Data.Monoid ((<>))
#endif
import Data.Maybe (isNothing)
import qualified Data.Set as Set
import Data.Some
import Data.Text (Text)
import qualified Data.Text as Text
import Language.Haskell.Exts.Syntax as HS

import Thrift.Compiler.GenStruct
import Thrift.Compiler.GenUtils
import Thrift.Compiler.Plugins.Haskell
import Thrift.Compiler.Types hiding (Decl(..))


-- | All things required to generate a Service.hs file

genServiceExports :: HS Service -> [HS.ExportSpec ()]
genServiceExports s =
  (if isEmptyService s
  then HS.EAbs () (NoNamespace ()) (unqualSym $ commandTypeName s)
  else HS.EThingWith () (EWildcard () 0) (unqualSym $ commandTypeName s) [])
  : map (HS.EVar () . UnQual () . Ident ())
    ["reqName'", "reqParser'", "respWriter'", "methodsInfo'"]

genServiceImports :: Text.Text -> HS Service -> Set.Set Import
genServiceImports this s@Service{..} =
  foldr (Set.union . importsForFunc) baseImports (getServiceFunctions s)
  where
    importsForFunc Function{..} =
      foldr (Set.union . getImports) (retImport funResolvedType) funArgs
      where
        retImport Nothing = Set.empty
        retImport (Just (Some i)) = typeToImport i

    getImports :: HS (Field u) -> Set.Set Import
    getImports Field{..} = typeToImport fieldResolvedType

    baseImports = Set.fromList
      [ QImport "Prelude" "Prelude"
      , QImport "Control.Exception" "Exception"
      , QImport "Control.Monad.ST.Trans" "ST"
      , QImport "Control.Monad.Trans.Class" "Trans"
      , QImport "Data.ByteString.Builder" "Builder"
      , QImport "Data.Default" "Default"
      , QImport "Data.HashMap.Strict" "HashMap"
      , QImport "Data.Map.Strict" "Map"
      , QImport "Data.Int" "Int"
      , QImport "Data.Proxy" "Proxy"
      , QImport "Data.Text" "Text"
      , QImport "Thrift.Binary.Parser" "Parser"
      , QImport "Thrift.Protocol.ApplicationException.Types" "Thrift"
      , QImport "Thrift.Processor" "Thrift"
      , SymImport "Prelude"
        [ "<$>", "<*>", "++", ".", "==" ]
      , SymImport "Control.Applicative" [ "<*", "*>" ]
      , SymImport "Data.Monoid" [ "<>" ]
      ] `Set.union`
      (case resolvedName . fst . supResolvedName <$> serviceSuper of
        Nothing -> Set.empty
        Just (UName n) -> mkImport this n
        Just (QName m n) -> mkImport m n)
    mkImport m n =
      Set.singleton $ QImport (Text.intercalate "." [m, n, "Service"]) n

genServiceDecls :: HS Service -> [HS.Decl ()]
genServiceDecls s =
  [ genCommandDT s
  , processorInstance
  ] ++ concat
  [ genReqName s
  , genReqParser s
  , genRespWriter s
  , genMethodsInfo s
  ]

  where
    processorInstance = HS.InstDecl () Nothing
      (HS.IRule () Nothing Nothing $
       HS.IHApp ()
         (HS.IHCon () $ qualSym "Thrift" "Processor")
         (HS.TyCon () $ unqualSym $ commandTypeName s)
      )
      (Just $ classFuns ++ additional)

    classFuns = flip map [ "reqName" , "reqParser" , "respWriter" ] $ \fn ->
      HS.InsDecl () $ HS.FunBind () [ HS.Match () (textToName fn) []
        (HS.UnGuardedRhs () $ con (fn <> "'")) Nothing
      ]

    additional =
      [ HS.InsDecl () $ HS.FunBind ()
        [ HS.Match () (textToName "methodsInfo") [HS.PWildCard ()]
          (HS.UnGuardedRhs () $ var "methodsInfo'") Nothing
        ]
      ]

-- | Generates a GADT for all functions that the service can implement
-- If extending a service, adds a "Super<Name>" constructor to forward
-- requests along
genCommandDT :: HS Service -> HS.Decl ()
genCommandDT s@Service{..} = HS.GDataDecl () (HS.DataType ()) Nothing
  (HS.DHApp ()
     (HS.DHead () (textToName $ commandTypeName s))
     (HS.UnkindedVar () (HS.Ident () "a")))
  Nothing
  genFunctions
  mzero
  where
    genFunctions = map genDTFunction (getServiceFunctions s) ++
      case fst . supResolvedName <$> serviceSuper of
        Nothing -> []
        Just Name{..} -> [genSuper $ localName resolvedName]

    genSuper superName =
      HS.GadtDecl () (textToName $ "Super" <> superName)
#if MIN_VERSION_haskell_src_exts(1,21,0)
        Nothing Nothing
#endif
        Nothing $
        HS.TyFun ()
        (qualType superName (superName <> "Command") `appT` simpleType "a")
        (simpleType (commandTypeName s) `appT` tvar "a")

    genDTFunction Function{..} =
      HS.GadtDecl () (textToName $ toCamel funName)
#if MIN_VERSION_haskell_src_exts(1,21,0)
        Nothing Nothing
#endif
        Nothing (mkArgs funArgs)
      where
        mkArgs = foldr
          (\Field{..} ->
            HS.TyFun () $ genType $ qualifyType "Types" fieldResolvedType)
          (simpleType (commandTypeName s) `appT`
            maybe (HS.unit_tycon ())
            (`withSome` (genType . qualifyType "Types")) funResolvedType)

-- | Generates a function that returns a Text name for the given function
genReqName :: HS Service -> [HS.Decl ()]
genReqName s@Service{..} =
  [ HS.TypeSig () [textToName "reqName'"] $ HS.TyFun ()
      (HS.TyApp () (simpleType $ commandTypeName s) (simpleType "a"))
      (qualType "Text" "Text")
  ] ++ map genFunction theseFunctions ++
    case fst . supResolvedName <$> serviceSuper of
      Nothing -> noParentBody (null theseFunctions)
      Just Name{..} -> [genSuper $ localName resolvedName]

  where
    theseFunctions = getServiceFunctions s
    genSuper superName = HS.FunBind ()
      [ HS.Match () (textToName "reqName'")
        [HS.PParen () $ HS.PApp ()
           (HS.UnQual () $ textToName $ "Super" <> superName) [pvar "x"]]
        (HS.UnGuardedRhs () $ qvar superName "reqName'" `app` var "x")
        Nothing
      ]

    genFunction Function{..} = HS.FunBind ()
      [ HS.Match () (textToName "reqName'")
        [HS.PApp () (unqualSym $ toCamel funName) $
          map (\Field{..} -> pvar ("__field__" <> fieldName)) funArgs]
        (HS.UnGuardedRhs () $ stringLit funName)
        Nothing
      ]

    noParentBody False = []

    noParentBody True = [ HS.FunBind ()
      [ HS.Match () (textToName "reqName'")
        [ HS.PWildCard () ]
        (HS.UnGuardedRhs () (stringLit "unknown function"))
        Nothing
      ]]

-- | Generates a function that parses the input message as appropriate
genReqParser :: HS Service -> [HS.Decl ()]
genReqParser s@Service{..} =
  [ HS.TypeSig () [textToName "reqParser'"] $
    HS.TyForall () Nothing
#if MIN_VERSION_haskell_src_exts(1,22,0)
      (Just $ HS.CxSingle () $ HS.TypeA () (HS.TyApp () (HS.TyCon () (qualSym "Thrift" "Protocol"))
        (tvar "p" ))) $
#else
      (Just $ HS.CxSingle () $ HS.ClassA () (qualSym "Thrift" "Protocol")
        [ tvar "p" ]) $
#endif
    HS.TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
    HS.TyFun () (qualType "Text" "Text") $
    HS.TyApp () (qualType "Parser" "Parser") $ HS.TyParen () $
      qualType "Thrift" "Some" `appT` simpleType (commandTypeName s)
  ] ++
  map genFunction (getServiceFunctions s) ++
  genSuper serviceSuper
  where
    genSuper Nothing = [ HS.FunBind ()
      [ HS.Match () (textToName "reqParser'")
        [ HS.PWildCard (), pvar "funName"]
        (HS.UnGuardedRhs () (qvar "Prelude" "errorWithoutStackTrace" `app` HS.Paren ()
           (infixApp "++" (stringLit "unknown function call: ")
           (qvar "Text" "unpack" `app` var "funName"))))
        Nothing
      ]]
    genSuper (Just Super{..}) =
      let n = localName $ resolvedName $ fst supResolvedName in
      [ HS.FunBind ()
      [ HS.Match () (textToName "reqParser'")
        [ pvar "_proxy", pvar "funName" ]
        (HS.UnGuardedRhs () $ HS.Do ()
          [ HS.Generator () (HS.PApp () (qualSym "Thrift" "Some") [ pvar "x" ])
             (qvar (toCamel n) "reqParser'" `app`
              var "_proxy" `app`
              var "funName")
          , HS.Qualifier ()
             (qvar "Prelude" "return" `app` HS.Paren ()
               (qcon "Thrift" "Some" `app`
                HS.Paren () (
                  con ("Super" <> n) `app`
                  var "x"
                )))
          ])
        Nothing
      ]]

    genFunction Function{..} = HS.FunBind ()
      [ HS.Match () (textToName "reqParser'")
        [ pvar "_proxy", stringP funName ]
        (genFieldParser (map (qualifyField "Types") funArgs)
         (toCamel funName)
         (app (qcon "Thrift" "Some")))
        Nothing
      ]

-- | Generates a function that builds the appropriate response type
genRespWriter :: HS Service -> [HS.Decl ()]
genRespWriter s@Service{..} =
  [ HS.TypeSig () [textToName "respWriter'"] $
    HS.TyForall () Nothing
#if MIN_VERSION_haskell_src_exts(1,22,0)
      (Just $ HS.CxSingle () $ HS.TypeA () (HS.TyApp () (HS.TyCon () (qualSym "Thrift" "Protocol"))
        (tvar "p" ))) $
#else
      (Just $ HS.CxSingle () $ HS.ClassA () (qualSym "Thrift" "Protocol")
        [ tvar "p" ]) $
#endif
    HS.TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
    HS.TyFun () (qualType "Int" "Int32") $
    HS.TyFun () (simpleType (commandTypeName s) `appT` tvar "a") $
    HS.TyFun ()
    (qualType "Prelude" "Either" `appT`
     qualType "Exception" "SomeException" `appT`
     tvar "a") $
    HS.TyTuple () Boxed
      [ qualType "Builder" "Builder"
      , qualType "Prelude" "Maybe" `appT`
          HS.TyTuple() Boxed
            [ qualType "Exception" "SomeException"
            , qualType "Thrift" "Blame"
            ]
      ]
  ] ++
  map genFunction theseFunctions ++
  genSuper serviceSuper (null theseFunctions)
  where
    theseFunctions = getServiceFunctions s
    genSuper Nothing True =
      [ HS.FunBind ()
        [ HS.Match () (textToName "respWriter'")
          [ HS.PWildCard () ]
          (HS.UnGuardedRhs () (qvar "Prelude" "errorWithoutStackTrace" `app` HS.Paren ()
            (stringLit "unknown function")))
          Nothing
        ]
      ]
    genSuper Nothing False = []
    genSuper (Just Super{..}) _ =
      let n = localName $ resolvedName $ fst supResolvedName
      in
        [ HS.FunBind ()
          [ HS.Match () (textToName "respWriter'")
            [ pvar "_proxy"
            , pvar "_seqNum"
            , HS.PParen () $ HS.PApp ()
              (HS.UnQual () $ textToName $ "Super" <> n) [pvar "_x"]
            , pvar "_r" ]
            (HS.UnGuardedRhs () $
             qvar n "respWriter'" `app`
             var "_proxy" `app`
             var "_seqNum" `app`
             var "_x" `app`
             var "_r"
            )
            Nothing
          ]]

    genFunction :: HS Function -> Decl ()
    genFunction Function{..} = HS.FunBind ()
      [ HS.Match () (textToName "respWriter'")
        [ pvar "_proxy"
        , pvar "_seqNum"
        , HS.PRec () (unqualSym $ toCamel funName) []
        , pvar "_r" ]
        (HS.UnGuardedRhs () $
          Tuple () Boxed
         [ infixApp "<>"
            (infixApp "<>"
              (protocolFun "genMsgBegin" `app`
              stringLit funName `app`
              var "_msgType" `app`
              var "_seqNum")
              (var "_msgBody"))
            (protocolFun "genMsgEnd")
          , var "_msgException" ])
        (Just $ BDecls ()
         [ PatBind () (PTuple () Boxed
            [ pvar "_msgType"
            , pvar "_msgBody"
            , pvar "_msgException" ])
           (UnGuardedRhs () $ Case () (var "_r")
            [ Alt () (PApp () (qualSym "Prelude" "Left") [ pvar "_ex" ])
              (GuardedRhss () $
               [ GuardedRhs ()
                 [ Generator ()
                   (PApp () (qualSym "Prelude" "Just")
                    [ PAsPat () (textToName "_e") $
                      PRec () (qualSym "Thrift" "ApplicationException") []
                    ]) $
                   qvar "Exception" "fromException" `app` var "_ex"
                 ] $
                 genRespTup genEXCEPTION
                  (protocolFun "buildStruct" `app` var "_e")
                  (Just (var "_ex", "ServerError"))
               ] ++
               map genExceptionCase funExceptions ++
               [ GuardedRhs () [ Qualifier () $ qvar "Prelude" "otherwise" ] $
                 Let ()
                  (BDecls ()
                    [ PatBind ()
                        (pvar "_e")
                        (UnGuardedRhs () $
                          qcon "Thrift" "ApplicationException" `app`
                          (qvar "Text" "pack" `app`
                            (qvar "Prelude" "show" `app` var "_ex")) `app`
                          qcon "Thrift"
                            "ApplicationExceptionType_InternalError")
                        Nothing ]) $
                  genRespTup genEXCEPTION
                    (protocolFun "buildStruct" `app` var "_e")
                    (Just
                      ( qvar "Exception" "toException" `app` var "_e"
                      , "ServerError" ))
               ])
              Nothing
            , Alt () (PApp () (qualSym "Prelude" "Right") [ pvar "_result" ])
              (UnGuardedRhs () $ genRespTup genREPLY
                (protocolFun "genStruct" `app`
                  HS.List ()
                    (genResp $ qualifyResolvedType "Types" funResolvedType))
                Nothing)
              Nothing
            ])
           Nothing
         ])
      ]
      where
        genExceptionCase :: HS (Field 'ThrowsField) -> GuardedRhs ()
        genExceptionCase Field{..} =
          case fieldTag of
            THROWS_RESOLVED -> GuardedRhs ()
              [ Generator ()
                (PApp () (qualSym "Prelude" "Just")
                 [ PAsPat () (textToName "_e") $
                   PRec () (genConstructor (Just "Types") fieldResolvedType) []
                 ])
                 (qvar "Exception" "fromException" `app` var "_ex")
              ] $
              genRespTup genREPLY
                (protocolFun "genStruct" `app`
                  HS.List ()
                  [ genFieldBase fieldResolvedType fieldName fieldId
                     (intLit (0 :: Int))
                   (  var "_e")
                  ])
                (Just (var "_ex", "ClientError"))

        genRespTup int body exc = Tuple () Boxed
          [ intLit int
          , body
          , case exc of
              Just (e,blame) -> qcon "Prelude" "Just" `app`
                Tuple () Boxed [ e, qcon "Thrift" blame ]
              Nothing -> qcon "Prelude" "Nothing"
          ]

        genResp Nothing = []
        genResp (Just (Some t)) =
          [ genFieldBase t "" 0 (intLit (0 :: Int)) (var "_result") ]

genMethodsInfo :: HS Service -> [HS.Decl ()]
genMethodsInfo service =
  [ signature
  , HS.FunBind ()
    [ HS.Match () (textToName methodsInfo') []
        (HS.UnGuardedRhs () infos) Nothing
    ]
  ]
  where
    signature = HS.TypeSig () [textToName methodsInfo'] $
      qualType "Map" "Map"
        `appT` qualType "Text" "Text"
        `appT` HS.TyCon () (qualSym "Thrift" "MethodInfo")

    methodsInfo' = "methodsInfo'"

    superMethodsInfo :: Maybe (HS.Exp ())
    superMethodsInfo = do
      (Name{..}, _) <- supResolvedName <$> serviceSuper service
      return $ qvar (localName resolvedName) methodsInfo'

    instanceMethodsInfo = qvar "Map" "fromList" `app`
      listE (genInfoTuple <$> getServiceFunctions service)

    infos = case superMethodsInfo of
      Nothing -> instanceMethodsInfo
      Just s -> qvar "Map" "union" `app` instanceMethodsInfo `app` s

    genInfoTuple f@Function{..} =
      Tuple () Boxed
        [ stringLit funName
        , genOneMethodInfo f
        ]

    genOneMethodInfo :: HS Function -> HS.Exp ()
    genOneMethodInfo Function{..} =
      qcon "Thrift" "MethodInfo"
        `app` qcon "Thrift" (Text.pack $ show funPriority)
        `app` qcon "Prelude" (Text.pack $ show funIsOneWay)

commandTypeName :: HS Service -> Text
commandTypeName Service{..} = serviceResolvedName <> "Command"

isEmptyService :: HS Service -> Bool
isEmptyService s@Service{..} = null (getServiceFunctions s) && isNothing serviceSuper