1#!/usr/bin/env python3
2
3# Copyright 2021 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import collections
18import os
19
20
21def to_inc(filename):
22    """Given filename, synthesize what should go in an include statement to get that file"""
23    if filename.startswith("include/"):
24        return '<%s>' % filename[len("include/"):]
25    return '"%s"' % filename
26
27
28def set_pragmas(filename, pragmas):
29    """Set the file-level IWYU pragma in filename"""
30    lines = []
31    saw_first_define = False
32    for line in open(filename).read().splitlines():
33        if line.startswith('// IWYU pragma: '):
34            continue
35        lines.append(line)
36        if not saw_first_define and line.startswith('#define '):
37            saw_first_define = True
38            lines.append('')
39            for pragma in pragmas:
40                lines.append('// IWYU pragma: %s' % pragma)
41            lines.append('')
42    open(filename, 'w').write('\n'.join(lines) + '\n')
43
44
45def set_exports(pub, cg):
46    """In file pub, mark the include for cg with IWYU pragma: export"""
47    lines = []
48    for line in open(pub).read().splitlines():
49        if line.startswith('#include %s' % to_inc(cg)):
50            lines.append('#include %s  // IWYU pragma: export' % to_inc(cg))
51        else:
52            lines.append(line)
53    open(pub, 'w').write('\n'.join(lines) + '\n')
54
55
56CG_ROOTS_GRPC = (
57    (r'sync', 'grpc/support/sync.h', False),
58    (r'atm', 'grpc/support/atm.h', False),
59    (r'grpc_types', 'grpc/grpc.h', True),
60    (r'gpr_types', 'grpc/grpc.h', True),
61    (r'compression_types', 'grpc/compression.h', True),
62    (r'connectivity_state', 'grpc/grpc.h', True),
63)
64
65CG_ROOTS_GRPCPP = [
66    (r'status_code_enum', 'grpcpp/support/status.h', False),
67]
68
69
70def fix_tree(tree, cg_roots):
71    """Fix one include tree"""
72    # Map of filename --> paths including that filename
73    reverse_map = collections.defaultdict(list)
74    # The same, but for things with '/impl/codegen' in their names
75    cg_reverse_map = collections.defaultdict(list)
76    for root, dirs, files in os.walk(tree):
77        root_map = cg_reverse_map if '/impl/codegen' in root else reverse_map
78        for filename in files:
79            root_map[filename].append(root)
80    # For each thing in '/impl/codegen' figure out what exports it
81    for filename, paths in cg_reverse_map.items():
82        print("****", filename)
83        # Exclude non-headers
84        if not filename.endswith('.h'):
85            continue
86        pragmas = []
87        # Check for our 'special' headers: if we see one of these, we just
88        # hardcode where they go to because there's some complicated rules.
89        for root, target, friend in cg_roots:
90            print(root, target, friend)
91            if filename.startswith(root):
92                pragmas = ['private, include <%s>' % target]
93                if friend:
94                    pragmas.append('friend "src/.*"')
95                if len(paths) == 1:
96                    path = paths[0]
97                    if filename.startswith(root + '.'):
98                        set_exports('include/' + target, path + '/' + filename)
99                    if filename.startswith(root + '_'):
100                        set_exports(path + '/' + root + '.h',
101                                    path + '/' + filename)
102        # If the path for a file in /impl/codegen is ambiguous, just don't bother
103        if not pragmas and len(paths) == 1:
104            path = paths[0]
105            # Check if we have an exporting candidate
106            if filename in reverse_map:
107                proper = reverse_map[filename]
108                # And that it too is unambiguous
109                if len(proper) == 1:
110                    # Build the two relevant pathnames
111                    cg = path + '/' + filename
112                    pub = proper[0] + '/' + filename
113                    # And see if the public file actually includes the /impl/codegen file
114                    if ('#include %s' % to_inc(cg)) in open(pub).read():
115                        # Finally, if it does, we'll set that pragma
116                        pragmas = ['private, include %s' % to_inc(pub)]
117                        # And mark the export
118                        set_exports(pub, cg)
119        # If we can't find a good alternative include to point people to,
120        # mark things private anyway... we don't want to recommend people include
121        # from impl/codegen
122        if not pragmas:
123            pragmas = ['private']
124        for path in paths:
125            set_pragmas(path + '/' + filename, pragmas)
126
127
128fix_tree('include/grpc', CG_ROOTS_GRPC)
129fix_tree('include/grpcpp', CG_ROOTS_GRPCPP)
130