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.
  -}
  
  {-# LANGUAGE TypeApplications #-}
  module Thrift.Compiler
    ( run
    , parseAll, parseThriftFile
    , parseAllE, parseThriftFileE
    , typecheckInput
    , typecheckInputE
    ) where
  
  import Control.Concurrent.Async
  import Control.Monad
  import Data.List
  import Data.Typeable
  import Language.Haskell.Exts hiding (parse, Decl, name)
  import System.Exit
  import System.FilePath
  import qualified Data.Map.Strict as Map
  
  -- Backends
  import Thrift.Compiler.GenHaskell
  import Thrift.Compiler.GenJSON
  import Thrift.Compiler.GenJSONLoc
  
  import Thrift.Compiler.Options
  import Thrift.Compiler.OptParse
  import Thrift.Compiler.Parser
  import Thrift.Compiler.Pretty
  import Thrift.Compiler.Typechecker
  import Thrift.Compiler.Typechecker.Monad
  import Thrift.Compiler.Types as Thrift
  
  -- | Run the thrift compiler and return a list of paths to the generated files
  run :: ThriftLanguage l => Options l -> IO [FilePath]
  run opts@Options{..} = do
    -- Parse and Typecheck
    (headModule, deps) <- typecheckInput opts =<< parseAll opts optsPath
  
    if optsLenient && not optsLenientStillGenCode then return []
    else do
      -- Generate Outputs
      genFiles <- getGenerator opts headModule deps
  
      -- Write THIRFT-MADE-GENERATED-FILES
      case optsThriftMade of
        Nothing -> return ()
        Just path -> writeFile path $ unlines genFiles
  
      return genFiles
  
  typecheckInput
    :: ThriftLanguage l
    => Options l
    -> ModuleMap
    -> IO (Program l Thrift.Loc, [Program l Thrift.Loc])
  typecheckInput opts moduleMap =
    case typecheckInputE opts moduleMap of
      Left es -> mapM_ putStrLn es >> exitFailure
      Right ms -> pure ms
  
  typecheckInputE
    :: ThriftLanguage l
    => Options l
    -> ModuleMap
    -> Either [String] (Program l Thrift.Loc, [Program l Thrift.Loc])
  typecheckInputE opts moduleMap =
    case typecheck opts moduleMap of
      Left es ->  Left $ map renderTypeErrorPlain (sortBy orderError es)
      Right ms -> Right ms
  
  getGenerator
    :: ThriftLanguage l
    => Options l
    -> Program l Thrift.Loc
    -> [Program l Thrift.Loc]
    -> IO [FilePath]
  getGenerator opts@Options{..} prog deps = case optsGenMode of
    Lint -> return []
    EmitCode
      | Just (hsopts, hsprogs) <- cast (opts, allProgs) ->
          concat <$> mapConcurrently @[] (writeHsCode hsopts) hsprogs
      | otherwise -> error
          "Code generation is only supported for Haskell. Try using --emit-json"
    EmitJSON WithoutLoc
      | optsSingleOutput -> (:[]) <$> writeJSON prog mdeps
      | otherwise -> mapConcurrently (`writeJSON` Nothing) allProgs
    EmitJSON WithLoc
      | optsSingleOutput -> (:[]) <$> writeJSONLoc prog mdeps
      | otherwise -> mapConcurrently (`writeJSONLoc` Nothing) allProgs
    where
      (allProgs, mdeps)
        | optsRecursive = (prog : deps, Just deps)
        | otherwise = ([prog], Nothing)
  
  -- | Return a parse error or the result
  parseAllE :: Options l -> FilePath -> IO (Either String ModuleMap)
  parseAllE Options{..} = parseItE (Right Map.empty)
    where
      parseItE :: Either String ModuleMap
               -> FilePath -> IO (Either String ModuleMap)
      parseItE err@Left{} _path = pure err -- stop at first error
      parseItE (Right tmap) path
        -- If we already parsed it, don't do it again
        | Map.member path tmap = pure (Right tmap)
        | otherwise = do
            e <- parseThriftFileE optsIncludePath path
            case e of
              Left err -> return (Left err)
              Right file@ThriftFile{..} -> do
                let newMap = Map.insert path file tmap
                foldM parseItE (Right newMap) (getIncludes thriftHeaders)
      getIncludes = foldr getInc []
      getInc HInclude{incType=Include,..} ps = incPath : ps
      getInc _ ps = ps
  
  -- | parse all the things recursively, starting with the input and traversing
  -- the includes
  parseAll :: Options l -> FilePath -> IO ModuleMap
  parseAll options fp = do
    e <- parseAllE options fp
    case e of
      Left err -> putStrLn err >> exitFailure
      Right x -> return x
  
  -- | Return a parse error or the result
  parseThriftFileE
    :: FilePath -> FilePath
    -> IO (Either String (ThriftFile SpliceFile Thrift.Loc))
  parseThriftFileE baseDir path = do
    result <- parse baseDir path
    case result of
      Left err -> return (Left err)
      Right file@ThriftFile{..} -> do
        spliceFile <- parseHsInclude baseDir thriftHeaders
        return (Right file { thriftSplice = spliceFile })
  
  -- | Parse a single thrift file
  parseThriftFile :: FilePath -> FilePath -> IO (ThriftFile SpliceFile Thrift.Loc)
  parseThriftFile baseDir path = do
    e <- parseThriftFileE baseDir path
    case e of
        Left err -> putStrLn err >> exitFailure
        Right x -> return x
  
  -- | Select the last hs_include header (there should be at most one)
  parseHsInclude :: FilePath -> [Header s l a] -> IO SpliceFile
  parseHsInclude baseDir headers = case foldl' getInc Nothing headers of
    Nothing -> return Nothing
    Just path -> Just <$> do
      let mode = defaultParseMode { parseFilename = path }
      fromParseResult <$> parseFileWithMode mode (baseDir </> path)
    where
      getInc _ HInclude{incType=HsInclude,..} = Just incPath
      getInc ns _ = ns
  