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.GenJSON
    ( genJSON
    , writeJSON
    , getAstPath
    ) where
  
  import Prelude hiding (Enum)
  import Data.Aeson
  import Data.Aeson.Encode.Pretty
  import Data.ByteString.Builder
  import qualified Data.ByteString.Lazy as LBS
  import qualified Data.HashMap.Strict as HashMap
  import Data.Proxy
  import Data.Some
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.Lazy as Lazy
  import qualified Data.Text.Lazy.Encoding as Lazy
  import GHC.TypeLits
  import System.Directory
  import System.FilePath
  
  import Util.Aeson
  
  import Thrift.Compiler.Typechecker
  import Thrift.Compiler.Types as Types
  import Thrift.Compiler.Options as Options
  import Thrift.Compiler.Plugin
  
  writeJSON
    :: Typecheckable l
    => Program l a         -- ^ Top level program ti generate
    -> Maybe [Program l a] -- ^ Dependencies if using recursive mode
    -> IO FilePath
  writeJSON prog deps = do
    createDirectoryIfMissing True dir
    LBS.writeFile path $ prettyJSON $ genJSON prog deps
    return path
    where
      path = dir </> file
      (dir, file) = getAstPath prog
      prettyJSON = encodePretty' defConfig { confCompare = compare }
  
  getAstPath :: Program l a -> (FilePath, FilePath)
  getAstPath Program{..} = (progOutPath, Text.unpack progName ++ ".ast")
  
  genJSON :: Typecheckable l => Program l a -> Maybe [Program l a] -> Value
  genJSON prog Nothing = Object $ genJSONProg prog
  genJSON prog (Just deps) = toJSON $ genJSONProg prog : map genJSONProg deps
  
  genJSONProg :: Typecheckable l => Program l a -> Object
  genJSONProg Program{..} = objectFromList
    [ "name"     .= progHSName
    , "path"     .= progPath
    , "includes" .= map Types.progPath progIncludes
    , "typedefs" .= map genTypedef dTdefs
    , "enums"    .= map genEnum dEnums
    , "consts"   .= map genConst dConsts
    , "structs"  .= map genStruct dStructs
    , "unions"   .= map genUnion dUnions
    , "services" .= map genService dServs
    , "options"  .= genOptions (options progEnv)
    ]
    where
      Decls{..} = partitionDecls progDecls
  
  -- Typedefs --------------------------------------------------------------------
  
  genTypedef :: Typecheckable l => Typedef 'Resolved l a -> Object
  genTypedef Typedef{..} = objectFromList
    [ "name" .= tdResolvedName
    , "type" .= genType tdResolvedType
    , "newtype" .= case tdTag of { IsNewtype -> True ; IsTypedef -> False }
    ]
  
  -- Enums -----------------------------------------------------------------------
  
  genEnum :: Typecheckable l => Enum 'Resolved l a -> Object
  genEnum Enum{..} = objectFromList
    [ "name"      .= enumResolvedName
    , "constants" .= map genEnumConst enumConstants
    , "flavour" .= case enumFlavour of
        SumTypeEnum{} -> "sum_type" :: Text
        PseudoEnum{} -> "pseudo"
    ]
  
  genEnumConst :: EnumValue 'Resolved l a -> Object
  genEnumConst EnumValue{..} = objectFromList
    [ "name"  .= evResolvedName
    , "value" .= evValue
    ]
  
  -- Constants -------------------------------------------------------------------
  
  genConst :: Typecheckable l => Const 'Resolved l a -> Object
  genConst Const{..} = objectFromList
    [ "name"  .= constResolvedName
    , "type"  .= genType constResolvedType
    , "value" .= genConstVal constResolvedType constResolvedVal
    ]
  
  -- Structs, Exceptions, and Unions ---------------------------------------------
  
  genStruct :: Typecheckable l => Struct 'Resolved l a -> Object
  genStruct Struct{..} = objectFromList
    [ "name" .= structResolvedName
    , "struct_type" .= case structType of
        StructTy    -> "STRUCT" :: Text
        ExceptionTy -> "EXCEPTION"
    , "fields" .= map genField structMembers
    ]
  
  genField :: Typecheckable l => Field u 'Resolved l a -> Object
  genField Field{..} = objectFromList $
    [ "name"  .= fieldResolvedName
    , "id"    .= fieldId
    , "type"  .= genType fieldResolvedType
    ] ++
    (case fieldResolvedVal of
       Nothing -> []
       Just val -> [ "default_value" .= genConstVal fieldResolvedType val ]) ++
    (case fieldTag of
       STRUCT_FIELD -> [ "requiredness" .=
                         case fieldRequiredness of
                           Default    -> "default" :: Text
                           Required{} -> "required"
                           Optional{} -> "optional"
                       ]
       _ -> [])
  
  genUnion :: Typecheckable l => Union 'Resolved l a -> Object
  genUnion Union{..} = objectFromList
    [ "name"   .= unionResolvedName
    , "fields" .= map genAlt unionAlts
    ]
  
  genAlt :: Typecheckable l => UnionAlt 'Resolved l a -> Object
  genAlt UnionAlt{..} = objectFromList
    [ "name" .= altResolvedName
    , "id"   .= altId
    , "type" .= genType altResolvedType
    ]
  
  -- Services and Functions ------------------------------------------------------
  
  genService :: Typecheckable l => Service 'Resolved l a -> Object
  genService s@Service{..} = objectFromList $
    [ "name"      .= serviceResolvedName
    , "functions" .= map genFunction (getServiceFunctions s)
    ] ++
    (case serviceSuper of
       Nothing -> []
       Just Super{..} -> [ "super" .= genName (fst supResolvedName) ])
  
  genFunction :: Typecheckable l => Function 'Resolved l a -> Object
  genFunction Function{..} = objectFromList
    [ "name" .= funResolvedName
    , "return_type" .= case funResolvedType of
        Nothing -> simpleType "void"
        Just ty -> withSome ty genType
    , "args"   .= map genField funArgs
    , "throws" .= map genField funExceptions
    , "oneway" .= funIsOneWay
    ]
  
  -- Types and Constants ---------------------------------------------------------
  
  genType :: Typecheckable l => Type l t -> Object
  
  -- Base Types
  genType I8  = simpleType "byte"
  genType I16 = simpleType "i16"
  genType I32 = simpleType "i32"
  genType I64 = simpleType "i64"
  genType TFloat  = simpleType "float"
  genType TDouble = simpleType "double"
  genType TBool   = simpleType "bool"
  genType TText   = simpleType "string"
  genType TBytes  = simpleType "binary"
  
  -- Collections
  genType (TSet u)       = collectionType "set" u
  genType (THashSet u)   = collectionType "hash_set" u
  genType (TList u)      = collectionType "list" u
  genType (TMap k v)     = mapType "map" k v
  genType (THashMap k v) = mapType "hash_map" k v
  
  -- Named Types
  genType (TStruct name _loc)    = namedType "struct" name
  genType (TException name _loc) = namedType "exception" name
  genType (TUnion name _loc)     = namedType "union" name
  genType (TEnum name _loc _) = namedType "enum" name
  genType (TTypedef name ty _loc) = objectFromList
    [ "type" .= ("typedef" :: Text)
    , "name" .= genName name
    , "inner_type" .= genType ty
    ]
  genType (TNewtype name ty _loc) = objectFromList
    [ "type" .= ("newtype" :: Text)
    , "name" .= genName name
    , "inner_type" .= genType ty
    ]
  genType (TSpecial ty) = case backTranslateType ty of
    (Some u, tag) -> genType u <> objectFromList [ "special" .= tag ]
  
  simpleType :: Text -> Object
  simpleType tyName = objectFromList ["type" .= (String tyName)]
  
  collectionType :: Typecheckable l => Text -> Type l t -> Object
  collectionType tyName u = objectFromList
    [ "type" .= tyName
    , "inner_type" .= genType u
    ]
  
  mapType :: Typecheckable l => Text -> Type l u -> Type l v -> Object
  mapType tyName k v = objectFromList
    [ "type" .= tyName
    , "key_type" .= genType k
    , "val_type" .= genType v
    ]
  
  namedType :: Text -> Name -> Object
  namedType tyName name = objectFromList
    [ "type" .= tyName
    , "name" .= genName name
    ]
  
  genName :: Name -> Object
  genName Name{..} = objectFromList $
    [ "name" .= localName resolvedName ] ++
    [ "src" .= m | QName m _ <- [sourceName] ]
  
  genConstVal :: Typecheckable l => Type l t -> TypedConst l t -> Object
  genConstVal ty (Literal x) =
    objectFromList [ "literal" .= genLiteral ty x ]
  genConstVal _ (Identifier name _ _loc) =
    objectFromList [ "named_constant" .= genName name ]
  genConstVal _ (WeirdEnumToInt _ name _ _loc) =
    objectFromList [ "named_constant_enumToInt" .= genName name ]
  
  genLiteral :: Typecheckable l => Type l t -> t -> Object
  
  -- Base Types
  genLiteral ty@I8  n = simpleLiteral ty n
  genLiteral ty@I16 n = simpleLiteral ty n
  genLiteral ty@I32 n = simpleLiteral ty n
  genLiteral ty@I64 n =
    -- We need to include the string representation because JSON does not support
    -- 64 bit integers
    simpleLiteral ty n <> objectFromList [ "string" .= show n ]
  genLiteral ty@TFloat n =
    simpleLiteral ty n <>
    objectFromList [ "binary" .= toLazyText (floatHexFixed n) ]
  genLiteral ty@TDouble n =
    simpleLiteral ty n <>
    objectFromList [ "binary" .= toLazyText (doubleHexFixed n) ]
  genLiteral ty@TBool b = simpleLiteral ty b
  genLiteral ty@TText s = simpleLiteral ty s
  -- Serialized as a hexidecimal string
  genLiteral ty@TBytes s = simpleLiteral ty $ toLazyText $ byteStringHex s
  
  -- Collections
  genLiteral (TSet u)       (Set xs)     = listLiteral "set" u xs
  genLiteral (THashSet u)   (HashSet xs) = listLiteral "hash_set" u xs
  genLiteral (TList u)      (List xs)    = listLiteral "list" u xs
  genLiteral (TMap k v)     (Map xs)     = mapLiteral "map" k v xs
  genLiteral (THashMap k v) (HashMap xs) = mapLiteral "hash_map" k v xs
  
  -- Named Types
  genLiteral TStruct{} (Some sval) = genStructVal sval
  genLiteral TException{} (Some (EV sval)) = genStructVal sval
  genLiteral TUnion{} (Some uval) = genUnionVal uval
  genLiteral TEnum{} (EnumVal name _loc) = objectFromList
    [ "type"  .= ("enum" :: Text)
    , "value" .= genName name
    ]
  genLiteral (TTypedef _ ty _loc) x = genLiteral ty x
  genLiteral (TNewtype _ ty _loc) (New x) = objectFromList
    [ "type"  .= ("newtype" :: Text)
    , "value" .= genLiteral ty x
    ]
  genLiteral st@(TSpecial ty) val = case backTranslateLiteral ty val of
    ThisLit u x -> objectFromList
      [ "type"  .= genType st
      , "value" .= genLiteral u x
      ]
  
  simpleLiteral :: (Typecheckable l, ToJSON a) => Type l t -> a -> Object
  simpleLiteral ty x = genType ty <> objectFromList [ "value" .= x ]
  
  listLiteral :: Typecheckable l => Text -> Type l t -> [TypedConst l t] -> Object
  listLiteral tyName ty xs = objectFromList
    [ "type"  .= tyName
    , "value" .= map (genConstVal ty) xs
    ]
  
  mapLiteral
    :: Typecheckable l
    => Text
    -> Type l k
    -> Type l v
    -> [(TypedConst l k, TypedConst l v)]
    -> Object
  mapLiteral tyName kt vt xs = objectFromList
    [ "type"  .= tyName
    , "value" .= map (genPair kt vt) xs
    ]
  
  genPair
    :: Typecheckable l
    => Type l k
    -> Type l v
    -> (TypedConst l k, TypedConst l v)
    -> Value
  genPair kt vt (k, v) = Object $ objectFromList
    [ "key" .= genConstVal kt k
    , "val" .= genConstVal vt v
    ]
  
  genStructVal :: Typecheckable l => StructVal l s -> Object
  genStructVal s = objectFromList
    [ "type"  .= ("struct" :: Text)
    , "value" .= genStructFields s
    ]
  
  genStructFields :: Typecheckable l => StructVal l s -> [Object]
  genStructFields Empty = []
  genStructFields (ConsVal proxy ty c s) =
    genFieldVal proxy ty c : genStructFields s
  genStructFields (ConsDefault proxy ty s) = objectFromList
    [ "field_name"  .= symbolVal proxy
    , "field_type"  .= genType ty
    , "field_value" .=
      HashMap.singleton ("default" :: Text) Null
    ] :
    genStructFields s
  genStructFields (ConsJust proxy ty c s) =
    genFieldVal proxy ty c : genStructFields s
  genStructFields (ConsNothing proxy s) = objectFromList
    [ "field_name"  .= symbolVal proxy
    , "field_value" .= Null
    ] :
    genStructFields s
  
  genUnionVal :: Typecheckable l => UnionVal l s -> Object
  genUnionVal (UnionVal proxy ty c _) = objectFromList
    -- This isn't technically a thrift type, but we'll use it anyway
    [ "type"  .= ("union" :: Text)
    , "value" .= genFieldVal proxy ty c
    ]
  
  genFieldVal
    :: (Typecheckable l, KnownSymbol s)
    => Proxy s
    -> Type l t
    -> TypedConst l t
    -> Object
  genFieldVal proxy ty c = objectFromList
    [ "field_name"  .= symbolVal proxy
    , "field_type"  .= genType ty
    , "field_value" .= genConstVal ty c
    ]
  
  toLazyText :: Builder -> Lazy.Text
  toLazyText = Lazy.decodeUtf8 . toLazyByteString
  
  -- Options ---------------------------------------------------------------------
  
  genOptions :: Options.Options l -> Object
  genOptions Options.Options{..} = objectFromList
    [ "path" .= optsPath
    , "out_path" .= optsOutPath
    , "include_path" .= optsIncludePath
    , "recursive" .= optsRecursive
    , "genfiles" .= optsThriftMade
    ]
  