Only import programs that are needed in runtime

Summary:
Currently we import all includes even if they are not referenced later.

However, in BuildGraph API, we won't generate runtime dependencies for modules that are not referenced later since it's unnecessary.

This discrepancy breaks Configerator since they won't materialize modules not in BuildGraph dependencies.

This diff changed the logic and removed unnecessary imports in Patch.

Basically it iterates all types in the root program to see whether it referenced programs outside root, and only imports those includes.

The same idea should work for thrift-python imports which I am experimenting in the next stacked diff

Reviewed By: pranavtbhat

Differential Revision: D77348074

fbshipit-source-id: 9686fd975734b2cbeea8ddadb9b5c2422f76169c

Author: Mizuchi

third-party/thrift/src/thrift/compiler/generate/python/util.cc

@@ -18,6 +18,7 @@
 
 #include <boost/algorithm/string.hpp>
 #include <boost/algorithm/string/replace.hpp>
+#include <thrift/compiler/ast/ast_visitor.h>
 #include <thrift/compiler/generate/cpp/util.h>
 
 namespace apache::thrift::compiler {
@@ -36,6 +37,53 @@ bool is_patch_program(const t_program* prog) {
               : false;
 }
 
+static void get_needed_includes_impl(
+    const t_program* root,
+    const t_type& type,
+    std::unordered_set<const t_type*>& seen,
+    std::unordered_set<const t_program*>& result) {
+  if (!seen.insert(&type).second) {
+    return;
+  }
+  if (type.get_program() && type.get_program() != root) {
+    // If type is not in root, root needs type's program in runtime.
+    result.insert(type.get_program());
+    return;
+  }
+  if (const t_typedef* asTypedef = type.try_as<t_typedef>()) {
+    get_needed_includes_impl(root, *asTypedef->get_type(), seen, result);
+  }
+  if (const t_list* asList = type.try_as<t_list>()) {
+    get_needed_includes_impl(root, *asList->get_elem_type(), seen, result);
+  }
+  if (const t_set* asSet = type.try_as<t_set>()) {
+    get_needed_includes_impl(root, *asSet->get_elem_type(), seen, result);
+  }
+  if (const t_map* asMap = type.try_as<t_map>()) {
+    get_needed_includes_impl(root, *asMap->get_key_type(), seen, result);
+    get_needed_includes_impl(root, *asMap->get_val_type(), seen, result);
+  }
+  if (const t_structured* asStructured = type.try_as<t_structured>()) {
+    for (auto& field : asStructured->fields()) {
+      get_needed_includes_impl(root, *field.type().get_type(), seen, result);
+    }
+  }
+}
+
+std::unordered_set<const t_program*> needed_includes_in_runtime(
+    const t_program* root) {
+  std::unordered_set<const t_program*> programs;
+  std::unordered_set<const t_type*> seen;
+  const_ast_visitor visitor;
+  visitor.add_root_definition_visitor([&](const t_named& def) {
+    if (auto type = dynamic_cast<const t_type*>(&def)) {
+      get_needed_includes_impl(root, *type, seen, programs);
+    }
+  });
+  visitor(*root);
+  return programs;
+}
+
 bool type_contains_patch(const t_type* type) {
   if (auto typedf = dynamic_cast<t_typedef const*>(type)) {
     return type_contains_patch(typedf->get_type());

third-party/thrift/src/thrift/compiler/generate/python/util.h

@@ -48,6 +48,10 @@ std::string get_py3_namespace_with_name_and_prefix(
     const std::string& prefix,
     const std::string& sep = ".");
 
+// List of programs that `root` needed in runtime
+std::unordered_set<const t_program*> needed_includes_in_runtime(
+    const t_program* root);
+
 inline const std::unordered_set<std::string>& get_python_reserved_names() {
   static const std::unordered_set<std::string> keywords = {
       "False",  "None",    "True",    "and",      "as",       "assert", "async",

third-party/thrift/src/thrift/compiler/generate/t_mstch_python_generator.cc

@@ -248,7 +248,8 @@ class python_mstch_program : public mstch_program {
           {"included_module_mangle", it->ns_mangle},
           {"has_services?", it->has_services},
           {"has_types?", it->has_types},
-          {"is_patch?", it->is_patch}});
+          {"is_patch?", it->is_patch},
+          {"needed_in_runtime?", it->needed_in_runtime}});
     }
     return a;
   }
@@ -330,9 +331,11 @@ class python_mstch_program : public mstch_program {
     bool has_services;
     bool has_types;
     bool is_patch;
+    bool needed_in_runtime;
   };
 
   void gather_included_program_namespaces() {
+    auto needed_includes = needed_includes_in_runtime(program_);
     for (const t_program* included_program :
          program_->get_includes_for_codegen()) {
       bool has_types =
@@ -347,7 +350,8 @@ class python_mstch_program : public mstch_program {
               included_program, get_option("root_module_prefix")),
           !included_program->services().empty(),
           has_types,
-          is_patch_program(included_program)};
+          is_patch_program(included_program),
+          static_cast<bool>(needed_includes.count(included_program))};
     }
   }
 
@@ -366,6 +370,7 @@ class python_mstch_program : public mstch_program {
       ns.has_services = false;
       ns.has_types = true;
       ns.is_patch = is_patch_program(prog);
+      ns.needed_in_runtime = true;
       include_namespaces_[path] = std::move(ns);
     }
   }

third-party/thrift/src/thrift/compiler/generate/templates/patch/thrift_patch.py.mustache

@@ -56,11 +56,11 @@ import {{program:safe_patch_module_path}}.thrift_types as _fbthrift_safe_patch_t
 {{/if (not program:safe_patch?)}}
 
 {{#program:include_namespaces}}
-{{#if has_types?}}{{#if (not is_patch?)}}
+{{#if has_types?}}{{#if (not is_patch?)}}{{#if needed_in_runtime?}}
 
 import {{included_module_path}}.{{> ../python/types/types_import_path}} as {{included_module_mangle}}__{{> ../python/types/types_import_path}}
 import {{included_module_path}}.thrift_patch
-{{/if (not is_patch?)}}{{/if has_types?}}
+{{/if needed_in_runtime?}}{{/if (not is_patch?)}}{{/if has_types?}}
 {{/program:include_namespaces}}
 
 {{#program:structs}}