Oracle chaining sketch

src/ir/possible-contents.cpp

@@ -1372,38 +1372,41 @@ struct TNHInfo {
   std::unordered_map<Expression*, PossibleContents> inferences;
 };
 
-class TNHOracle : public ModuleUtils::ParallelFunctionAnalysis<TNHInfo> {
+class TNHOracleAnalysis
+  : public ModuleUtils::ParallelFunctionAnalysis<TNHInfo> {
   const PassOptions& options;
+  std::shared_ptr<const Oracle> prevOracle;
 
 public:
   using Parent = ModuleUtils::ParallelFunctionAnalysis<TNHInfo>;
-  TNHOracle(Module& wasm, const PassOptions& options)
+  TNHOracleAnalysis(Module& wasm,
+                    const PassOptions& options,
+                    std::shared_ptr<const Oracle> prevOracle)
     : Parent(wasm,
              [this, &options](Function* func, TNHInfo& info) {
                scan(func, info, options);
              }),
-      options(options) {
+      options(options), prevOracle(prevOracle) {
 
     // After the scanning phase that we run in the constructor, continue to the
     // second phase of analysis: inference.
     infer();
   }
 
   // Get the type we inferred was possible at a location.
-  PossibleContents getContents(Expression* curr) {
-    auto naiveContents = PossibleContents::fullConeType(curr->type);
-
-    // If we inferred nothing, use the naive type.
-    auto iter = inferences.find(curr);
-    if (iter == inferences.end()) {
-      return naiveContents;
+  PossibleContents getContents(Expression* curr) const {
+    if (auto iter = inferences.find(curr); iter != inferences.end()) {
+      // We only store useful contents that improve on the naive estimate that
+      // uses the type in the IR.
+      [[maybe_unused]] auto naiveContents =
+        PossibleContents::fullConeType(curr->type);
+      auto contents = iter->second;
+      assert(contents != naiveContents);
+      return contents;
     }
 
-    auto& contents = iter->second;
-    // We only store useful contents that improve on the naive estimate that
-    // uses the type in the IR.
-    assert(contents != naiveContents);
-    return contents;
+    // We inferred nothing ourselves, so defer to the previous oracle.
+    return prevOracle->getExprContents(curr);
   }
 
 private:
@@ -1424,11 +1427,38 @@ class TNHOracle : public ModuleUtils::ParallelFunctionAnalysis<TNHInfo> {
                          const CastParams& targetCastParams,
                          const analysis::CFGBlockIndexes& blockIndexes,
                          TNHInfo& info);
+
+  // Get the information known by the previous oracle.
+  PossibleContents getPrevContents(Expression* expr) {
+    return prevOracle->getExprContents(expr);
+  }
 };
 
-void TNHOracle::scan(Function* func,
-                     TNHInfo& info,
-                     const PassOptions& options) {
+// Wraps around a TNHOracleAnalysis in order to provide a standard Oracle
+// interface (this avoids multiple inheritance).
+class TNHOracle : public Oracle {
+  TNHOracleAnalysis analysis;
+
+public:
+  TNHOracle(Module& wasm,
+            const PassOptions& options,
+            std::shared_ptr<const Oracle> prevOracle)
+    : Oracle(wasm, options), analysis(wasm, options, prevOracle) {}
+
+  // Get the type we inferred was possible at a location.
+  PossibleContents getContents(Location location) const override {
+    if (auto* exprLoc = std::get_if<ExpressionLocation>(&location)) {
+      return analysis.getContents(exprLoc->expr);
+    }
+
+    // We do not infer about non-expression locations yet.
+    return PossibleContents::many();
+  }
+};
+
+void TNHOracleAnalysis::scan(Function* func,
+                             TNHInfo& info,
+                             const PassOptions& options) {
   if (func->imported()) {
     return;
   }
@@ -1532,7 +1562,7 @@ void TNHOracle::scan(Function* func,
   scanner.walkFunction(func);
 }
 
-void TNHOracle::infer() {
+void TNHOracleAnalysis::infer() {
   // Phase 2: Inside each function, optimize calls based on the cast params of
   // the called function (which we noted during phase 1).
   //
@@ -1620,7 +1650,7 @@ void TNHOracle::infer() {
     }
 
     for (auto* call : info.callRefs) {
-      auto targetType = call->target->type;
+      auto targetType = getPrevContents(call->target).getType();
       if (!targetType.isRef()) {
         // This is unreachable or null, and other passes will optimize that.
         continue;
@@ -1652,7 +1682,8 @@ void TNHOracle::infer() {
         // If any of our operands will fail a cast, then we will trap.
         bool traps = false;
         for (auto& [castIndex, castType] : targetInfo.castParams) {
-          auto operandType = call->operands[castIndex]->type;
+          auto operandType =
+            getPrevContents(call->operands[castIndex]).getType();
           auto result = GCTypeUtils::evaluateCastCheck(operandType, castType);
           if (result == GCTypeUtils::Failure) {
             traps = true;
@@ -1733,11 +1764,12 @@ void TNHOracle::infer() {
   }
 }
 
-void TNHOracle::optimizeCallCasts(Expression* call,
-                                  const ExpressionList& operands,
-                                  const CastParams& targetCastParams,
-                                  const analysis::CFGBlockIndexes& blockIndexes,
-                                  TNHInfo& info) {
+void TNHOracleAnalysis::optimizeCallCasts(
+  Expression* call,
+  const ExpressionList& operands,
+  const CastParams& targetCastParams,
+  const analysis::CFGBlockIndexes& blockIndexes,
+  TNHInfo& info) {
   // Optimize in the same basic block as the call: all instructions still in
   // that block will definitely execute if the call is reached. We will do that
   // by going backwards through the call's operands and fallthrough values, and
@@ -1773,26 +1805,15 @@ void TNHOracle::optimizeCallCasts(Expression* call,
     auto* curr = operand;
     while (1) {
       // Note the type if it is useful.
-      if (castType != curr->type) {
-        // There are two constraints on this location: any value there must
-        // be of the declared type (curr->type) and also the cast type, so
-        // we know only their intersection can appear here.
-        auto declared = PossibleContents::fullConeType(curr->type);
-        auto intersection = PossibleContents::fullConeType(castType);
-        intersection.intersect(declared);
-        if (intersection.isConeType()) {
-          auto intersectionType = intersection.getType();
-          if (intersectionType != curr->type) {
-            // We inferred a more refined type.
-            info.inferences[curr] = intersection;
-          }
-        } else {
-          // Otherwise, the intersection can be a null (if the heap types are
-          // incompatible, but a null is allowed), or empty. We can apply
-          // either.
-          assert(intersection.isNull() || intersection.isNone());
-          info.inferences[curr] = intersection;
-        }
+      //
+      // There are two constraints on this location: what we know about it
+      // already (from the previous oracle), and also the cast type, so we know
+      // only their intersection can appear here.
+      auto known = getPrevContents(curr);
+      auto intersection = PossibleContents::fullConeType(castType);
+      intersection.intersect(known);
+      if (intersection != known) {
+        info.inferences[curr] = intersection;
       }
 
       auto* next = Properties::getImmediateFallthrough(curr, options, wasm);
@@ -1864,7 +1885,7 @@ struct Flower {
       // No oracle; just use the type in the IR.
       return PossibleContents::fullConeType(curr->type);
     }
-    return tnhOracle->getContents(curr);
+    return tnhOracle->getExprContents(curr);
   }
 
 private:
@@ -2045,7 +2066,12 @@ Flower::Flower(Module& wasm, const PassOptions& options)
 #ifdef POSSIBLE_CONTENTS_DEBUG
     std::cout << "tnh phase\n";
 #endif
-    tnhOracle = std::make_unique<TNHOracle>(wasm, options);
+
+    // Start with an empty base oracle, infer TNH once in the middle, and a
+    // second time.
+    auto base = std::make_shared<Oracle>(wasm, options);
+    auto mid = std::make_shared<TNHOracle>(wasm, options, base);
+    tnhOracle = std::make_unique<TNHOracle>(wasm, options, mid);
   }
 
 #ifdef POSSIBLE_CONTENTS_DEBUG

src/ir/possible-contents.h

@@ -599,6 +599,46 @@ template<> struct hash<wasm::ConeReadLocation> {
 
 namespace wasm {
 
+// A generic oracle, which provides an interface to query for contents at
+// locations.
+//
+// This is not an abstract class, but can be instantiated in order to get a
+// "empty" oracle, i.e., one that returns no insightful responses.
+//
+// A common pattern in users is to chain oracles. Each oracle can benefit from
+// the inferences in the previous one before it, and it adds new inferences on
+// top of that.
+class Oracle {
+protected:
+  Module& wasm;
+  const PassOptions& options;
+
+public:
+  Oracle(Module& wasm, const PassOptions& options)
+    : wasm(wasm), options(options) {}
+
+  virtual ~Oracle() {}
+
+  // Get the contents possible at a location.
+  virtual PossibleContents getContents(Location location) const {
+    // Nothing useful is known in this uninsightful oracle. If we have an
+    // expression then we have its type, at least.
+    // TODO: We could also repot types of other Locations.
+    if (auto* exprLoc = std::get_if<ExpressionLocation>(&location)) {
+      return PossibleContents::fullConeType(exprLoc->expr->type);
+    }
+    return PossibleContents::many();
+  }
+
+  // Helper for the common case of an expression location that is not a
+  // multivalue.
+  PossibleContents getExprContents(Expression* curr) const {
+    assert(curr->type.size() == 1);
+    auto ret = getContents(ExpressionLocation{curr, 0});
+    return ret;
+  }
+};
+
 // Analyze the entire wasm file to find which contents are possible in which
 // locations. This assumes a closed world and starts from roots - newly created
 // values - and propagates them to the locations they reach. After the
@@ -632,20 +672,17 @@ namespace wasm {
 // caller is assumed to know the wasm IR type anyhow, and also other
 // optimization passes work on the types in the IR, so we do not focus on that
 // here.
-class ContentOracle {
-  Module& wasm;
-  const PassOptions& options;
-
+class ContentOracle : public Oracle {
   void analyze();
 
 public:
   ContentOracle(Module& wasm, const PassOptions& options)
-    : wasm(wasm), options(options) {
+    : Oracle(wasm, options) {
     analyze();
   }
 
   // Get the contents possible at a location.
-  PossibleContents getContents(Location location) {
+  PossibleContents getContents(Location location) const override {
     auto iter = locationContents.find(location);
     if (iter == locationContents.end()) {
       // We know of no possible contents here.
@@ -654,14 +691,10 @@ class ContentOracle {
     return iter->second;
   }
 
-  // Helper for the common case of an expression location that is not a
-  // multivalue.
-  PossibleContents getContents(Expression* curr) {
-    assert(curr->type.size() == 1);
-    return getContents(ExpressionLocation{curr, 0});
-  }
-
 private:
+  // Maps locations to the content we inferred there. If a location is not
+  // here then no content is possible there (which is the result of the flow
+  // analysis not sending anything here at all).
   std::unordered_map<Location, PossibleContents> locationContents;
 };
 

src/passes/GUFA.cpp

@@ -97,7 +97,7 @@ struct GUFAOptimizer
   std::unordered_map<Expression*, PossibleContents> newContents;
 
   Expression* replaceCurrent(Expression* rep) {
-    newContents[rep] = oracle.getContents(getCurrent());
+    newContents[rep] = oracle.getExprContents(getCurrent());
 
     return WalkerPass<
       PostWalker<GUFAOptimizer,
@@ -111,7 +111,7 @@ struct GUFAOptimizer
       return iter->second;
     }
 
-    return oracle.getContents(curr);
+    return oracle.getExprContents(curr);
   }
 
   void visitExpression(Expression* curr) {

test/lit/passes/gufa-tnh.wast

@@ -1915,6 +1915,39 @@
   )
 )
 
+;; An inference that requires two internal passes of the TNH oracle to achieve.
+;; TODO: For this to pass, we need to not only update the oracle result for
+;;       the cast in $middle, but also the ref.cast itself (or read from
+;;       the oracle for the cast).
+(module
+  (func $called (param $x (ref null func))
+    (drop
+      (ref.cast func
+        (local.get $x)
+      )
+    )
+  )
+
+  (func $middle (param $x (ref null func))
+    (call $called
+      (ref.cast null func
+        (local.get $x)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $caller (type $none_=>_none)
+  ;; CHECK-NEXT:  (call $called)
+  ;; CHECK-NEXT: )
+  (func $caller (param $x (ref null func)) (export "out")
+    (call $middle
+      (ref.cast null func
+        (local.get $x)
+      )
+    )
+  )
+)
+
 ;; A CallRef with no non-trapping targets must be unreachable if traps never
 ;; happen. However, this requires closed world, so we do nothing here. (This
 ;; test is mirrored in gufa-tnh-closed, where closed world *is* enabled.)