Simplify codemods

Summary:
Reduce the amount of boilerplate in Thrift codemods by moving the top-level code needed to run a codemod to a separate function, imaginatively named `run_codemod`. It handles command-line arguments (which we don't actually use except for the file name and includes), parses the Thrift files and invokes the provided function that does the actual work on a parsed AST.

Before:

```
int main(int argc, char** argv) {
  auto source_mgr = source_manager();
  auto program_bundle = parse_and_get_program(
      source_mgr, std::vector<std::string>(argv, argv + argc));
  if (!program_bundle) {
    return 1;
  }
  auto program = program_bundle->root_program();
  add_package(source_mgr, *program).run();

  return 0;
}
```

After:

```
int main(int argc, char** argv) {
  return apache::thrift::compiler::run_codemod(
      argc, argv, [](source_manager& sm, t_program& p) {
        add_package(sm, p).run();
      });
}
```

Fix a few minor issues such as wrong includes.

Reviewed By: yukselakinci

Differential Revision: D68466951

fbshipit-source-id: c8f79b1ac93a72197003350548c59bc8e6af41ce

third-party/thrift/src/thrift/compiler/codemod/add_package.cc

@@ -14,12 +14,10 @@
  * limitations under the License.
  */
 
-#include <vector>
-
 #include <thrift/compiler/ast/t_program_bundle.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
 #include <thrift/compiler/codemod/package_generator.h>
-#include <thrift/compiler/compiler.h>
 
 using apache::thrift::compiler::source_manager;
 using apache::thrift::compiler::t_program;
@@ -143,14 +141,8 @@ class add_package {
 } // namespace
 
 int main(int argc, char** argv) {
-  auto source_mgr = source_manager();
-  auto program_bundle = parse_and_get_program(
-      source_mgr, std::vector<std::string>(argv, argv + argc));
-  if (!program_bundle) {
-    return 1;
-  }
-  auto program = program_bundle->root_program();
-  add_package(source_mgr, *program).run();
-
-  return 0;
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        add_package(sm, p).run();
+      });
 }

third-party/thrift/src/thrift/compiler/codemod/annotate_non_optional_cpp_ref_fields.cc

@@ -14,17 +14,19 @@
  * limitations under the License.
  */
 
-#include <algorithm>
 #include <string>
 #include <string_view>
 #include <fmt/core.h>
 #include <thrift/compiler/ast/ast_visitor.h>
 #include <thrift/compiler/ast/t_program_bundle.h>
 #include <thrift/compiler/ast/t_struct.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
-#include <thrift/compiler/compiler.h>
 #include <thrift/compiler/generate/cpp/util.h>
 
