-
Notifications
You must be signed in to change notification settings - Fork 108
/
Plugin.hs
130 lines (116 loc) · 5.04 KB
/
Plugin.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
-- 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 OverloadedStrings #-}
module Data.ProtoLens.Compiler.Plugin
( ProtoFileName
, ProtoFile(..)
, analyzeProtoFiles
, collectEnvFromDeps
, outputFilePath
, moduleName
, moduleNameStr
) where
import Data.Char (toUpper)
import Data.List (foldl', intercalate)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map, unions, (!))
import Data.Monoid ((<>))
import Data.String (fromString)
import qualified Data.Text as T
import Data.Text (Text)
import Lens.Family2
import Proto.Google.Protobuf.Descriptor (FileDescriptorProto)
import Proto.Google.Protobuf.Descriptor_Fields (name, dependency, publicDependency)
import System.FilePath (dropExtension, splitDirectories)
import Data.ProtoLens.Compiler.Definitions
import Data.ProtoLens.Compiler.Combinators (ModuleName, Name, QName)
-- | The filename of an input .proto file.
type ProtoFileName = Text
data ProtoFile = ProtoFile
{ descriptor :: FileDescriptorProto
, haskellModule :: ModuleName
, definitions :: Env Name
, services :: [ServiceInfo]
-- | The names of proto files exported (transitively, via "import public"
-- decl) by this file.
, exports :: [ProtoFileName]
, exportedEnv :: Env QName
}
-- Given a list of FileDescriptorProtos, collect information about each file
-- into a map of 'ProtoFile's keyed by 'ProtoFileName'.
analyzeProtoFiles :: Text -> [FileDescriptorProto] -> Map ProtoFileName ProtoFile
analyzeProtoFiles modulePrefix files =
Map.fromList [ (f ^. name, ingestFile f) | f <- files ]
where
filesByName = Map.fromList [(f ^. name, f) | f <- files]
moduleNames = fmap (moduleName modulePrefix) filesByName
-- The definitions in each input proto file, indexed by filename.
definitionsByName = fmap collectDefinitions filesByName
servicesByName = fmap collectServices filesByName
-- The exports from each .proto file (including any "public import"
-- dependencies), as they appear to other modules that are importing them;
-- i.e., qualified by module name.
exportsByName = transitiveExports files
localExports = Map.intersectionWith qualifyEnv moduleNames definitionsByName
exportedEnvs = fmap (\es -> unions [localExports ! e | e <- es]) exportsByName
ingestFile f = ProtoFile
{ descriptor = f
, haskellModule = moduleNames ! n
, definitions = definitionsByName ! n
, services = servicesByName ! n
, exports = exportsByName ! n
, exportedEnv = exportedEnvs ! n
}
where
n = f ^. name
collectEnvFromDeps :: [ProtoFileName] -> Map ProtoFileName ProtoFile -> Env QName
collectEnvFromDeps deps filesByName =
unions $ fmap (exportedEnv . (filesByName !)) deps
-- | Get the output file path (for CodeGeneratorResponse.File) for a Haskell
-- ModuleName.
outputFilePath :: String -> Text
outputFilePath n = T.replace "." "/" (T.pack n) <> ".hs"
-- | Get the Haskell 'ModuleName' corresponding to a given .proto file.
moduleName :: Text -> FileDescriptorProto -> ModuleName
moduleName modulePrefix fd
= fromString $ moduleNameStr (T.unpack modulePrefix) (T.unpack $ fd ^. name)
-- | Get the Haskell module name corresponding to a given .proto file.
moduleNameStr :: String -> FilePath -> String
moduleNameStr prefix path = fixModuleName rawModuleName
where
fixModuleName "" = ""
-- Characters allowed in Bazel filenames but not in module names:
fixModuleName ('.':c:cs) = '.' : toUpper c : fixModuleName cs
fixModuleName ('_':c:cs) = toUpper c : fixModuleName cs
fixModuleName ('-':c:cs) = toUpper c : fixModuleName cs
fixModuleName (c:cs) = c : fixModuleName cs
rawModuleName = intercalate "."
. (prefix :)
. splitDirectories $ dropExtension
$ path
-- | 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.
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