1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include <folly/functional/Partial.h>
18 
19 #include <thrift/compiler/ast/ast_visitor.h>
20 #include <thrift/compiler/ast/t_struct.h>
21 #include <thrift/compiler/codemod/file_manager.h>
22 #include <thrift/compiler/compiler.h>
23 #include <thrift/compiler/lib/cpp2/util.h>
24 
25 using namespace apache::thrift::compiler;
26 
only_annotations(const t_node & node,std::vector<std::string> annotations)27 static bool only_annotations(
28     const t_node& node, std::vector<std::string> annotations) {
29   if (node.annotations().size() != annotations.size()) {
30     return false;
31   }
32 
33   std::vector<std::string> node_annotations;
34 
35   for (const auto& annotation : node.annotations()) {
36     node_annotations.push_back(annotation.first);
37   }
38 
39   std::sort(annotations.begin(), annotations.end());
40 
41   return std::equal(
42       annotations.begin(), annotations.end(), node_annotations.begin());
43 }
44 
45 // Migrates cpp.ref and cpp2.ref unstructured annotations to its
46 // newly implemented structured versions.
47 // NOTE: Rely on automated formatting to fix formatting issues.
cppref_to_structured(codemod::file_manager & fm,const t_field & field)48 static void cppref_to_structured(
49     codemod::file_manager& fm, const t_field& field) {
50   if (!field.annotations().count("cpp.ref") &&
51       !field.annotations().count("cpp2.ref")) {
52     return;
53   }
54 
55   fm.add_include("thrift/annotation/cpp.thrift");
56 
57   const auto field_begin_offset = field.src_range().begin().offset();
58 
59   fm.add(
60       {field_begin_offset,
61        field_begin_offset,
62        "@cpp.Ref{type = cpp.RefType.Unique}\n"});
63 
64   if (only_annotations(field, {"cpp.ref", "cpp2.ref"})) {
65     fm.remove_all_annotations(field);
66     return;
67   }
68 
69   for (const auto& annotation : field.annotations()) {
70     if (annotation.first == "cpp.ref" || annotation.first == "cpp2.ref") {
71       fm.remove(annotation);
72     }
73   }
74 }
75 
main(int argc,char ** argv)76 int main(int argc, char** argv) {
77   auto program_bundle =
78       parse_and_get_program(std::vector<std::string>(argv, argv + argc));
79 
80   if (!program_bundle) {
81     return 0;
82   }
83 
84   auto program = program_bundle->root_program();
85 
86   /* Currently other languages except C++ don't support `@cpp.Ref` */
87   for (const auto& n : program->namespaces()) {
88     if (n.first != "cpp" && n.first != "cpp2") {
89       return 0;
90     }
91   }
92 
93   codemod::file_manager fm(*program);
94 
95   const_ast_visitor visitor;
96   visitor.add_field_visitor(folly::partial(cppref_to_structured, std::ref(fm)));
97   visitor(*program);
98 
99   fm.apply_replacements();
100 
101   return 0;
102 }
103