+using apache::thrift::compiler::source_manager;
+using apache::thrift::compiler::t_program;
+
 namespace apache::thrift::compiler {
 namespace {
 
@@ -83,7 +85,6 @@ class AnnonateNonOptionalCppRefFields final {
  private:
   source_manager& source_manager_;
   t_program& program_;
-  std::unique_ptr<t_program_bundle> program_bundle_;
   codemod::file_manager file_manager_;
 
   bool maybe_annotate_field(const t_field& field) {
@@ -115,24 +116,12 @@ class AnnonateNonOptionalCppRefFields final {
   }
 };
 
-int run_main(int argc, char** argv) {
-  source_manager src_manager;
-  const std::unique_ptr<t_program_bundle> program_bundle =
-      parse_and_get_program(
-          src_manager, std::vector<std::string>(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);
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        apache::thrift::compiler::AnnonateNonOptionalCppRefFields(sm, p).run();
+      });
 }

third-party/thrift/src/thrift/compiler/codemod/codemod.cc

@@ -0,0 +1,42 @@
+/*
+ * 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 <fmt/core.h>
+#include <thrift/compiler/codemod/codemod.h>
+#include <thrift/compiler/compiler.h>
+
+namespace apache::thrift::compiler {
+
+int run_codemod(
+    int argc,
+    char** argv,
+    std::function<void(source_manager&, t_program&)> codemod) {
+  if (argc <= 1) {
+    fmt::print(stderr, "Usage: {} <thrift-file>\n", argv[0]);
+    return 1;
+  }
+
+  auto source_mgr = source_manager();
+  auto program_bundle = parse_and_get_program(
+      source_mgr, std::vector<std::string>(argv, argv + argc));
+  if (!program_bundle) {
+    return 1;
+  }
+  codemod(source_mgr, *program_bundle->root_program());
+  return 0;
+}
+
+} // namespace apache::thrift::compiler

third-party/thrift/src/thrift/compiler/codemod/codemod.h

@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <functional>
+#include <thrift/compiler/ast/t_program.h>
+#include <thrift/compiler/source_location.h>
+
+namespace apache::thrift::compiler {
+
+// Parses a Thrift file specified in the command-line arguments and runs
+// `codemod`, a function implementing the codemod, on the AST representation of
+// this file.
+[[nodiscard]] int run_codemod(
+    int argc,
+    char** argv,
+    std::function<void(source_manager&, t_program&)> codemod);
+
+} // namespace apache::thrift::compiler

third-party/thrift/src/thrift/compiler/codemod/cppref_to_structured.cc

@@ -19,8 +19,8 @@
 #include <thrift/compiler/ast/ast_visitor.h>
 #include <thrift/compiler/ast/t_program_bundle.h>
 #include <thrift/compiler/ast/t_struct.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
-#include <thrift/compiler/compiler.h>
 #include <thrift/compiler/generate/cpp/util.h>
 
 using namespace apache::thrift::compiler;
@@ -75,23 +75,15 @@ static void cppref_to_structured(
 }
 
 int main(int argc, char** argv) {
-  auto source_mgr = source_manager();
-  auto program_bundle = parse_and_get_program(
-      source_mgr, std::vector<std::string>(argv, argv + argc));
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        codemod::file_manager fm(sm, p);
 
-  if (!program_bundle) {
-    return 0;
-  }
-
-  auto program = program_bundle->root_program();
-
-  codemod::file_manager fm(source_mgr, *program);
-
-  const_ast_visitor visitor;
-  visitor.add_field_visitor(folly::partial(cppref_to_structured, std::ref(fm)));
-  visitor(*program);
-
-  fm.apply_replacements();
+        const_ast_visitor visitor;
+        visitor.add_field_visitor(
+            folly::partial(cppref_to_structured, std::ref(fm)));
+        visitor(p);
 
-  return 0;
+        fm.apply_replacements();
+      });
 }

third-party/thrift/src/thrift/compiler/codemod/hoist_annotated_types.cc

@@ -24,8 +24,8 @@
 #include <thrift/compiler/ast/ast_visitor.h>
 #include <thrift/compiler/ast/t_program_bundle.h>
 #include <thrift/compiler/ast/t_struct.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
-#include <thrift/compiler/compiler.h>
 #include <thrift/compiler/generate/cpp/util.h>
 
 using namespace apache::thrift::compiler;
@@ -350,27 +350,21 @@ class hoist_annotated_types {
   }
 
  private:
-  struct Typedef {
+  struct typedef_info {
     std::string type;
     t_typedef* ptr;
     std::string structured = "";
   };
-  std::map<std::string, Typedef> typedefs_;
+  std::map<std::string, typedef_info> typedefs_;
   codemod::file_manager fm_;
   source_manager sm_;
   t_program& prog_;
 };
 } // namespace
 
 int main(int argc, char** argv) {
-  auto source_mgr = source_manager();
-  auto program_bundle = parse_and_get_program(
-      source_mgr, std::vector<std::string>(argv, argv + argc));
-  if (!program_bundle) {
-    return 1;
-  }
-  auto program = program_bundle->root_program();
-  hoist_annotated_types(source_mgr, *program).run();
-
-  return 0;
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        hoist_annotated_types(sm, p).run();
+      });
 }

third-party/thrift/src/thrift/compiler/codemod/remove_cpp_noexcept_move.cc

@@ -17,10 +17,9 @@
 #include <folly/functional/Partial.h>
 
 #include <thrift/compiler/ast/ast_visitor.h>
-#include <thrift/compiler/ast/t_program_bundle.h>
 #include <thrift/compiler/ast/t_structured.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
-#include <thrift/compiler/compiler.h>
 
 using namespace apache::thrift::compiler;
 
@@ -35,26 +34,17 @@ static void remove_cpp_noexcept_move(
 }
 
 int main(int argc, char** argv) {
-  auto source_mgr = source_manager();
-  auto program_bundle = parse_and_get_program(
-      source_mgr, std::vector<std::string>(argv, argv + argc));
-
-  if (!program_bundle) {
-    return 0;
-  }
-
-  auto program = program_bundle->root_program();
-
-  codemod::file_manager fm(source_mgr, *program);
-
-  const_ast_visitor visitor;
-  visitor.add_struct_visitor(
-      folly::partial(remove_cpp_noexcept_move, std::ref(fm)));
-  visitor.add_union_visitor(
-      folly::partial(remove_cpp_noexcept_move, std::ref(fm)));
-  visitor(*program);
-
-  fm.apply_replacements();
-
-  return 0;
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        codemod::file_manager fm(sm, p);
+
+        const_ast_visitor visitor;
+        visitor.add_struct_visitor(
+            folly::partial(remove_cpp_noexcept_move, std::ref(fm)));
+        visitor.add_union_visitor(
+            folly::partial(remove_cpp_noexcept_move, std::ref(fm)));
+        visitor(p);
+
+        fm.apply_replacements();
+      });
 }

third-party/thrift/src/thrift/compiler/codemod/remove_redundant_field_custom_default_values.cc

@@ -14,19 +14,19 @@
  * limitations under the License.
  */
 
-#include <algorithm>
-#include <string>
-#include <string_view>
 #include <thrift/compiler/ast/ast_visitor.h>
 #include <thrift/compiler/ast/t_program_bundle.h>
 #include <thrift/compiler/ast/t_struct.h>
+#include <thrift/compiler/codemod/codemod.h>
 #include <thrift/compiler/codemod/file_manager.h>
-#include <thrift/compiler/compiler.h>
 #include <thrift/compiler/diagnostic.h>
 #include <thrift/compiler/generate/cpp/util.h>
 #include <thrift/compiler/sema/standard_validator.h>
 #include <thrift/compiler/source_location.h>
 
+using apache::thrift::compiler::source_manager;
+using apache::thrift::compiler::t_program;
+
 namespace apache::thrift::compiler {
 namespace {
 
@@ -117,25 +117,13 @@ class RemoveRedundantFieldCustomDefaultValues final {
   }
 };
 
-int run_main(int argc, char** argv) {
-  source_manager src_manager;
-  const std::unique_ptr<t_program_bundle> program_bundle =
-      parse_and_get_program(
-          src_manager, std::vector<std::string>(argv, argv + argc));
-
-  if (program_bundle == nullptr) {
-    return 1;
-  }
-
-  RemoveRedundantFieldCustomDefaultValues(
-      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);
+  return apache::thrift::compiler::run_codemod(
+      argc, argv, [](source_manager& sm, t_program& p) {
+        apache::thrift::compiler::RemoveRedundantFieldCustomDefaultValues(sm, p)
+            .run();
+      });
 }