[wasm-split] Do not split out functions referring to segments (#6517)

Since data and elem segments cannot be imported or exported, there is no way to
access them from the secondary module, so functions that need to refer to them
cannot be split out.

Fixes #6512.

2 files changed, 156 insertions, 4 deletions

diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp
index 8ef0e7b8c..27b6a0902 100644
--- a/src/ir/module-splitting.cpp
+++ b/src/ir/module-splitting.cpp
@@ -282,7 +282,7 @@ struct ModuleSplitter {
   // Initialization helpers
   static std::unique_ptr<Module> initSecondary(const Module& primary);
   static std::pair<std::set<Name>, std::set<Name>>
-  classifyFunctions(const Module& primary, const Config& config);
+  classifyFunctions(Module& primary, const Config& config);
   static std::map<Name, Name> initExportedPrimaryFuncs(const Module& primary);
 
   // Other helpers
@@ -343,7 +343,64 @@ std::unique_ptr<Module> ModuleSplitter::initSecondary(const Module& primary) {
 }
 
 std::pair<std::set<Name>, std::set<Name>>
-ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) {
+ModuleSplitter::classifyFunctions(Module& primary, const Config& config) {
+  // Find functions that refer to data or element segments. These functions must
+  // remain in the primary module because segments cannot be exported to be
+  // accessed from the secondary module.
+  //
+  // TODO: Investigate other options, such as moving the segments to the
+  // secondary module or replacing the segment-using instructions in the
+  // secondary module with calls to imports.
+  ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>>
+    segmentReferrerCollector(
+      primary, [&](Function* func, std::vector<Name>& segmentReferrers) {
+        if (func->imported()) {
+          return;
+        }
+
+        struct SegmentReferrerCollector
+          : PostWalker<SegmentReferrerCollector,
+                       UnifiedExpressionVisitor<SegmentReferrerCollector>> {
+          bool hasSegmentReference = false;
+
+          void visitExpression(Expression* curr) {
+
+#define DELEGATE_ID curr->_id
+
+#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>();
+#define DELEGATE_GET_FIELD(id, field) cast->field
+#define DELEGATE_FIELD_TYPE(id, field)
+#define DELEGATE_FIELD_HEAPTYPE(id, field)
+#define DELEGATE_FIELD_CHILD(id, field)
+#define DELEGATE_FIELD_OPTIONAL_CHILD(id, field)
+#define DELEGATE_FIELD_INT(id, field)
+#define DELEGATE_FIELD_LITERAL(id, field)
+#define DELEGATE_FIELD_NAME(id, field)
+#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
+#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
+#define DELEGATE_FIELD_ADDRESS(id, field)
+
+#define DELEGATE_FIELD_NAME_KIND(id, field, kind)                              \
+  if (kind == ModuleItemKind::DataSegment ||                                   \
+      kind == ModuleItemKind::ElementSegment) {                                \
+    hasSegmentReference = true;                                                \
+  }
+
+#include "wasm-delegations-fields.def"
+          }
+        };
+        SegmentReferrerCollector collector;
+        collector.walkFunction(func);
+        if (collector.hasSegmentReference) {
+          segmentReferrers.push_back(func->name);
+        }
+      });
+
+  std::unordered_set<Name> segmentReferrers;
+  for (auto& [_, referrers] : segmentReferrerCollector.map) {
+    segmentReferrers.insert(referrers.begin(), referrers.end());
+  }
+
   std::set<Name> primaryFuncs, secondaryFuncs;
   for (auto& func : primary.functions) {
     // In JSPI mode exported functions cannot be moved to the secondary
@@ -351,14 +408,15 @@ ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) {
     // wrapper. Exported JSPI functions can still benefit from splitting though
     // since only the JSPI wrapper stub will remain in the primary module.
     if (func->imported() || config.primaryFuncs.count(func->name) ||
-        (config.jspi && ExportUtils::isExported(primary, *func))) {
+        (config.jspi && ExportUtils::isExported(primary, *func)) ||
+        segmentReferrers.count(func->name)) {
       primaryFuncs.insert(func->name);
     } else {
       assert(func->name != primary.start && "The start function must be kept");
       secondaryFuncs.insert(func->name);
     }
   }
-  return std::make_pair(primaryFuncs, secondaryFuncs);
+  return std::make_pair(std::move(primaryFuncs), std::move(secondaryFuncs));
 }
 
 std::map<Name, Name>
diff --git a/test/lit/wasm-split/segments.wast b/test/lit/wasm-split/segments.wast
new file mode 100644
index 000000000..7a14afc42
--- /dev/null
+++ b/test/lit/wasm-split/segments.wast
@@ -0,0 +1,94 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-split %s -all --keep-funcs=foo -g -o1 %t.1.wasm -o2 %t.2.wasm
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+(module
+
+ ;; PRIMARY:      (type $0 (func))
+
+ ;; PRIMARY:      (type $data-array (array i8))
+ (type $data-array (array i8))
+
+ ;; PRIMARY:      (type $elem-array (array externref))
+ (type $elem-array (array externref))
+
+ ;; PRIMARY:      (memory $mem 0)
+ (memory $mem 0)
+
+ ;; PRIMARY:      (data $data "hello world")
+ (data $data "hello world")
+
+ ;; PRIMARY:      (elem $elem externref)
+ (elem $elem externref)
+
+ ;; PRIMARY:      (export "memory" (memory $mem))
+
+ ;; PRIMARY:      (func $data.drop
+ ;; PRIMARY-NEXT:  (data.drop $data)
+ ;; PRIMARY-NEXT: )
+ (func $data.drop
+  (data.drop $data)
+ )
+
+ ;; PRIMARY:      (func $memory.init
+ ;; PRIMARY-NEXT:  (memory.init $data
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $memory.init
+  (memory.init $mem $data
+   (i32.const 0)
+   (i32.const 0)
+   (i32.const 0)
+  )
+ )
+
+ ;; PRIMARY:      (func $array.new_data
+ ;; PRIMARY-NEXT:  (drop
+ ;; PRIMARY-NEXT:   (array.new_data $data-array $data
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:   )
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $array.new_data
+  (drop
+   (array.new_data $data-array $data
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+ )
+
+ ;; PRIMARY:      (func $array.new_elem
+ ;; PRIMARY-NEXT:  (drop
+ ;; PRIMARY-NEXT:   (array.new_elem $elem-array $elem
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:   )
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $array.new_elem
+  (drop
+   (array.new_elem $elem-array $elem
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+ )
+
+ ;; SECONDARY:      (type $0 (func))
+
+ ;; SECONDARY:      (import "primary" "memory" (memory $mem 0))
+
+ ;; SECONDARY:      (func $no-segment
+ ;; SECONDARY-NEXT:  (nop)
+ ;; SECONDARY-NEXT: )
+ (func $no-segment
+  (nop)
+ )
+)