gn2bp: Add genrules to their respective architecture sink

Test: m cronet_aml_base_base
Change-Id: I410f05320bdb12a070154fbfc004ea2ce9b1343b
diff --git a/tools/gn2bp/gen_android_bp b/tools/gn2bp/gen_android_bp
index e919406..747c45a 100755
--- a/tools/gn2bp/gen_android_bp
+++ b/tools/gn2bp/gen_android_bp
@@ -211,6 +211,8 @@
     self.cppflags = set()
     self.local_include_dirs = set()
     self.export_system_include_dirs = set()
+    self.generated_headers = set()
+    self.export_generated_headers = set()
 
   def to_string(self, output):
     nested_out = []
@@ -226,6 +228,8 @@
     self._output_field(nested_out, 'cppflags')
     self._output_field(nested_out, 'local_include_dirs')
     self._output_field(nested_out, 'export_system_include_dirs')
+    self._output_field(nested_out, 'generated_headers')
+    self._output_field(nested_out, 'export_generated_headers')
 
     if nested_out:
       output.append('    %s: {' % self.name)
@@ -968,6 +972,15 @@
     if '-fexceptions' in flag:
       module.cppflags.add('-fexceptions')
 
+def add_genrule_per_arch(module, dep_module, type):
+  module.generated_headers.update(dep_module.genrule_headers)
+  # If the module is a static library, export all the generated headers.
+  if type == 'cc_library_static':
+    module.export_generated_headers.update(dep_module.genrule_headers)
+  module.srcs.update(dep_module.genrule_srcs)
+  module.shared_libs.update(dep_module.genrule_shared_libs)
+  module.header_libs.update(dep_module.genrule_header_libs)
+
 def set_module_include_dirs(module, cflags, include_dirs):
   for flag in cflags:
     if '-isystem' in flag:
@@ -1089,11 +1102,6 @@
     module.local_include_dirs = [d for d in module.local_include_dirs
                                  if d not in local_include_dirs_denylist]
 
-
-  # If the module is a static library, export all the generated headers.
-  if module.type == 'cc_library_static':
-    module.export_generated_headers = module.generated_headers
-
   # dep_name is an unmangled GN target name (e.g. //foo:bar(toolchain)).
   # Currently, only one module is generated from target even target has multiple toolchains.
   # And module is generated based on the first visited target.
@@ -1134,10 +1142,13 @@
     elif dep_module.type == 'cc_defaults':
       module.defaults.append(dep_module.name)
     elif dep_module.type == 'genrule':
-      module.generated_headers.update(dep_module.genrule_headers)
-      module.srcs.update(dep_module.genrule_srcs)
-      module.shared_libs.update(dep_module.genrule_shared_libs)
-      module.header_libs.update(dep_module.genrule_header_libs)
+      for arch_name, arch in target.arch.items():
+        if dep_module.name.endswith(arch_name):
+          add_genrule_per_arch(module.target[arch_name], dep_module, module.type)
+      if dep_module.name.endswith("_gen"):
+        module.srcs.update(dep_module.genrule_srcs)
+        module.generated_headers.update(dep_module.genrule_headers)
+        module.export_generated_headers.update(dep_module.genrule_headers)
     elif dep_module.type == 'cc_binary':
       continue  # Ignore executables deps (used by cmdline integration tests).
     else: