fix(tools): read preprocessor branch combinations lazily in import_audit

tools/import_audit.py

@@ -1,5 +1,6 @@
 import argparse
 import hashlib
+import itertools
 import json
 import os
 import re
@@ -536,22 +537,25 @@ def branch_variants(body):
     root.append(body[position:])
 
     def expand(items):
-        variants = [""]
+        parts = []
         for item in items:
             if isinstance(item, str):
-                variants = [variant + item for variant in variants]
+                parts.append((item,))
                 continue
             choices = [choice for alternative in item["alternatives"] for choice in expand(alternative)]
-            variants = [variant + choice for variant in variants for choice in choices]
-            if len(variants) > BRANCH_VARIANT_LIMIT:
+            if len(choices) > BRANCH_VARIANT_LIMIT:
                 raise AuditError(f"more than {BRANCH_VARIANT_LIMIT} preprocessor branch combinations in one function")
-        return variants
+            parts.append(choices)
+        for count, variant in enumerate(itertools.product(*parts), 1):
+            if count > BRANCH_VARIANT_LIMIT:
+                raise AuditError(f"more than {BRANCH_VARIANT_LIMIT} preprocessor branch combinations in one function")
+            yield "".join(variant)
 
     return expand(root)
 
 
 def variant_calls(body):
-    return [unconditional_calls(variant) for variant in branch_variants(body)]
+    return (unconditional_calls(variant) for variant in branch_variants(body))
 
 
 def unconditional_calls(body):

tools/tests/test_import_audit.py

@@ -1601,6 +1601,24 @@ def test_guarded_exports_are_partial_and_unconditional_ones_are_stubs(self):
         self.assertEqual(partial, {(library, nid_names.compute_nid(name)) for library in ("libGuard.prx", "libShared.prx")
                                    for name, kind in kinds.items() if kind == "partial"})
 
+    def test_many_host_branches_only_fail_when_every_combination_must_be_read(self):
+        blocks = "".join(f"#ifdef HOST_{index}\n    value += {index};\n#endif\n" for index in range(9))
+        source = self.work.source
+        write_source(source, "libHosts",
+                     "void Fail() { NotImplemented_nid_no_patch(__func__); }\n"
+                     "int APS5_VABI sceStub(int value) { NotImplemented_nid_no_patch(__func__); return 0; }\n"
+                     f"int APS5_VABI sceHosts(int value) {{\n{blocks}    return value;\n}}\n"
+                     f"int APS5_VABI sceLastHost(int value) {{\n{blocks}#ifdef HOST_LAST\n    Fail();\n#endif\n    return value;\n}}\n")
+        stubs, partial = import_audit.source_stub_classes(source)
+        stub, hosts, last = (("libHosts.prx", nid_names.compute_nid(name)) for name in ("sceStub", "sceHosts", "sceLastHost"))
+        self.assertEqual(set(stubs), {stub, last})
+        self.assertEqual(partial, {last})
+        write_source(source, "libHosts",
+                     "void Fail() { NotImplemented_nid_no_patch(__func__); }\n"
+                     f"int APS5_VABI sceEveryHost(int value) {{\n{blocks}    Fail();\n    return value;\n}}\n")
+        with self.assertRaisesRegex(import_audit.AuditError, "more than 256 preprocessor branch combinations"):
+            import_audit.source_stub_classes(source)
+
     def test_partial_import_is_its_own_class(self):
         direct, guarded = nid_names.compute_nid("sceDirect"), nid_names.compute_nid("sceGuarded")
         stubs = {("libA.prx", direct): "sceDirect", ("libA.prx", guarded): "sceGuarded", ("libB.prx", guarded): "sceGuarded"}