RemoveUnusedModuleElements: Handle null table segments

src/passes/RemoveUnusedModuleElements.cpp

@@ -49,6 +49,7 @@
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/insert_ordered.h"
+#include "support/space.h"
 #include "support/stdckdint.h"
 #include "support/utilities.h"
 #include "wasm-builder.h"
@@ -295,6 +296,19 @@ struct Analyzer {
 
     // All elems, regardless of type.
     std::unordered_set<Name> allElems;
+
+    // Whether this table's elems may overlap, i.e., trample each other. In such
+    // a case, we can optimize less: if segment A writes a function, and B
+    // tramples it with a null or with a function of another type, then if we
+    // call that index with the right type for A, the trampling causes a trap -
+    // so we cannot remove the trampling segment, even though it has nothing
+    // can can be called, forcing us to keep it just to preserve the trap.
+    //
+    // We could be more precise here, as currently we "give up" on analyzing
+    // exact overlap, and just stop removing element segments. But such an exact
+    // analysis would be complex, and is rarely needed: when trapsNeverHappen
+    // this is not an issue, and also segment overlap is rare.
+    bool mayHaveOverlappingElems = false;
   };
   std::unordered_map<Name, FlatTableInfo> flatTableInfoMap;
 
@@ -316,6 +330,10 @@ struct Analyzer {
   }
 
   void prepare() {
+    // The spans of element segments for each table. We use this to compute
+    // |mayHaveOverlappingElems|.
+    std::unordered_map<Name, DisjointSpans> tableElemSpans;
+
     auto notePossibleFunc =
       [&](FlatTableInfo& info, Expression* item, Name elem = Name()) {
         if (auto* refFunc = item->dynCast<RefFunc>()) {
@@ -345,6 +363,34 @@ struct Analyzer {
       info.allElems.insert(elem->name);
     }
 
+    // Compute elem overlap.
+    for (auto& elem : module->elementSegments) {
+      if (elem->isPassive()) {
+        continue;
+      }
+
+      auto& info = flatTableInfoMap[elem->table];
+
+      // There is at least one elem, this one.
+      assert(!info.allElems.empty());
+      if (info.allElems.size() <= 1) {
+        // This table has just this one elem, so no overlap is possible, even
+        // if the single elem has a non-constant offset.
+        continue;
+      }
+
+      if (auto* c = elem->offset->dynCast<Const>()) {
+        auto start = c->value.getUnsigned();
+        auto end = start + elem->data.size();
+        if (tableElemSpans[elem->table].addAndCheckOverlap({start, end})) {
+          info.mayHaveOverlappingElems = true;
+        }
+      } else {
+        // Offset is not constant, so it might overlap with others.
+        info.mayHaveOverlappingElems = true;
+      }
+    }
+
     // If a table has an initial value, it is callable as well.
     for (auto& table : module->tables) {
       if (table->init) {
@@ -461,11 +507,17 @@ struct Analyzer {
     // if it has the right type, we would not trap. (Note that we were using the
     // fact that we do not distinguish types of traps, in the case without an
     // initial value - we turned a trap on the wrong type to one on a null.)
-    if (!options.trapsNeverHappen && module->getTable(table)->init) {
-      // We only need to reference the elem, so it writes its functions - the
-      // functions are not actually called, as we just trap.
-      for (auto& elem : info.allElems) {
-        reference({ModuleElementKind::ElementSegment, elem});
+    if (!options.trapsNeverHappen) {
+      // A related situation is a table with element segments that trample
+      // valid values with nulls or functions of the wrong type for a call:
+      // when they overwrite a valid function, they cause a trap, which we
+      // must preserve if traps can.
+      if (module->getTable(table)->init || info.mayHaveOverlappingElems) {
+        // We only need to reference the elem, so it writes its functions - the
+        // functions are not actually called, as we just trap.
+        for (auto& elem : info.allElems) {
+          reference({ModuleElementKind::ElementSegment, elem});
+        }
       }
     }
 

test/lit/passes/remove-unused-module-elements-closed-tnh.wast

@@ -0,0 +1,131 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --remove-unused-module-elements --closed-world -tnh -all -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --remove-unused-module-elements                     -all -S -o - | filecheck %s --check-prefix OPEN_WORLD
+
+;; A table with a single elem with non-constant offset that we can remove.
+;; Usually a non-constant segment suggests it might overlap with others, which
+;; causes us to keep all elems, but here there is just one, so we can optimize.
+;; We require closed world (to know about which calls can reach it - the only
+;; call is to type $func, so elem $elem which contains $other is not needed).
+;; We also require traps-never-happen, as an elem with non-constant offset can
+;; trap, normally. When both closed world and tnh, we remove elem $elem.
+(module
+  ;; TNH__:      (type $func (func))
+  ;; CHECK:      (type $func (func))
+  ;; OPEN_WORLD:      (type $func (func))
+  (type $func (func))
+
+  ;; OPEN_WORLD:      (type $other (func (result i32)))
+  (type $other (func (result i32)))
+
+  ;; OPEN_WORLD:      (import "a" "b" (global $offset i32))
+  (import "a" "b" (global $offset i32))
+
+  ;; TNH__:      (table $table 6 6 funcref)
+  ;; CHECK:      (table $table 6 6 funcref)
+  ;; OPEN_WORLD:      (table $table 6 6 funcref)
+  (table $table 6 6 funcref)
+
+  ;; OPEN_WORLD:      (elem $elem (global.get $offset) $other)
+  (elem $elem (global.get $offset) $other)
+
+  ;; OPEN_WORLD:      (export "export" (func $export))
+
+  ;; OPEN_WORLD:      (func $other (type $other) (result i32)
+  ;; OPEN_WORLD-NEXT:  (i32.const 42)
+  ;; OPEN_WORLD-NEXT: )
+  (func $other (type $other) (result i32)
+    (i32.const 42)
+  )
+
+  ;; TNH__:      (export "export" (func $export))
+
+  ;; TNH__:      (func $export (type $func)
+  ;; TNH__-NEXT:  (call_indirect $table (type $func)
+  ;; TNH__-NEXT:   (i32.const 0)
+  ;; TNH__-NEXT:  )
+  ;; TNH__-NEXT: )
+  ;; CHECK:      (export "export" (func $export))
+
+  ;; CHECK:      (func $export (type $func)
+  ;; CHECK-NEXT:  (call_indirect $table (type $func)
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; OPEN_WORLD:      (func $export (type $func)
+  ;; OPEN_WORLD-NEXT:  (call_indirect $table (type $func)
+  ;; OPEN_WORLD-NEXT:   (i32.const 0)
+  ;; OPEN_WORLD-NEXT:  )
+  ;; OPEN_WORLD-NEXT: )
+  (func $export (export "export")
+    ;; Call with type $func, but the table contains $other.
+    (call_indirect $table (type $func)
+      (i32.const 0)
+    )
+  )
+)
+
+;; As above, but now we have two segments, one constant. We don't know if they
+;; overlap, so we normally assume the worst. However, with tnh, we can assume no
+;; traps, and remove elem $elem. (We cannot remove elem $second in either case,
+;; as it contains a function we have an indirect call for its type.)
+(module
+  ;; TNH__:      (type $func (func))
+  ;; CHECK:      (type $func (func))
+  ;; OPEN_WORLD:      (type $func (func))
+  (type $func (func))
+
+  ;; OPEN_WORLD:      (type $other (func (result i32)))
+  (type $other (func (result i32)))
+
+  ;; OPEN_WORLD:      (import "a" "b" (global $offset i32))
+  (import "a" "b" (global $offset i32))
+
+  ;; TNH__:      (table $table 6 6 funcref)
+  ;; CHECK:      (table $table 6 6 funcref)
+  ;; OPEN_WORLD:      (table $table 6 6 funcref)
+  (table $table 6 6 funcref)
+
+  ;; OPEN_WORLD:      (elem $elem (global.get $offset) $other)
+  (elem $elem (global.get $offset) $other)
+
+  ;; CHECK:      (elem $second (i32.const 0) $export)
+  ;; OPEN_WORLD:      (elem $second (i32.const 0) $export)
+  (elem $second (i32.const 0) $export)
+
+  ;; OPEN_WORLD:      (export "export" (func $export))
+
+  ;; OPEN_WORLD:      (func $other (type $other) (result i32)
+  ;; OPEN_WORLD-NEXT:  (i32.const 42)
+  ;; OPEN_WORLD-NEXT: )
+  (func $other (type $other) (result i32)
+    (i32.const 42)
+  )
+
+  ;; TNH__:      (export "export" (func $export))
+
+  ;; TNH__:      (func $export (type $func)
+  ;; TNH__-NEXT:  (call_indirect $table (type $func)
+  ;; TNH__-NEXT:   (i32.const 0)
+  ;; TNH__-NEXT:  )
+  ;; TNH__-NEXT: )
+  ;; CHECK:      (export "export" (func $export))
+
+  ;; CHECK:      (func $export (type $func)
+  ;; CHECK-NEXT:  (call_indirect $table (type $func)
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; OPEN_WORLD:      (func $export (type $func)
+  ;; OPEN_WORLD-NEXT:  (call_indirect $table (type $func)
+  ;; OPEN_WORLD-NEXT:   (i32.const 0)
+  ;; OPEN_WORLD-NEXT:  )
+  ;; OPEN_WORLD-NEXT: )
+  (func $export (export "export")
+    (call_indirect $table (type $func)
+      (i32.const 0)
+    )
+  )
+)
+