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.GenUnion
    ( genUnionDecl
    , genUnionImports
    ) where
  
  import Prelude hiding (exp)
  import Data.Maybe
  import Data.Text (Text)
  import Language.Haskell.Exts.Syntax hiding
    (Name, Type, Annotation, Decl)
  import qualified Data.Set as Set
  import qualified Language.Haskell.Exts.Syntax as HS
  
  import Thrift.Compiler.GenStruct
  import Thrift.Compiler.GenUtils
  import Thrift.Compiler.Plugins.Haskell
  import Thrift.Compiler.Types
  
  -- Generate Datatype -----------------------------------------------------------
  
  genUnionImports :: HS Union -> Set.Set Import
  genUnionImports Union{..} =
    foldr (Set.union . getImports) baseImports unionAlts
    where
      getImports :: HS UnionAlt -> Set.Set Import
      getImports UnionAlt{..} = typeToImport altResolvedType
      baseImports = Set.fromList
                    [ QImport "Prelude" "Prelude"
                    , QImport "Control.DeepSeq" "DeepSeq"
                    , QImport "Control.Exception" "Exception"
                    , QImport "Data.Aeson" "Aeson"
                    , QImport "Data.Aeson.Types" "Aeson"
                    , QImport "Data.Default" "Default"
                    , QImport "Data.Hashable" "Hashable"
                    , QImport "Data.HashMap.Strict" "HashMap"
                    , QImport "Data.List" "List"
                    , QImport "Data.Ord" "Ord"
                    , SymImport "Prelude" [ ".", "<$>", "<*>", ">>=", "==", "++" ]
                    , SymImport "Data.Aeson" [ ".:", ".=" ]
                    , SymImport "Data.Monoid" [ "<>" ]
                    ]
  
  genUnionDecl :: HS Union -> [HS.Decl ()]
  genUnionDecl union@Union{..} =
    -- Union Declaration
    [ DataDecl () (DataType ()) Nothing
      (DHead () $ textToName unionResolvedName)
      -- Data Constructors
      (map genAltDecl unionAlts ++
       case unionHasEmpty of
         NonEmpty -> []
         HasEmpty ->
           [ QualConDecl () Nothing Nothing $
             ConDecl () (textToName unionEmptyName) []
           ])
  
      -- Deriving
      (pure $ deriving_ $ map (IRule () Nothing Nothing . IHCon ()) $
             [ qualSym "Prelude" "Eq"
             , qualSym "Prelude" "Show"
             ] ++
             [ qualSym "Prelude" "Ord" | deriveOrd ])
    -- Aeson Instances
    , genToJSONInst union
    -- ThriftStruct Instance
    , genThriftStruct union
    -- Other Instances
    , genNFData union
    , genDefault union
    , genHashable union
    ] ++
    [ genOrd union | not deriveOrd ]
    where
      deriveOrd =
        all (\UnionAlt{..} -> isNothing $ mkOrd altResolvedType) unionAlts
  
  genAltDecl :: HS UnionAlt -> QualConDecl ()
  genAltDecl UnionAlt{..} =
    QualConDecl () Nothing Nothing $
    ConDecl () (textToName altResolvedName) [ genType altResolvedType ]
  
  -- Generate Aeson Instances ----------------------------------------------------
  
  genToJSONInst :: HS Union -> HS.Decl ()
  genToJSONInst Union{..} =
    InstDecl () Nothing
      (IRule () Nothing Nothing $
       IHApp ()
         (IHCon () $ qualSym "Aeson" "ToJSON")
         (TyCon () $ unqualSym unionResolvedName))
      (Just [ InsDecl () $ FunBind () $
              map genToJSONAlt unionAlts ++
              case unionHasEmpty of
                NonEmpty -> []
                HasEmpty -> [ genToJSONEmpty unionEmptyName ]
            ])
  
  genToJSONAlt :: HS UnionAlt -> Match ()
  genToJSONAlt UnionAlt{..} =
    Match () (textToName "toJSON")
    [ PApp () (unqualSym altResolvedName) [ pvar arg ] ]
    (UnGuardedRhs () $
       qvar "Aeson" "object" `app`
       HS.List ()
       [ infixApp ".=" (stringLit altName) $
         case fixToJSONValue altResolvedType of
           Nothing -> var arg
           Just f  -> f `app` var arg
       ])
    Nothing
    where
      arg = "__" <> altName
  
  genToJSONEmpty :: Text -> Match ()
  genToJSONEmpty uname =
    Match () (textToName "toJSON")
    [ PApp () (unqualSym uname) [] ]
    (UnGuardedRhs () $
       qvar "Aeson" "object" `app` HS.List () [])
    Nothing
  
  -- Generate ThriftStruct Instance ----------------------------------------------
  
  genThriftStruct :: HS Union -> HS.Decl ()
  genThriftStruct union@Union{..} =
    InstDecl () Nothing
      (IRule () Nothing Nothing $
       IHApp ()
         (IHCon () $ qualSym "Thrift" "ThriftStruct")
         (TyCon () $ unqualSym unionResolvedName))
      (Just $ map (InsDecl ())
       [ genBuilder union
       , genParser union
       ])
  
  genBuilder :: HS Union -> HS.Decl ()
  genBuilder Union{..} =
    FunBind () $
    (flip map unionAlts $ \UnionAlt{..} ->
      let arg = "__" <> altName in
      mkMatch (PApp () (unqualSym altResolvedName) [ pvar arg ])
      [ genFieldBase altResolvedType altName altId lastId (var arg) ]) ++
    case unionHasEmpty of
      NonEmpty -> []
      HasEmpty -> [ mkMatch (PApp () (unqualSym unionEmptyName) []) [] ]
    where
      lastId = intLit (0 :: Int)
      mkMatch pat fields =
        Match () (textToName "buildStruct")
        [ pvar "_proxy", pat ]
        (UnGuardedRhs () $
         protocolFun "genStruct" `app`
         HS.List () fields)
        Nothing
  
  genParser :: HS Union -> HS.Decl ()
  genParser Union{..} =
    FunBind ()
    [ Match ()
      (textToName "parseStruct")
      [ pvar "_proxy" ]
      (UnGuardedRhs () $ Do ()
       [ Generator () (pvar "_fieldBegin") $
         protocolFun "parseFieldBegin" `app` lastId `app` var "_idMap"
       , Qualifier () $
         Case () (var "_fieldBegin")
         [ Alt ()
           (PApp () (qualSym "Thrift" "FieldBegin")
            [ pvar "_type", pvar "_id", pvar "_bool" ])
           (UnGuardedRhs () $ Do ()
            [ Qualifier () $ Case () (var "_id") $
              map genParseValue unionAlts ++
              [ Alt () (PWildCard ())
                (UnGuardedRhs () $ case unionHasEmpty of
                  NonEmpty ->
                    qvar "Prelude" "fail" `app`
                    (infixApp "++"
                     (stringLit
                      ("unrecognized alternative for union '" <>
                       unionName <>
                       "': "))
                     (qvar "Prelude" "show" `app` var "_id"))
                  HasEmpty -> Do ()
                    [ Qualifier () $
                        protocolFun "parseSkip" `app`
                        var "_type" `app`
                        qcon "Prelude" "Nothing"
                    , Qualifier () $ protocolFun "parseStop"
                    , Qualifier () $
                        qvar "Prelude" "return" `app` con unionEmptyName
                    ])
                Nothing
              ]
            ])
           Nothing
         , Alt () (PApp () (qualSym "Thrift" "FieldEnd") [])
           (UnGuardedRhs () $
            case unionHasEmpty of
              NonEmpty -> qvar "Prelude" "fail" `app`
                          stringLit ("union '" <> unionName <> "' is empty")
              HasEmpty -> qvar "Prelude" "return" `app` con unionEmptyName)
           Nothing
         ]
       ])
      (Just $ BDecls ()
       [ PatBind () (pvar "_idMap")
         (UnGuardedRhs () $
          qvar "HashMap" "fromList" `app`
          HS.List ()
          (map (\UnionAlt{..} ->
           Tuple () Boxed [ stringLit altName, intLit altId ])
           unionAlts))
         Nothing
       ])
    ]
    where
      lastId = intLit (0 :: Int)
  
  genParseValue :: HS UnionAlt -> Alt ()
  genParseValue UnionAlt{..} =
    Alt ()
    (PLit ()
     (if altId < 0 then Negative () else Signless ())
     (Int () (abs $ fromIntegral altId) (show altId)))
    (GuardedRhss ()
     -- check that the parsed type is correct
     [ GuardedRhs ()
       [ Qualifier () $
         infixApp "==" (var "_type") (genThriftType altResolvedType)
       ] $ Do ()
       [ Generator () (pvar "_val") $ genParseType P_FieldMode altResolvedType
       , Qualifier () $ protocolFun "parseStop"
       , Qualifier () $
         qvar "Prelude" "return" `app`
         (con altResolvedName `app` var "_val")
       ]
     ])
    Nothing
  
  -- Generate NFData Instance ----------------------------------------------------
  
  genNFData :: HS Union -> HS.Decl ()
  genNFData Union{..} =
    InstDecl () Nothing
      (IRule () Nothing Nothing $
       IHApp ()
         (IHCon () $ qualSym "DeepSeq" "NFData")
         (TyCon () $ unqualSym unionResolvedName))
      (Just
       [ InsDecl () $ FunBind () $
         (flip map unionAlts $ \UnionAlt{..} ->
           Match () (textToName "rnf")
           [ PApp () (unqualSym altResolvedName) [ pvar ("__" <> altName) ] ]
           (UnGuardedRhs () $ qvar "DeepSeq" "rnf" `app` var ("__" <> altName))
           Nothing) ++
         case unionHasEmpty of
           NonEmpty -> []
           HasEmpty ->
             [ Match () (textToName "rnf")
               [ PApp () (unqualSym unionEmptyName) [] ]
               (UnGuardedRhs () $ unit_con ())
               Nothing
             ]
       ])
  
  -- Generate Default Instance ---------------------------------------------------
  
  genDefault :: HS Union -> HS.Decl ()
  genDefault Union{..} =
    InstDecl () Nothing
      (IRule () Nothing Nothing $
       IHApp ()
         (IHCon () $ qualSym "Default" "Default")
         (TyCon () $ unqualSym unionResolvedName))
      (Just
       [ InsDecl () $ FunBind ()
         [ Match () (textToName "def") []
           (UnGuardedRhs () $
            case (unionHasEmpty, unionAlts) of
              (HasEmpty, _) -> con unionEmptyName
              (NonEmpty, UnionAlt{..} : _) ->
                con altResolvedName `app` typeToDefault altResolvedType
              (NonEmpty, []) ->
                qvar "Exception" "throw" `app`
                (qcon "Thrift" "ProtocolException" `app`
                 stringLit
                 ("def: no default value for empty union '" <> unionResolvedName
                  <> "'")))
           Nothing
         ]
       ])
  
  -- Generate Default Instance ---------------------------------------------------
  
  genHashable :: HS Union -> HS.Decl ()
  genHashable Union{..} =
    InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
     (IHCon () $ qualSym "Hashable" "Hashable")
     (TyCon () $ unqualSym unionResolvedName))
    (Just
     [ InsDecl () $ FunBind () $
       map genHashWithSalt unionAlts ++
       case unionHasEmpty of
         NonEmpty -> []
         HasEmpty -> [ genHashWithSaltEmpty unionEmptyName ]
     ])
  
  genHashWithSalt :: HS UnionAlt -> Match ()
  genHashWithSalt UnionAlt{..} =
    Match () (textToName "hashWithSalt")
    [ pvar "__salt"
    , PApp () (unqualSym altResolvedName)
      [ pvar $ "_" <> altName ]
    ]
    (UnGuardedRhs () $
     qvar "Hashable" "hashWithSalt" `app`
     var "__salt" `app`
     (qvar "Hashable" "hashWithSalt" `app`
      intLit altId `app`
      transformValue mkHashable Default altResolvedType (var $ "_" <> altName)))
    Nothing
  
  genHashWithSaltEmpty :: Text -> Match ()
  genHashWithSaltEmpty uname =
    Match () (textToName "hashWithSalt")
    [ pvar "__salt"
    , PApp () (unqualSym uname) []
    ]
    (UnGuardedRhs () $
     qvar "Hashable" "hashWithSalt" `app`
     var "__salt" `app`
     ExpTypeSig () (intLit (0 ::Int)) (qualType "Prelude" "Int"))
    Nothing
  
  -- Generate Default Instance ---------------------------------------------------
  
  genOrd :: HS Union -> HS.Decl ()
  genOrd Union{..} =
    InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
     (IHCon () $ qualSym "Ord" "Ord")
     (TyCon () $ unqualSym unionName))
    (Just
     [ InsDecl () $ FunBind () $
       concatMap genCompare unionAlts ++
       case unionHasEmpty of
         NonEmpty -> []
         HasEmpty -> genCompareEmpty unionEmptyName
     ])
  
  genCompare :: HS UnionAlt -> [Match ()]
  genCompare UnionAlt{..} =
    [ Match () (textToName "compare")
      [ PApp () (unqualSym altResolvedName) [pvar "__a"]
      , PApp () (unqualSym altResolvedName) [pvar "__b"]
      ]
      (UnGuardedRhs () $
       qvar "Ord" "compare" `app`
       transform (var "__a") `app`
       transform (var "__b"))
      Nothing
    , Match () (textToName "compare")
      [ PApp () (unqualSym altResolvedName) [PWildCard ()]
      , PWildCard ()
      ]
      (UnGuardedRhs () $ qcon "Ord" "LT")
      Nothing
    ]
    where
      transform = transformValue mkOrd Default altResolvedType
  
  genCompareEmpty :: Text -> [Match ()]
  genCompareEmpty uname =
    [ Match () (textToName "compare")
      [ PApp () (unqualSym uname) []
      , PApp () (unqualSym uname) []
      ]
      (UnGuardedRhs () $ qcon "Ord" "EQ")
      Nothing
    , Match () (textToName "compare")
      [ PApp () (unqualSym uname) []
      , PWildCard ()
      ]
      (UnGuardedRhs () $ qcon "Ord" "GT")
      Nothing
    ]
  