diff --git a/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields.cc b/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields.cc new file mode 100644 index 00000000000000..568db5771821e7 --- /dev/null +++ b/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields.cc @@ -0,0 +1,138 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace apache::thrift::compiler { +namespace { + +constexpr std::string_view kAnnotation = "@cpp.AllowLegacyNonOptionalRef"; +constexpr const char* const kAnnotationUri = kCppAllowLegacyNonOptionalRefUri; + +bool should_annotate_field(const t_field& field) { + if (field.qualifier() == t_field_qualifier::optional) { + // Field is optional + return false; + } + + if (field.find_structured_annotation_or_null(kAnnotationUri) != nullptr) { + // kAnnotation is already present. + return false; + } + + // Field is not optional, and not already annotated. + // Is it a @cpp.Ref / cpp[2].ref[_type] ? + + if (field.find_structured_annotation_or_null(kCppRefUri) != nullptr) { + return true; + } + + if (field.has_annotation( + {"cpp.ref", "cpp2.ref", "cpp.ref_type", "cpp2.ref_type"})) { + return true; + } + + // Not a C++ reference field + return false; +} + +class AnnonateNonOptionalCppRefFields final { + public: + AnnonateNonOptionalCppRefFields( + source_manager& src_manager, t_program& program) + : source_manager_(src_manager), + program_(program), + file_manager_(source_manager_, program) {} + + void run() { + bool any_annotated = false; + + const_ast_visitor visitor; + visitor.add_field_visitor([&](const t_field& field) { + any_annotated |= maybe_annotate_field(field); + }); + visitor(program_); + + if (any_annotated) { + file_manager_.apply_replacements(); + } + } + + private: + source_manager& source_manager_; + t_program& program_; + std::unique_ptr program_bundle_; + codemod::file_manager file_manager_; + + bool maybe_annotate_field(const t_field& field) { + if (!should_annotate_field(field)) { + return false; + } + + // field is a cpp.Ref field, but not optional => add exemption annotation + + // First, make sure the include is present + file_manager_.add_include("thrift/annotation/cpp.thrift"); + + // Second, add a line above the field (with the correct indentation), with + // the annotation. + const source_range line_leading_whitespace = + file_manager_.get_line_leading_whitespace( + field.type().src_range().begin); + + const uint_least32_t start_of_line_offset = + line_leading_whitespace.begin.offset(); + file_manager_.add( + {.begin_pos = start_of_line_offset, + .end_pos = start_of_line_offset, + .new_content = fmt::format( + "{}{}\n", + source_manager_.get_text_range(line_leading_whitespace), + kAnnotation)}); + return true; + } +}; + +int run_main(int argc, char** argv) { + source_manager src_manager; + const std::unique_ptr program_bundle = + parse_and_get_program( + src_manager, std::vector(argv, argv + argc)); + + if (program_bundle == nullptr) { + return 1; + } + + AnnonateNonOptionalCppRefFields(src_manager, *program_bundle->root_program()) + .run(); + return 0; +} + +} // namespace +} // namespace apache::thrift::compiler + +int main(int argc, char** argv) { + return apache::thrift::compiler::run_main(argc, argv); +} diff --git a/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields_test.py b/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields_test.py new file mode 100644 index 00000000000000..66b9f0742921c6 --- /dev/null +++ b/third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields_test.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pyre-unsafe + +import os +import shutil +import tempfile +import textwrap +import unittest + +import pkg_resources + +from thrift.compiler.codemod.test_utils import read_file, run_binary, write_file + + +class AnnotateNonOptionalCppRefFieldsTest(unittest.TestCase): + def setUp(self): + tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, tmp, True) + self.tmp = tmp + self.addCleanup(os.chdir, os.getcwd()) + os.chdir(self.tmp) + self.maxDiff = None + + def test_basic_replace(self): + write_file( + "foo.thrift", + textwrap.dedent( + """\ + struct NoRefField { + // No-op + 1: i32 a; + } + + // Should annotate + struct A { + 1: string field1 (cpp.ref = "true"); + + 2: string field2 (cpp2.ref = "true"); + + 3: string field3 (cpp.ref_type = "shared"); + + 4: string field4 (cpp2.ref_type = "shared"); + + 5: optional string optional_field1 (cpp.ref = "true"); + + 6: optional string optional_field2 (cpp2.ref = "true"); + + 7: optional string optional_field3 (cpp.ref_type = "shared"); + + 8: optional string optional_field4 (cpp2.ref_type = "shared"); + } + """ + ), + ) + + binary = pkg_resources.resource_filename(__name__, "codemod") + run_binary(binary, "foo.thrift") + + self.assertEqual( + read_file("foo.thrift"), + textwrap.dedent( + """\ + include "thrift/annotation/cpp.thrift" + + struct NoRefField { + // No-op + 1: i32 a; + } + + // Should annotate + struct A { + @cpp.AllowLegacyNonOptionalRef + 1: string field1 (cpp.ref = "true"); + + @cpp.AllowLegacyNonOptionalRef + 2: string field2 (cpp2.ref = "true"); + + @cpp.AllowLegacyNonOptionalRef + 3: string field3 (cpp.ref_type = "shared"); + + @cpp.AllowLegacyNonOptionalRef + 4: string field4 (cpp2.ref_type = "shared"); + + 5: optional string optional_field1 (cpp.ref = "true"); + + 6: optional string optional_field2 (cpp2.ref = "true"); + + 7: optional string optional_field3 (cpp.ref_type = "shared"); + + 8: optional string optional_field4 (cpp2.ref_type = "shared"); + } + """ + ), + )