content: {-
    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.
  -}
  
  module Thrift.Compiler.GenFunction
    ( genFunctionDecls
    , genFunctionImports
    ) where
  
  import Data.Maybe
  import Data.Set (union)
  import Data.Some
  import Data.Text (Text)
  import Language.Haskell.Exts.Syntax hiding (Type)
  import qualified Data.Set as Set
  import qualified Language.Haskell.Exts.Syntax as HS
  
  import Thrift.Compiler.Types
  import Thrift.Compiler.GenStruct
  import Thrift.Compiler.GenUtils
  import Thrift.Compiler.Plugins.Haskell
  import Util.HSE
  
  -- Generate Imports and Decls --------------------------------------------------
  
  genFunctionImports :: HS Function -> Set.Set Import
  genFunctionImports Function{..} =
    foldr (union . getImports)
    (foldr (union . getImports) baseImports funExceptions)
    funArgs `union`
    maybe Set.empty (`withSome` typeToImport) funResolvedType
    where
      getImports :: HS (Field u) -> Set.Set Import
      getImports Field{..} = typeToImport fieldResolvedType
      baseImports = Set.fromList
        [ QImport "Prelude" "Prelude"
        , QImport "Control.Arrow" "Arrow"
        , QImport "Control.Exception" "Exception"
        , QImport "Control.Concurrent" "Concurrent"
        , QImport "Control.Monad" "Monad"
        , QImport "Control.Monad.Trans.Class" "Trans"
        , QImport "Control.Monad.Trans.Reader" "Reader"
        , QImport "Data.ByteString.Builder" "ByteString"
        , QImport "Data.ByteString.Lazy" "LBS"
        , QImport "Data.Int" "Int"
        , QImport "Data.HashMap.Strict" "HashMap"
        , QImport "Data.List" "List"
        , QImport "Data.Proxy" "Proxy"
        , QImport "Thrift.Binary.Parser" "Parser"
        , QImport "Thrift.Protocol.ApplicationException.Types" "Thrift"
        , SymImport "Prelude" [ "==", "=<<", ">>=", "<$>", "." ]
        , SymImport "Data.Monoid" [ "<>" ]
        ]
  
  genFunctionDecls
    :: HS Service -> HS Function -> [HS.Decl ()]
  genFunctionDecls service func = concatMap ($ func)
    [ genFunctionThrift service
    , genFunctionIO service
    , genFunctionSend service
    , genFunctionRecv
    , genBuildCall
    , genParseResponse
    ]
  
  -- Generate the Function Call In the Thrift Monad ------------------------------
  
  genFunctionThrift
    :: HS Service -> HS Function -> [HS.Decl ()]
  genFunctionThrift Service{..} Function{..} =
    -- Type Signature
    [ TypeSig () [ name ] $
      TyForall () Nothing (commonCtx serviceResolvedName) $
      genArgTypes funArgs
      (qualType "Thrift" "ThriftM" `appT`
       tvar "p" `appT`
       tvar "c" `appT`
       tvar "s" `appT`
       maybe (unit_tycon ()) (`withSome` genType) funResolvedType)
    -- Function Body
    , FunBind ()
      [ Match () name
        (map (pvar . ("__field__" <>) . fieldName) funArgs)
        (UnGuardedRhs () $ Do ()
         [ Generator ()
           (PApp () (qualSym "Thrift" "ThriftEnv")
            (map pvar [ "_proxy", "_channel", "_opts", "_counter" ])) $
           qvar "Reader" "ask"
         , Qualifier () $
           qvar "Trans" "lift" `app`
           genApplyArgs funArgs
           (var (funResolvedName <> "IO") `app`
            var "_proxy" `app`
            var "_channel" `app`
            var "_counter" `app`
            var "_opts")
         ])
        Nothing
      ]
    ]
    where
      name = textToName funResolvedName
  
  -- Generate the Function Call in the IO Monad ----------------------------------
  
  genFunctionIO
    :: HS Service -> HS Function -> [HS.Decl ()]
  genFunctionIO Service{..} fun@Function{..} =
    -- Type Signature
    [ TypeSig () [ textToName (funResolvedName <> "IO") ] $
      TyForall () Nothing (commonCtx serviceResolvedName) $
      TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
      TyFun () (tvar "c" `appT` tvar "s") $
      TyFun () (qualType "Thrift" "Counter") $
      TyFun () (qualType "Thrift" "RpcOptions") $
      genArgTypes funArgs
      (qualType "Prelude" "IO" `appT`
       maybe (unit_tycon ()) (`withSome` genType) funResolvedType)
    -- Function Body
    , FunBind ()
      [ Match () (textToName $ funResolvedName <> "IO")
        (map pvar $
         [ "_proxy", "_channel", "_counter", "_opts" ] ++
         map (("__field__" <>) . fieldName) funArgs)
        (UnGuardedRhs () $ genFunctionIOBody fun)
        Nothing
      ]
    ]
  
  genFunctionIOBody :: HS Function -> Exp ()
  genFunctionIOBody fun@Function{..}
    | funIsOneWay = Do () $
      [ Generator () (pvar "_handle") $ qvar "Concurrent" "newEmptyMVar"
      , Qualifier () $
        genApplyArgs funArgs $
        var (sendName fun) `app`
        var "_proxy" `app`
        var "_channel" `app`
        var "_counter" `app`
        (qvar "Concurrent" "putMVar" `app` var "_handle") `app`
        var "_opts"
      , Qualifier () $ infixApp ">>="
        (qvar "Concurrent" "takeMVar" `app` var "_handle")
        (qvar "Prelude" "maybe" `app`
         (qvar "Prelude" "return" `app` unit_con ()) `app`
         (qvar "Exception" "throw"))
      ]
    | otherwise = Do () $
      [ Generator ()
        (PTuple () Boxed [ pvar "_handle", pvar "_sendCob", pvar "_recvCob" ]) $
        qvar "Thrift" "mkCallbacks" `app`
        (var (recvName fun) `app` var "_proxy")
      , Qualifier () $
        genApplyArgs funArgs $
        var (sendName fun) `app`
        var "_proxy" `app`
        var "_channel" `app`
        var "_counter" `app`
        var "_sendCob" `app`
        var "_recvCob" `app`
        var "_opts"
      , Qualifier () $ qvar "Thrift" "wait" `app` var "_handle"
      ]
  
  -- Generate Send Call ----------------------------------------------------------
  
  genFunctionSend
    :: HS Service -> HS Function -> [HS.Decl ()]
  genFunctionSend Service{..} fun@Function{..} =
    -- Type Signature
    [ TypeSig () [ textToName (sendName fun) ] $
      TyForall () Nothing (commonCtx serviceResolvedName) $
      TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
      TyFun () (tvar "c" `appT` tvar "s") $
      TyFun () (qualType "Thrift" "Counter") $
      TyFun () (qualType "Thrift" "SendCallback") $
      (if funIsOneWay
       then id
       else TyFun () (qualType "Thrift" "RecvCallback")) $
      TyFun () (qualType "Thrift" "RpcOptions") $
      genArgTypes funArgs $
      qualType "Prelude" "IO" `appT` unit_tycon ()
    -- Function Body
    , FunBind ()
      [ Match () (textToName $ sendName fun)
        (map pvar $
         [ "_proxy", "_channel", "_counter", "_sendCob" ] ++
         (if funIsOneWay
          then []
          else [ "_recvCob" ]) ++
         "_rpcOpts" :
         map (("__field__" <>) . fieldName) funArgs)
        (UnGuardedRhs () $ Do ()
         [ Generator () (pvar "_seqNum") (var "_counter")
         , LetStmt () $ BDecls ()
           [ PatBind () (pvar "_callMsg")
             (UnGuardedRhs () $
              qvar "LBS" "toStrict" `app`
              (qvar "ByteString" "toLazyByteString" `app`
               genApplyArgs funArgs
               (var (buildName fun) `app`
                var "_proxy" `app`
                var "_seqNum")))
             Nothing
           ]
         , Qualifier () $
           if funIsOneWay
           then
             qvar "Thrift" "sendOnewayRequest" `app`
             var "_channel" `app`
             (qcon "Thrift" "Request" `app`
              var "_callMsg" `app`
              rpcAst) `app`
             var "_sendCob"
           else
             qvar "Thrift" "sendRequest" `app`
             var "_channel" `app`
             (qcon "Thrift" "Request" `app`
              var "_callMsg" `app`
              rpcAst) `app`
             var "_sendCob" `app`
             var "_recvCob"
         ])
        Nothing
      ]
    ]
    where
      rpcAst = App ()
        (App ()
          (Var () (Qual () (ModuleName () "Thrift") (Ident () "setRpcPriority")))
          (Var () (UnQual () (Ident () "_rpcOpts")))
        )
        (Var () (Qual () (ModuleName () "Thrift") (Ident () priorityString)))
      priorityString = show $ case funPriority of
        NPriorities -> NormalPriority
        _           -> funPriority
  
  sendName :: HS Function -> Text
  sendName Function{..} = "send_" <> funResolvedName
  
  -- Generate Recv Call ----------------------------------------------------------
  
  genFunctionRecv :: HS Function -> [HS.Decl ()]
  genFunctionRecv fun@Function{..}
    | funIsOneWay = []
    | otherwise =
    -- Type Signature
    [ TypeSig () [ textToName (recvName fun) ] $
      TyForall () Nothing
      (Just $ CxTuple ()
       [ classA () (qualSym "Thrift" "Protocol") [ tvar "p" ]
       ]) $
      TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
      TyFun () (qualType "Thrift" "Response") $
      qualType "Prelude" "Either" `appT`
      qualType "Exception" "SomeException" `appT`
      maybe (unit_tycon ()) (`withSome` genType) funResolvedType
  
    -- Function Body
    , FunBind ()
      [ Match () (textToName $ recvName fun)
        [ pvar "_proxy"
        , PApp () (qualSym "Thrift" "Response")
          [ pvar "_response", PWildCard () ]
        ]
        (UnGuardedRhs () $
         app (qvar "Monad" "join") $
         qvar "Arrow" "left" `app`
         (qvar "Exception" "SomeException" `compose`
          qcon "Thrift" "ProtocolException") `app`
         (qvar "Parser" "parse" `app`
          (var (parseName fun) `app` var "_proxy") `app`
          var "_response"))
         Nothing
      ]
    ]
  
  recvName :: HS Function -> Text
  recvName Function{..} = "recv_" <> funResolvedName
  
  -- Generate the Function Call Message Builder ----------------------------------
  
  genBuildCall :: HS Function -> [HS.Decl ()]
  genBuildCall fun@Function{..} =
    -- Type Signature
    [ TypeSig () [ textToName (buildName fun) ] $
      TyForall () Nothing
      (Just $ CxSingle () $
       classA () (qualSym "Thrift" "Protocol") [ tvar "p" ]) $
      TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
      TyFun () (qualType "Int" "Int32") $
      genArgTypes funArgs (qualType "ByteString" "Builder")
    -- Function Body
    , FunBind ()
      [ Match () (textToName (buildName fun))
        (pvar "_proxy" :
         pvar "_seqNum" :
         map (pvar . ("__field__" <>) . fieldName) funArgs)
        (UnGuardedRhs () $
         infixApp "<>"
         (infixApp "<>"
          (protocolFun "genMsgBegin" `app`
           stringLit funName `app`
           intLit genCALL `app`
           var "_seqNum")
          (genBuildFields funArgs))
         (protocolFun "genMsgEnd"))
        Nothing
      ]
    ]
  
  buildName :: HS Function -> Text
  buildName Function{..} = "_build_" <> funResolvedName
  
  -- Generate the Response Parser ------------------------------------------------
  
  genParseResponse :: HS Function -> [HS.Decl ()]
  genParseResponse func@Function{..}
    | funIsOneWay = []
    | otherwise =
    -- Type Signature
    [ TypeSig () [ textToName (parseName func) ] $
      TyForall () Nothing
      (Just $ CxSingle () $
       classA () (qualSym "Thrift" "Protocol") [ tvar "p" ]) $
      TyFun () (qualType "Proxy" "Proxy" `appT` tvar "p") $
      qualType "Parser" "Parser" `appT`
      (qualType "Prelude" "Either" `appT`
       qualType "Exception" "SomeException" `appT`
       case funResolvedType of
         Nothing -> unit_tycon ()
         Just (Some ty) -> genType ty)
    -- Function Body
    , FunBind ()
      [ Match () (textToName (parseName func)) [ pvar "_proxy" ]
        (UnGuardedRhs () $ Do ()
         -- Parse the beginning of the message
         [ Generator ()
           (PApp () (qualSym "Thrift" "MsgBegin")
            [ pvar "_name", pvar "_msgTy", PWildCard () ]) $
           protocolFun "parseMsgBegin"
         , Generator () (pvar "_result") $
           Case () (var "_msgTy")
           -- Case 1: CALL
           [ Alt () (intP genCALL)
             (UnGuardedRhs () $
              qvar "Prelude" "fail" `app`
              stringLit (funName <> ": expected reply but got function call"))
             Nothing
           -- Case 2: REPLY
           , Alt () (intP genREPLY)
             (GuardedRhss ()
              [ GuardedRhs ()
                [ Qualifier () $ infixApp "==" (var "_name") (stringLit funName) ]
                (genParseReply func)
              , GuardedRhs () [ Qualifier () $ qvar "Prelude" "otherwise" ] $
                qvar "Prelude" "fail" `app`
                stringLit "reply function does not match"
              ])
             Nothing
           -- Case 3: EXCEPTION
           , Alt () (intP genEXCEPTION)
             (UnGuardedRhs () $
              qvar "Prelude" "fmap" `app`
              (qcon "Prelude" "Left" `compose`
               qcon "Exception" "SomeException") `app`
              ExpTypeSig ()
              (protocolFun "parseStruct")
              (qualType "Parser" "Parser" `appT`
               qualType "Thrift" "ApplicationException"))
             Nothing
           -- Case 4: ONEWAY
           , Alt () (intP genONEWAY)
             (UnGuardedRhs () $
              qvar "Prelude" "fail" `app`
              stringLit
              (funName <> ": expected reply but got oneway function call"))
             Nothing
           -- Catch All
           , Alt () (PWildCard ())
             (UnGuardedRhs () $
              qvar "Prelude" "fail" `app`
              stringLit (funName <> ": invalid message type"))
             Nothing
           ]
         , Qualifier () $ protocolFun "parseMsgEnd"
         , Qualifier () $ qvar "Prelude" "return" `app` var "_result"
         ])
        Nothing
      ]
    ]
  
  parseName :: HS Function -> Text
  parseName Function{..} = "_parse_" <> funResolvedName
  
  -- This assumes that the response is basically a union, only one response is
  -- possible.
  genParseReply :: HS Function -> Exp ()
  genParseReply Function{..} = Do ()
    [ LetStmt () $ BDecls ()
      [ PatBind () (pvar "_idMap")
        (UnGuardedRhs () $
         qvar "HashMap" "fromList" `app`
         HS.List ()
         (map (\Field{..} ->
          Tuple () Boxed [ stringLit fieldName, intLit fieldId ])
          fields))
        Nothing
      ]
    , Generator () (pvar "_fieldBegin") $
      protocolFun "parseFieldBegin" `app` intLit (0 :: Int)  `app` var "_idMap"
    , Qualifier () $
      Case () (var "_fieldBegin")
      [ Alt ()
        (PApp () (qualSym "Thrift" "FieldBegin")
         [ pvar "_type", pvar "_id", pvar "_bool" ])
        (UnGuardedRhs () $ Do ()
         [ Qualifier () $ Case () (var "_id") $
           maybeToList ((`withSome` genParseResult) <$> funResolvedType) ++
           map genParseException funExceptions ++
           [ Alt () (PWildCard ())
             (UnGuardedRhs () $
              qvar "Prelude" "fail" `app`
              (qvar "Prelude" "unwords" `app`
               HS.List ()
               [ stringLit "unrecognized exception, type:"
               , qvar "Prelude" "show" `app` var "_type"
               , stringLit "field id:"
               , qvar "Prelude" "show" `app` var "_id"
               ]))
             Nothing
           ]
         ])
        Nothing
      , Alt () (PApp () (qualSym "Thrift" "FieldEnd") [])
        (UnGuardedRhs () $
         case funResolvedType of
          Nothing ->
            qvar "Prelude" "return" `app`
            (qcon "Prelude" "Right" `app` unit_con ())
          Just{} -> qvar "Prelude" "fail" `app` stringLit "no response")
        Nothing
      ]
    ]
    where
      resultField :: HSType t -> HS (Field 'StructField)
      resultField ty = Field
        { fieldId     = 0
        , fieldName   = funName <> "_success"
        , fieldResolvedType   = ty
        , fieldRequiredness = Default
        , fieldTag = STRUCT_FIELD
        -- The following are placeholders since the fields aren't needed
        , fieldResolvedName = ""
        , fieldType = AnnotatedType I8 Nothing (Arity0Loc nlc)
        , fieldResolvedVal = Nothing
        , fieldVal = Nothing
        , fieldLaziness = Lazy
        , fieldLoc = FieldLoc nlc "0" nlc nlc Nothing NoSep
        , fieldAnns = Nothing
        , fieldSAnns = []
        }
      structify
        :: HS (Field 'ThrowsField)
        -> HS (Field 'StructField)
      structify Field{ fieldRequiredness = Default, ..} =
        Field{ fieldTag = STRUCT_FIELD, fieldRequiredness = Default, .. }
      fields = maybe id (`withSome` ((:) . resultField)) funResolvedType $
               map structify funExceptions
  
  genParseResult :: HSType t -> Alt ()
  genParseResult ty =
    Alt () (intP (0 :: Int))
    (GuardedRhss ()
     [ GuardedRhs ()
       [ Qualifier () $ infixApp "==" (var "_type") (genThriftType ty) ] $
       qvar "Prelude" "fmap" `app`
       qcon "Prelude" "Right" `app`
       genParseType P_FieldMode ty
     ])
    Nothing
  
  genParseException :: HS (Field 'ThrowsField) -> Alt ()
  genParseException Field{..} =
    Alt () (intP fieldId)
    (GuardedRhss ()
     [ GuardedRhs ()
       [ Qualifier () $ infixApp "==" (var "_type")
         (genThriftType fieldResolvedType) ] $
       qvar "Prelude" "fmap" `app`
       (qvar "Prelude" "Left" `compose` qvar "Exception" "SomeException") `app`
       ExpTypeSig ()
       (genParseType P_FieldMode fieldResolvedType)
       (qualType "Parser" "Parser" `appT` genType fieldResolvedType)
     ])
    Nothing
  
  -- Helpers ---------------------------------------------------------------------
  
  commonCtx :: Text -> Maybe (HS.Context ())
  commonCtx sname = Just $ CxTuple ()
    -- Protocol Constraint
    [ classA () (qualSym "Thrift" "Protocol") [ tvar "p" ]
    -- ClientChannel Constraint
    , classA () (qualSym "Thrift" "ClientChannel")  [ tvar "c" ]
    -- Service Subtyping Constraint
    , classA ()
      (Qual () (ModuleName () "Thrift") (Symbol () "<:"))
      [ tvar "s"
      , simpleType sname
      ]
    ]
  
  genArgTypes :: [HS (Field 'Argument)] -> HS.Type () -> HS.Type ()
  genArgTypes = flip $ foldr (\Field{..} -> TyFun () $ genType fieldResolvedType)
  
  genApplyArgs :: [HS (Field 'Argument)] -> Exp () -> Exp ()
  genApplyArgs =
    flip $ foldl (\f Field{..} -> f `app` var ("__field__" <> fieldName))
  