Partition programs with multi-line dispatch tables as text (#1626)

native/llvm-codegen/src/partition.cpp

@@ -359,7 +359,34 @@ std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
     } else if (Line.front() == '@') {
       Symbol S;
       S.Header = Line;
-      if (!parseGlobal(Line, S) || !Add(std::move(S)))
+      if (!parseGlobal(Line, S))
+        return std::nullopt;
+      // The emitter writes an interface dispatch table one element per
+      // line: the definition opens its aggregate initializer at the end of
+      // the line, and a line that is only the closing bracket ends it.
+      StringRef Opening = Line.rtrim();
+      if (S.Kind == SymbolKind::Variable &&
+          (Opening.ends_with(" [") || Opening.ends_with(" {"))) {
+        StringRef Close = Opening.back() == '[' ? "]" : "}";
+        size_t Cursor = Next;
+        while (true) {
+          if (Cursor >= Text.size())
+            return std::nullopt;
+          size_t LineEnd = Text.find('\n', Cursor);
+          if (LineEnd == StringRef::npos)
+            LineEnd = Text.size();
+          StringRef Element = Text.slice(Cursor, LineEnd);
+          if (Element == Close) {
+            S.Body = Text.slice(Next, LineEnd);
+            Next = LineEnd + 1;
+            break;
+          }
+          if (!Element.starts_with("  "))
+            return std::nullopt;
+          Cursor = LineEnd + 1;
+        }
+      }
+      if (!Add(std::move(S)))
         return std::nullopt;
     } else if (Line.starts_with("declare ")) {
       Symbol S;

packages/compiler/test/native-partitions.test.ts

@@ -162,6 +162,31 @@ describe.runIf(helper !== null)("LLVM program partitions", () => {
     });
   }
 
+  test("a table written one element per line keeps the text partitioner and its reuse", async () => {
+    const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-table-"));
+    dirs.push(dir);
+    // The emitter writes interface dispatch tables one element per line.
+    const table = [
+      "@table = internal constant [2 x ptr] [",
+      "  ptr @step3,",
+      "  ptr @step9",
+      "]",
+    ].join("\n");
+    const source = (edit: number) =>
+      program(edit, table).replace(
+        "  %total = load i64, ptr @total",
+        "  %entry = load ptr, ptr @table\n  call void %entry(i64 0)\n  %total = load i64, ptr @total",
+      );
+    const cache = join(dir, "cache");
+    const first = await emit(dir, source(0), { outputs: 4, optLevel: "2", cache, name: "first" });
+    expect(await exitCode(dir, first, "first")).toBe(expectedExit());
+    const sources = await cacheEntries(cache, ".bc");
+    const edited = await emit(dir, source(7), { outputs: 4, optLevel: "2", cache, name: "edited" });
+    const after = await cacheEntries(cache, ".bc");
+    expect([...after].filter((entry) => !sources.has(entry))).toHaveLength(1);
+    expect(await exitCode(dir, edited, "edited")).toBe(expectedExit(7));
+  });
+
   test("inputs outside the text partitioner's grammar are still partitioned", async () => {
     const dir = await mkdtemp(join(tmpdir(), "scriptc-partitions-fallback-"));
     dirs.push(dir);