Propagate to exact types in struct-utils.h (#7889)

Key collected struct information on both heap type and exactness,
allowing queries for exact types to return more precise results than
queries on the corresponding inexact types.

Use this to fix a bug in CFP where it failed to take into account
exactness and would unnecessarily and incorrectly emit a select between
two values of different types where a single exact type was expected.

Author: tlively

src/ir/struct-utils.h

@@ -19,6 +19,7 @@
 
 #include "ir/properties.h"
 #include "ir/subtypes.h"
+#include "wasm-type.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -83,20 +84,23 @@ template<typename T> struct StructValues : public std::vector<T> {
   T desc;
 };
 
-// Maps heap types to a StructValues for that heap type.
+// Maps heap types to a StructValues for that heap type. Includes exactness in
+// the key to allow differentiating between values for exact and inexact
+// references to each type.
 //
 // Also provides a combineInto() helper that combines one map into another. This
 // depends on the underlying T defining a combine() method.
 template<typename T>
-struct StructValuesMap : public std::unordered_map<HeapType, StructValues<T>> {
+struct StructValuesMap
+  : public std::unordered_map<std::pair<HeapType, Exactness>, StructValues<T>> {
   // When we access an item, if it does not already exist, create it with a
   // vector of the right length for that type.
-  StructValues<T>& operator[](HeapType type) {
-    assert(type.isStruct());
+  StructValues<T>& operator[](std::pair<HeapType, Exactness> type) {
+    assert(type.first.isStruct());
     auto inserted = this->insert({type, {}});
     auto& values = inserted.first->second;
     if (inserted.second) {
-      values.resize(type.getStruct().fields.size());
+      values.resize(type.first.getStruct().fields.size());
     }
     return values;
   }
@@ -113,7 +117,8 @@ struct StructValuesMap : public std::unordered_map<HeapType, StructValues<T>> {
   void dump(std::ostream& o) {
     o << "dump " << this << '\n';
     for (auto& [type, vec] : (*this)) {
-      o << "dump " << type << " " << &vec << ' ';
+      o << "dump " << type.first << (type.second == Exact ? " exact " : " ")
+        << &vec << ' ';
       for (auto x : vec) {
         x.dump(o);
         o << " ";
@@ -203,7 +208,8 @@ struct StructScanner
     // Note writes to all the fields of the struct.
     auto heapType = type.getHeapType();
     auto& fields = heapType.getStruct().fields;
-    auto& infos = functionNewInfos[this->getFunction()][heapType];
+    auto ht = std::make_pair(heapType, Exact);
+    auto& infos = functionNewInfos[this->getFunction()][ht];
     for (Index i = 0; i < fields.size(); i++) {
       if (curr->isWithDefault()) {
         self().noteDefault(fields[i].type, heapType, i, infos[i]);
@@ -224,11 +230,12 @@ struct StructScanner
     }
 
     // Note a write to this field of the struct.
-    noteExpressionOrCopy(curr->value,
-                         type.getHeapType(),
-                         curr->index,
-                         functionSetGetInfos[this->getFunction()]
-                                            [type.getHeapType()][curr->index]);
+    auto ht = std::make_pair(type.getHeapType(), type.getExactness());
+    noteExpressionOrCopy(
+      curr->value,
+      type.getHeapType(),
+      curr->index,
+      functionSetGetInfos[this->getFunction()][ht][curr->index]);
   }
 
   void visitStructGet(StructGet* curr) {
@@ -237,11 +244,11 @@ struct StructScanner
       return;
     }
 
-    auto heapType = type.getHeapType();
+    auto ht = std::make_pair(type.getHeapType(), type.getExactness());
     auto index = curr->index;
-    self().noteRead(heapType,
+    self().noteRead(type.getHeapType(),
                     index,
-                    functionSetGetInfos[this->getFunction()][heapType][index]);
+                    functionSetGetInfos[this->getFunction()][ht][index]);
   }
 
   void visitStructRMW(StructRMW* curr) {
@@ -251,9 +258,9 @@ struct StructScanner
     }
 
     auto heapType = type.getHeapType();
+    auto ht = std::make_pair(heapType, type.getExactness());
     auto index = curr->index;
-    auto& info =
-      functionSetGetInfos[this->getFunction()][type.getHeapType()][index];
+    auto& info = functionSetGetInfos[this->getFunction()][ht][index];
 
     if (curr->op == RMWXchg) {
       // An xchg is really like a read and write combined.
@@ -274,9 +281,9 @@ struct StructScanner
     }
 
     auto heapType = type.getHeapType();
+    auto ht = std::make_pair(heapType, type.getExactness());
     auto index = curr->index;
-    auto& info =
-      functionSetGetInfos[this->getFunction()][type.getHeapType()][curr->index];
+    auto& info = functionSetGetInfos[this->getFunction()][ht][index];
 
     // A cmpxchg is like a read and conditional write.
     self().noteRead(heapType, index, info);
@@ -310,11 +317,12 @@ struct StructScanner
       return;
     }
     auto heapType = type.getHeapType();
+    auto ht = std::make_pair(heapType, type.getExactness());
     if (heapType.isStruct()) {
       // Any subtype of the reference here may be read from.
       self().noteRead(heapType,
                       DescriptorIndex,
-                      functionSetGetInfos[this->getFunction()][heapType].desc);
+                      functionSetGetInfos[this->getFunction()][ht].desc);
       return;
     }
   }
@@ -399,46 +407,58 @@ template<typename T> class TypeHierarchyPropagator {
   void propagate(StructValuesMap<T>& combinedInfos,
                  bool toSubTypes,
                  bool toSuperTypes) {
-    UniqueDeferredQueue<HeapType> work;
-    for (auto& [type, _] : combinedInfos) {
-      work.push(type);
+    UniqueDeferredQueue<std::pair<HeapType, Exactness>> work;
+    for (auto& [ht, _] : combinedInfos) {
+      work.push(ht);
     }
     while (!work.empty()) {
-      auto type = work.pop();
-      auto& infos = combinedInfos[type];
+      auto [type, exactness] = work.pop();
+      auto& infos = combinedInfos[{type, exactness}];
 
       if (toSuperTypes) {
-        // Propagate shared fields to the supertype.
-        if (auto superType = type.getDeclaredSuperType()) {
-          auto& superInfos = combinedInfos[*superType];
-          auto& superFields = superType->getStruct().fields;
-          for (Index i = 0; i < superFields.size(); i++) {
+        // Propagate shared fields to the supertype, which may be the inexact
+        // version of the same type.
+        std::optional<std::pair<HeapType, Exactness>> super;
+        if (exactness == Exact) {
+          super = {type, Inexact};
+        } else if (auto superType = type.getDeclaredSuperType()) {
+          super = {*superType, Inexact};
+        }
+        if (super) {
+          auto& superInfos = combinedInfos[*super];
+          const auto& superFields = &super->first.getStruct().fields;
+          for (Index i = 0; i < superFields->size(); i++) {
             if (superInfos[i].combine(infos[i])) {
-              work.push(*superType);
+              work.push(*super);
             }
           }
           // Propagate the descriptor to the super, if the super has one.
-          if (superType->getDescriptorType() &&
+          if (super->first.getDescriptorType() &&
               superInfos.desc.combine(infos.desc)) {
-            work.push(*superType);
+            work.push(*super);
           }
         }
       }
 
-      if (toSubTypes) {
-        // Propagate shared fields to the subtypes.
+      if (toSubTypes && exactness == Inexact) {
+        // Propagate shared fields to the subtypes, which may just be the exact
+        // version of the same type.
         auto numFields = type.getStruct().fields.size();
-        for (auto subType : subTypes.getImmediateSubTypes(type)) {
-          auto& subInfos = combinedInfos[subType];
+        auto handleSubtype = [&](std::pair<HeapType, Exactness> sub) {
+          auto& subInfos = combinedInfos[sub];
           for (Index i = 0; i < numFields; i++) {
             if (subInfos[i].combine(infos[i])) {
-              work.push(subType);
+              work.push(sub);
             }
           }
           // Propagate the descriptor.
           if (subInfos.desc.combine(infos.desc)) {
-            work.push(subType);
+            work.push(sub);
           }
+        };
+        handleSubtype({type, Exact});
+        for (auto subType : subTypes.getImmediateSubTypes(type)) {
+          handleSubtype({subType, Inexact});
         }
       }
     }

src/passes/ConstantFieldPropagation.cpp

@@ -53,7 +53,6 @@
 
 #include "ir/bits.h"
 #include "ir/gc-type-utils.h"
-#include "ir/module-utils.h"
 #include "ir/possible-constant.h"
 #include "ir/struct-utils.h"
 #include "ir/utils.h"
@@ -115,14 +114,8 @@ struct FunctionOptimizer : public WalkerPass<PostWalker<FunctionOptimizer>> {
   }
 
   PossibleConstantValues getInfo(HeapType type, Index index, Exactness exact) {
-    // If the reference is exact and the field immutable, then we are reading
-    // exactly what was written to struct.news and nothing else.
-    auto mutable_ = index == StructUtils::DescriptorIndex
-                      ? Immutable
-                      : GCTypeUtils::getField(type, index)->mutable_;
-    auto& infos =
-      (exact == Inexact || mutable_ == Mutable) ? propagatedInfos : rawNewInfos;
-    if (auto it = infos.find(type); it != infos.end()) {
+    if (auto it = propagatedInfos.find({type, exact});
+        it != propagatedInfos.end()) {
       // There is information on this type, fetch it.
       return it->second[index];
     }
@@ -290,7 +283,7 @@ struct FunctionOptimizer : public WalkerPass<PostWalker<FunctionOptimizer>> {
         return;
       }
 
-      auto iter = rawNewInfos.find(type);
+      auto iter = rawNewInfos.find({type, Exact});
       if (iter == rawNewInfos.end()) {
         // This type has no struct.news, so we can ignore it: it is abstract.
         return;
@@ -454,7 +447,8 @@ struct PCVScanner
   void noteCopy(HeapType type, Index index, PossibleConstantValues& info) {
     // Note copies, as they must be considered later. See the comment on the
     // propagation of values below.
-    functionCopyInfos[getFunction()][type][index] = true;
+    // TODO: Take into account exactness here.
+    functionCopyInfos[getFunction()][{type, Inexact}][index] = true;
   }
 
   void noteRead(HeapType type, Index index, PossibleConstantValues& info) {

src/passes/GlobalTypeOptimization.cpp

@@ -229,8 +229,13 @@ struct GlobalTypeOptimization : public Pass {
         continue;
       }
       auto& fields = type.getStruct().fields;
-      auto& dataFromSubsAndSupers = dataFromSubsAndSupersMap[type];
-      auto& dataFromSupers = dataFromSupersMap[type];
+      // Use the exact entry because information from the inexact entry in
+      // dataFromSupersMap will have been propagated down into it but not vice
+      // versa. (This doesn't matter or dataFromSubsAndSupers because the exact
+      // and inexact entries will have the same data.)
+      auto ht = std::make_pair(type, Exact);
+      auto& dataFromSubsAndSupers = dataFromSubsAndSupersMap[ht];
+      auto& dataFromSupers = dataFromSupersMap[ht];
 
       // Process immutability.
       for (Index i = 0; i < fields.size(); i++) {

src/passes/TypeRefining.cpp

@@ -193,7 +193,8 @@ struct TypeRefining : public Pass {
     for (auto type : allTypes) {
       if (type.isStruct()) {
         auto& fields = type.getStruct().fields;
-        auto& infos = finalInfos[type];
+        // Update the inexact entry because that's what we will query later.
+        auto& infos = finalInfos[{type, Inexact}];
         for (Index i = 0; i < fields.size(); i++) {
           auto gufaType = oracle.getContents(DataLocation{type, i}).getType();
           // Do not introduce new exact fields that might requires invalid
@@ -223,7 +224,7 @@ struct TypeRefining : public Pass {
         }
 
         auto type = structNew->type.getHeapType();
-        auto& infos = finalInfos[type];
+        auto& infos = finalInfos[{type, Inexact}];
         auto& fields = type.getStruct().fields;
         for (Index i = 0; i < fields.size(); i++) {
           // We are in a situation like this:
@@ -287,7 +288,9 @@ struct TypeRefining : public Pass {
       auto& fields = type.getStruct().fields;
       for (Index i = 0; i < fields.size(); i++) {
         auto oldType = fields[i].type;
-        auto& info = finalInfos[type][i];
+        // Use inexact because exact info will have been propagated up to
+        // inexact entries but not necessarily vice versa.
+        auto& info = finalInfos[{type, Inexact}][i];
         if (!info.noted()) {
           info = LUBFinder(oldType);
         }
@@ -301,11 +304,11 @@ struct TypeRefining : public Pass {
           // public, unchanged since we cannot optimize it
           Type newSuperType;
           if (!publicTypesSet.count(*super)) {
-            newSuperType = finalInfos[*super][i].getLUB();
+            newSuperType = finalInfos[{*super, Inexact}][i].getLUB();
           } else {
             newSuperType = superFields[i].type;
           }
-          auto& info = finalInfos[type][i];
+          auto& info = finalInfos[{type, Inexact}][i];
           auto newType = info.getLUB();
           if (!Type::isSubType(newType, newSuperType)) {
             // To ensure we are a subtype of the super's field, simply copy that
@@ -340,7 +343,7 @@ struct TypeRefining : public Pass {
       // After all those decisions, see if we found anything to optimize.
       for (Index i = 0; i < fields.size(); i++) {
         auto oldType = fields[i].type;
-        auto& lub = finalInfos[type][i];
+        auto& lub = finalInfos[{type, Inexact}][i];
         auto newType = lub.getLUB();
         if (newType != oldType) {
           canOptimize = true;
@@ -384,7 +387,8 @@ struct TypeRefining : public Pass {
         Type newFieldType;
         if (!curr->ref->type.isNull()) {
           auto oldType = curr->ref->type.getHeapType();
-          newFieldType = parent.finalInfos[oldType][curr->index].getLUB();
+          newFieldType =
+            parent.finalInfos[{oldType, Inexact}][curr->index].getLUB();
         }
 
         if (curr->ref->type.isNull() || newFieldType == Type::unreachable ||
@@ -449,7 +453,8 @@ struct TypeRefining : public Pass {
           if (!oldType.isRef()) {
             continue;
           }
-          auto newType = parent.finalInfos[oldStructType][i].getLUB();
+          auto newType =
+            parent.finalInfos[{oldStructType, Inexact}][i].getLUB();
           newFields[i].type = getTempType(newType);
         }
       }

test/lit/passes/cfp-reftest-desc.wast

@@ -161,7 +161,7 @@
     )
     (drop
       (struct.new_default $sub
-      (global.get $B)
+        (global.get $B)
       )
     )
     ;; We read from an exact $super here, so the type of the ref.get_desc is
@@ -185,4 +185,3 @@
     (i32.const 42)
   )
 )
-

test/lit/passes/cfp-reftest.wast

@@ -1510,4 +1510,3 @@
     )
   )
 )
-

test/lit/passes/cfp.wast

@@ -2546,8 +2546,13 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (struct.get $A 0
-  ;; CHECK-NEXT:    (local.get $A-exact)
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (local.get $A-exact)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 10)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
@@ -2591,7 +2596,7 @@
         (local.get $B)
       )
     )
-    ;; We should be able to optimize both exact references TODO.
+    ;; We should be able to optimize both exact references.
     (drop
       (struct.get $A 0
         (local.get $A-exact)