Track indirect call types in RemoveUnusedModuleElements (#7728)

An indirect call to a type in a table now only forces functions of
that type to be marked as used: functions of other types are
left alone, potentially leaving them unreached.

This is more precise than assuming any indirect call can
reach anywhere, which is more or less what we did before.

The downside is that this makes the pass is 10% slower. This is one
of our faster passes, however, so the cost seems worth it, given
that this can noticeably help in some cases, e.g., 6.5% code size
reduction on Poppler, and 20% on an embind testcase. In
general, however, most real-world codebases benefit little if at
all, because there is too much overlap in indirect call signatures:
Statistically, when you generate random graphs of size `n` and
chance for each edge to exist `p`, then even if `p` decreases to 0 the
graph will tend to end up fully connected [1]. And, in wasm, `p`
does not even decrease to 0:

* Consider some common signature like `{i32} -> {}` (i32 param, no
result).
* In real-world code, there is some chance `q>0` for that signature to
be called,
  and some chance `r>0` for that signature to exist in the code.
* `p >= O(rq) > 0` because all it takes for a connection to exist is
that that signature exists on one side and is called on the other.

That is, in large codebases there is an overlap in signatures, and
statistically this means that all the code will end up reachable, at
least in the limit. But some programs do get lucky and benefit
from this PR.

[1]
https://en.wikipedia.org/wiki/Erd%C5%91s%E2%80%93R%C3%A9nyi_model#Properties_of_G(n,_p)

Author: kripken

src/passes/RemoveUnusedModuleElements.cpp

@@ -24,7 +24,8 @@
 //  * No references at all. We can simply remove it.
 //  * References, but no uses. We can't remove it, but we can change it (see
 //    below).
-//  * Uses (which imply references). We must keep it as it is.
+//  * Uses (which imply references). We must keep it as it is, because it is
+//    fully used (e.g. for a function, it is called and may execute).
 //
 // An example of something with a reference but *not* a use is a RefFunc to a
 // function that has no corresponding CallRef to that type. We cannot just
@@ -62,25 +63,38 @@ using ModuleElementKind = ModuleItemKind;
 // name of the particular element.
 using ModuleElement = std::pair<ModuleElementKind, Name>;
 
+// Information from an indirect call: the name of the table, and the heap type.
+using IndirectCall = std::pair<Name, HeapType>;
+
 // Visit or walk an expression to find what things are referenced.
 struct ReferenceFinder
   : public PostWalker<ReferenceFinder,
                       UnifiedExpressionVisitor<ReferenceFinder>> {
   // Our findings are placed in these data structures, which the user of this
-  // code can then process.
-  std::vector<ModuleElement> elements;
+  // code can then process. We mark both uses and references, and also note
+  // uses of specific things that require special handling, like refFuncs.
+  std::vector<ModuleElement> used, referenced;
   std::vector<HeapType> callRefTypes;
   std::vector<Name> refFuncs;
   std::vector<StructField> structFields;
+  std::vector<IndirectCall> indirectCalls;
 
   // Add an item to the output data structures.
-  void note(ModuleElement element) { elements.push_back(element); }
-  void noteCallRef(HeapType type) { callRefTypes.push_back(type); }
-  void noteRefFunc(Name refFunc) { refFuncs.push_back(refFunc); }
-  void note(StructField structField) { structFields.push_back(structField); }
-
-  // Generic visitor
+  void use(ModuleElement element) { used.push_back(element); }
+  void reference(ModuleElement element) { referenced.push_back(element); }
+  void useCallRef(HeapType type) { callRefTypes.push_back(type); }
+  void useRefFunc(Name refFunc) { refFuncs.push_back(refFunc); }
+  void useStructField(StructField structField) {
+    structFields.push_back(structField);
+  }
+  void useIndirectCall(Name table, HeapType type) {
+    indirectCalls.push_back({table, type});
+  }
 
+  // Generic visitor: Use all the things referenced. This handles e.g. using the
+  // table of a table.get. When we do not want such unconditional use, we
+  // override (e.g. for call_indirect, we don't want to mark the entire table as
+  // used, see below).
   void visitExpression(Expression* curr) {
 #define DELEGATE_ID curr->_id
 
@@ -101,7 +115,7 @@ struct ReferenceFinder
 
 #define DELEGATE_FIELD_NAME_KIND(id, field, kind)                              \
   if (cast->field.is()) {                                                      \
-    note({kind, cast->field});                                                 \
+    use({kind, cast->field});                                                  \
   }
 
 #include "wasm-delegations-fields.def"
@@ -110,7 +124,7 @@ struct ReferenceFinder
   // Specific visitors
 
   void visitCall(Call* curr) {
-    note({ModuleElementKind::Function, curr->target});
+    use({ModuleElementKind::Function, curr->target});
 
     if (Intrinsics(*getModule()).isCallWithoutEffects(curr)) {
       // A call-without-effects receives a function reference and calls it, the
@@ -137,12 +151,15 @@ struct ReferenceFinder
   }
 
   void visitCallIndirect(CallIndirect* curr) {
-    note({ModuleElementKind::Table, curr->table});
+    // We refer to the table, but may not use all parts of it, that depends on
+    // the heap type we call with.
+    reference({ModuleElementKind::Table, curr->table});
+    useIndirectCall(curr->table, curr->heapType);
     // Note a possible call of a function reference as well, as something might
     // be written into the table during runtime. With precise tracking of what
     // is written into the table we could do better here; we could also see
     // which tables are immutable. TODO
-    noteCallRef(curr->heapType);
+    useCallRef(curr->heapType);
   }
 
   void visitCallRef(CallRef* curr) {
@@ -151,17 +168,17 @@ struct ReferenceFinder
       return;
     }
 
-    noteCallRef(curr->target->type.getHeapType());
+    useCallRef(curr->target->type.getHeapType());
   }
 
-  void visitRefFunc(RefFunc* curr) { noteRefFunc(curr->func); }
+  void visitRefFunc(RefFunc* curr) { useRefFunc(curr->func); }
 
   void visitStructGet(StructGet* curr) {
     if (curr->ref->type == Type::unreachable || curr->ref->type.isNull()) {
       return;
     }
     auto type = curr->ref->type.getHeapType();
-    note(StructField{type, curr->index});
+    useStructField(StructField{type, curr->index});
   }
 };
 
@@ -262,9 +279,12 @@ struct Analyzer {
       ReferenceFinder finder;
       finder.setModule(module);
       finder.visit(curr);
-      for (auto element : finder.elements) {
+      for (auto element : finder.used) {
         use(element);
       }
+      for (auto element : finder.referenced) {
+        reference(element);
+      }
       for (auto type : finder.callRefTypes) {
         useCallRefType(type);
       }
@@ -274,6 +294,9 @@ struct Analyzer {
       for (auto structField : finder.structFields) {
         useStructField(structField);
       }
+      for (auto call : finder.indirectCalls) {
+        useIndirectCall(call);
+      }
 
       // Scan the children to continue our work.
       scanChildren(curr);
@@ -316,6 +339,37 @@ struct Analyzer {
     }
   }
 
+  std::unordered_set<IndirectCall> usedIndirectCalls;
+
+  void useIndirectCall(IndirectCall call) {
+    auto [_, inserted] = usedIndirectCalls.insert(call);
+    if (!inserted) {
+      return;
+    }
+
+    // TODO: use structured bindings with c++20, needed for the capture below
+    auto table = call.first;
+    auto type = call.second;
+
+    // Any function in the table of that signature may be called.
+    ModuleUtils::iterTableSegments(
+      *module, table, [&](ElementSegment* segment) {
+        auto segmentReferenced = false;
+        for (auto* item : segment->data) {
+          if (auto* refFunc = item->dynCast<RefFunc>()) {
+            auto* func = module->getFunction(refFunc->func);
+            if (HeapType::isSubType(func->type, type)) {
+              use({ModuleElementKind::Function, refFunc->func});
+              segmentReferenced = true;
+            }
+          }
+        }
+        if (segmentReferenced) {
+          reference({ModuleElementKind::ElementSegment, segment->name});
+        }
+      });
+  }
+
   void useRefFunc(Name func) {
     if (!options.closedWorld) {
       // The world is open, so assume the worst and something (inside or outside
@@ -341,7 +395,7 @@ struct Analyzer {
       // We've never seen a CallRef for this, but might see one later.
       uncalledRefFuncMap[type].insert(func);
 
-      referenced.insert(element);
+      reference(element);
     }
   }
 
@@ -554,34 +608,11 @@ struct Analyzer {
     finder.setModule(module);
     finder.walk(curr);
 
-    for (auto element : finder.elements) {
-      // Avoid repeated work. Note that globals with multiple references to
-      // previous globals can lead to exponential work, so this is important.
-      // (If C refers twice to B, and B refers twice to A, then when we process
-      // C we would, naively, scan B twice and A four times.)
-      auto [_, inserted] = referenced.insert(element);
-      if (!inserted) {
-        continue;
-      }
-
-      auto& [kind, value] = element;
-      if (kind == ModuleElementKind::Global) {
-        // Like functions, (non-imported) globals have contents. For functions,
-        // things are simple: if a function ends up with references but no uses
-        // then we can simply empty out the function (by setting its body to an
-        // unreachable). We don't have a simple way to do the same for globals,
-        // unfortunately. For now, scan the global's contents and add references
-        // as needed.
-        // TODO: We could try to empty the global out, for example, replace it
-        //       with a null if it is nullable, or replace all gets of it with
-        //       something else, but that is not trivial.
-        auto* global = module->getGlobal(value);
-        if (!global->imported()) {
-          // Note that infinite recursion is not a danger here since a global
-          // can only refer to previous globals.
-          addReferences(global->init);
-        }
-      }
+    for (auto element : finder.used) {
+      reference(element);
+    }
+    for (auto element : finder.referenced) {
+      reference(element);
     }
 
     for (auto func : finder.refFuncs) {
@@ -594,7 +625,7 @@ struct Analyzer {
       // just adding a reference to the function, and not actually using the
       // RefFunc. (Only useRefFunc() + a CallRef of the proper type are enough
       // to make a function itself used.)
-      referenced.insert({ModuleElementKind::Function, func});
+      reference({ModuleElementKind::Function, func});
     }
 
     // Note: nothing to do with |callRefTypes| and |structFields|, which only
@@ -603,6 +634,44 @@ struct Analyzer {
     // handled in an entirely different way in Binaryen IR, and we don't need to
     // worry about it.)
   }
+
+  void reference(ModuleElement element) {
+    // Avoid repeated work. Note that globals with multiple references to
+    // previous globals can lead to exponential work, so this is important.
+    // (If C refers twice to B, and B refers twice to A, then when we process
+    // C we would, naively, scan B twice and A four times.)
+    auto [_, inserted] = referenced.insert(element);
+    if (!inserted) {
+      return;
+    }
+
+    // Some references force references to their internals, just by being
+    // referenced and present in the output.
+    auto& [kind, value] = element;
+    if (kind == ModuleElementKind::Global) {
+      // Like functions, (non-imported) globals have contents. For functions,
+      // things are simple: if a function ends up with references but no uses
+      // then we can simply empty out the function (by setting its body to an
+      // unreachable). We don't have a simple way to do the same for globals,
+      // unfortunately. For now, scan the global's contents and add references
+      // as needed.
+      // TODO: We could try to empty the global out, for example, replace it
+      //       with a null if it is nullable, or replace all gets of it with
+      //       something else, but that is not trivial.
+      auto* global = module->getGlobal(value);
+      if (!global->imported()) {
+        // Note that infinite recursion is not a danger here since a global
+        // can only refer to previous globals.
+        addReferences(global->init);
+      }
+    } else if (kind == ModuleElementKind::ElementSegment) {
+      // TODO: We could empty out parts of the segment we don't need.
+      auto* segment = module->getElementSegment(value);
+      for (auto* item : segment->data) {
+        addReferences(item);
+      }
+    }
+  }
 };
 
 struct RemoveUnusedModuleElements : public Pass {
@@ -709,13 +778,6 @@ struct RemoveUnusedModuleElements : public Pass {
         }
       });
 
-    // For now, all functions that can be called indirectly are marked as roots.
-    // TODO: Compute this based on which ElementSegments are actually used,
-    //       and which functions have a call_indirect of the proper type.
-    ElementUtils::iterAllElementFunctionNames(module, [&](Name name) {
-      roots.emplace_back(ModuleElementKind::Function, name);
-    });
-
     // Just as out-of-bound segments may cause observable traps at instantiation
     // time, so can struct.new instructions with null descriptors cause traps in
     // global or element segment initializers.

test/ctor-eval/bad-indirect-call3.wast.out

@@ -1,19 +1,9 @@
 (module
- (type $0 (func (param externref)))
- (type $1 (func))
+ (type $0 (func))
  (type $funcref_=>_none (func (param funcref)))
- (memory $0 256 256)
- (data $0 (i32.const 10) "waka waka waka waka waka")
  (table $0 1 1 funcref)
- (elem $implicit-elem (i32.const 0) $callee)
  (export "sig_mismatch" (func $sig_mismatch))
- (func $callee (type $0) (param $0 externref)
-  (i32.store8
-   (i32.const 40)
-   (i32.const 67)
-  )
- )
- (func $sig_mismatch (type $1)
+ (func $sig_mismatch (type $0)
   (call_indirect $0 (type $funcref_=>_none)
    (ref.null nofunc)
    (i32.const 0)

test/ctor-eval/basics-flatten.wast

@@ -1,6 +1,6 @@
 (module
   (type $v (func))
-  (memory 256 256)
+  (memory $m 256 256)
   ;; test flattening of multiple segments
   (data (i32.const 10) "waka ")
   (data (i32.const 15) "waka") ;; skip a byte here
@@ -10,6 +10,7 @@
   (export "test1" (func $test1))
   (export "test2" (func $test2))
   (export "test3" (func $test3))
+  (export "memory" (memory $m)) ;; export memory so we can see the flattened data
   (func $test1
     (drop (i32.const 0)) ;; no work at all, really
     (call $safe-to-call) ;; safe to call

test/ctor-eval/basics-flatten.wast.out

@@ -1,11 +1,5 @@
 (module
- (type $v (func))
- (memory $0 256 256)
+ (memory $m 256 256)
  (data $0 (i32.const 10) "nas\00\00\00aka\00yzkx waka wakm\00\00\00\00\00\00C")
- (func $call-indirect (type $v)
-  (i32.store8
-   (i32.const 40)
-   (i32.const 67)
-  )
- )
+ (export "memory" (memory $m))
 )

test/ctor-eval/basics.wast

@@ -1,12 +1,13 @@
 (module
   (type $v (func))
-  (memory 256 256)
+  (memory $m 256 256)
   (data (i32.const 10) "waka waka waka waka waka")
   (table 1 1 funcref)
   (elem (i32.const 0) $call-indirect)
   (export "test1" (func $test1))
   (export "test2" (func $test2))
   (export "test3" (func $test3))
+  (export "memory" (memory $m)) ;; export memory so we can see the updated data
   (func $test1
     (drop (i32.const 0)) ;; no work at all, really
     (call $safe-to-call) ;; safe to call

test/ctor-eval/basics.wast.out

@@ -1,11 +1,5 @@
 (module
- (type $v (func))
- (memory $0 256 256)
+ (memory $m 256 256)
  (data $0 (i32.const 10) "nas\00\00\00aka yzkx waka wakm\00\00\00\00\00\00C")
- (func $call-indirect (type $v)
-  (i32.store8
-   (i32.const 40)
-   (i32.const 67)
-  )
- )
+ (export "memory" (memory $m))
 )

test/ctor-eval/indirect-call3.wast

@@ -1,11 +1,12 @@
 (module
   (import "env" "_abort" (func $_abort))
   (type $v (func))
-  (memory 256 256)
+  (memory $m 256 256)
   (data (i32.const 10) "waka waka waka waka waka")
   (table 2 2 funcref)
   (elem (i32.const 0) $_abort $call-indirect)
   (export "test1" (func $test1))
+  (export "memory" (memory $m)) ;; export memory so we can see the updated data
   (func $test1
     (call_indirect (type $v) (i32.const 1)) ;; safe to call
     (i32.store8 (i32.const 20) (i32.const 120))