content: -- Copyright 2016 Google Inc. All Rights Reserved.
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  --
  -- Code for writing protocol compiler plugins.
  
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE OverloadedLabels #-}
  module Data.ProtoLens.Compiler.Plugin
      ( ProtoFileName
      , ProtoFile(..)
      , analyzeProtoFiles
      , collectEnvFromDeps
      ) where
  
  import Data.List (foldl')
  import qualified Data.Map.Strict as Map
  import Data.Map.Strict (Map, unions, (!))
  import Data.String (fromString)
  import qualified Data.Text as T
  import Data.Text (Text)
  import Lens.Family2
  import Proto.Google.Protobuf.Descriptor (FileDescriptorProto)
  
  import Data.ProtoLens.Compiler.Definitions
  import Data.ProtoLens.Compiler.ModuleName
  
  import GHC.SourceGen (ModuleNameStr, OccNameStr, RdrNameStr)
  
  -- | The filename of an input .proto file.
  type ProtoFileName = Text
  
  data ProtoFile = ProtoFile
      { descriptor :: FileDescriptorProto
      , haskellModule :: ModuleNameStr
      , definitions :: Env OccNameStr
      , services :: [ServiceInfo]
      , exportedEnv :: Env RdrNameStr
      , publicImports :: [ModuleNameStr]
      }
  
  -- Given a list of FileDescriptorProtos, collect information about each file
  -- into a map of 'ProtoFile's keyed by 'ProtoFileName'.
  analyzeProtoFiles :: [FileDescriptorProto] -> Either Text (Map ProtoFileName ProtoFile)
  analyzeProtoFiles files = do
      -- The definitions in each input proto file, indexed by filename.
      definitionsByName <- mapM collectDefinitions filesByName
      let servicesByName = fmap collectServices filesByName
      let exportsByName = transitiveExports files
      let exportedEnvs = fmap (foldMap (definitionsByName !)) exportsByName
  
      let ingestFile f = ProtoFile
            { descriptor = f
            , haskellModule = m
            , definitions = definitionsByName ! n
            , services = servicesByName ! n
            , exportedEnv = qualifyEnv m $ exportedEnvs ! n
            , publicImports = [moduleNames ! i | i <- reexported]
            }
            where
              n = f ^. #name
              m = moduleNames ! n
              reexported =
                [ (f ^. #dependency) !! fromIntegral i
                | i <- f ^. #publicDependency
                ]
  
      return $ Map.fromList [ (f ^. #name, ingestFile f) | f <- files ]
    where
      filesByName = Map.fromList [(f ^. #name, f) | f <- files]
      moduleNames = fmap fdModuleName filesByName
  
  collectEnvFromDeps :: [ProtoFileName] -> Map ProtoFileName ProtoFile -> Env RdrNameStr
  collectEnvFromDeps deps filesByName =
      unions $ fmap (exportedEnv . (filesByName !)) deps
  
  -- | Get the Haskell 'ModuleName' corresponding to a given .proto file.
  fdModuleName :: FileDescriptorProto -> ModuleNameStr
  fdModuleName fd
        = fromString $ protoModuleName (T.unpack $ fd ^. #name)
  
  -- | Given a list of .proto files (topologically sorted), determine which
  -- files' definitions are exported by which files.
  --
  -- Files only export their own definitions, along with the definitions exported
  -- by any "import public" declarations.  (And any definitions that *those* files
  -- "import public", etc.)
  transitiveExports :: [FileDescriptorProto] -> Map ProtoFileName [ProtoFileName]
  -- Accumulate the transitive dependencies by folding over the files in
  -- topological order.
  transitiveExports = foldl' setExportsFromFile Map.empty
    where
      setExportsFromFile :: Map ProtoFileName [ProtoFileName]
                         -> FileDescriptorProto
                         -> Map ProtoFileName [ProtoFileName]
      setExportsFromFile prevExports fd
          = flip (Map.insert n) prevExports $
              n : concat [ prevExports ! ((fd ^. #dependency) !! fromIntegral i)
                         -- Note that publicDependency is a list of indices into
                         -- the dependency list.
                         | i <- fd ^. #publicDependency
                         ]
        where n = fd ^. #name
  