Speed up local variable analysis by reusing previous results

src/ir/LocalGraph.cpp

@@ -118,10 +118,12 @@ struct LocalGraphFlower
     // scanning them linearly is efficient, avoiding hash computations (while
     // in Info, it's convenient to have a map so we can assign them easily,
     // where the last one seen overwrites the previous; and, we do that O(1)).
-    // TODO: If we also stored gets here then we could use the sets for a get
-    //       we already computed, for a get that we are computing, and stop that
-    //       part of the flow.
     std::vector<std::pair<Index, LocalSet*>> lastSets;
+
+    // Completed queries at the start of this block. If the block has no set
+    // for an index, these also describe its outgoing value. Reuse them instead
+    // of repeatedly traversing the same predecessors for later queries.
+    std::vector<std::pair<Index, LocalGet*>> computedGets;
   };
 
   // All the flow blocks.
@@ -285,12 +287,23 @@ struct LocalGraphFlower
                                       });
           if (lastSet != pred->lastSets.end()) {
             // There is a set here, apply it, and stop the flow.
-            // TODO: If we find a computed get, apply its sets and stop? That
-            //       could help but it requires more info on FlowBlock.
             for (auto* get : gets) {
               getSetsMap[get].insert(lastSet->second);
             }
           } else {
+            auto computed = std::find_if(
+              pred->computedGets.begin(),
+              pred->computedGets.end(),
+              [&](const auto& value) { return value.first == index; });
+            if (computed != pred->computedGets.end()) {
+              const auto& sets = getSetsMap.at(computed->second);
+              for (auto* get : gets) {
+                for (auto* set : sets) {
+                  getSetsMap[get].insert(set);
+                }
+              }
+              continue;
+            }
             // Keep on flowing.
             work.push_back(pred);
           }
@@ -300,6 +313,11 @@ struct LocalGraphFlower
 
     // Bump the current iteration for the next time we are called.
     currentIteration++;
+
+    // Publish only after the entire search has finished: a loop may revisit
+    // this block while its query is still incomplete.
+    getSetsMap.try_emplace(gets[0]);
+    block->computedGets.emplace_back(index, gets[0]);
   }
 
   // When the LocalGraph is in lazy mode we do not compute all of getSetsMap

test/gtest/local-graph.cpp

@@ -1,4 +1,5 @@
 #include "ir/local-graph.h"
+#include "ir/find_all.h"
 #include "parser/wat-parser.h"
 #include "wasm.h"
 
@@ -8,6 +9,89 @@ using LocalGraphTest = ::testing::Test;
 
 using namespace wasm;
 
+TEST_F(LocalGraphTest, ReuseCompletedCrossBlockQueries) {
+  Module wasm;
+  ASSERT_FALSE(WATParser::parseModule(wasm, R"wasm(
+    (module
+      (func $overwritten (param $c i32) (local $x i32)
+        (local.set $x (i32.const 10))
+        (if (local.get $c) (then (drop (local.get $x))))
+        (if (local.get $c) (then (drop (local.get $x))))
+        (local.set $x (i32.const 20))
+        (if (local.get $c) (then (drop (local.get $x))))
+        (drop (local.get $x)))
+      (func $loop (param $c i32) (local $x i32)
+        (local.set $x (i32.const 10))
+        (loop $again
+          (drop (local.get $x))
+          (if (local.get $c)
+            (then (drop (local.get $x)))
+            (else (local.set $x (i32.const 20))))
+          (drop (local.get $x))
+          (br_if $again (local.get $c)))
+        (drop (local.get $x)))
+      (func $initial (param $c i32) (local $x i32)
+        (loop $again
+          (drop (local.get $x))
+          (if (local.get $c)
+            (then (drop (local.get $x)))
+            (else (local.set $x (i32.const 20))))
+          (drop (local.get $x))
+          (br_if $again (local.get $c)))
+        (drop (local.get $x)))
+      (func $unreachable (param $c i32) (local $x i32)
+        (local.set $x (i32.const 10))
+        (unreachable)
+        (if (local.get $c) (then (drop (local.get $x))))
+        (drop (local.get $x)))
+    )
+  )wasm")
+                 .getErr());
+
+  for (auto& func : wasm.functions) {
+    SCOPED_TRACE(func->name.toString());
+    FindAll<LocalSet> sets(func->body);
+    FindAll<LocalGet> allGets(func->body);
+    std::vector<LocalGet*> gets;
+    for (auto* get : allGets.list) {
+      if (get->index == 1) {
+        gets.push_back(get);
+      }
+    }
+    auto check = [&](const LocalGraphBase::Sets& actual, size_t i) {
+      SCOPED_TRACE(i);
+      if (func->name == "overwritten") {
+        ASSERT_EQ(actual.size(), 1U);
+        EXPECT_TRUE(actual.contains(sets.list[i < 2 ? 0 : 1]));
+      } else if (func->name == "loop") {
+        ASSERT_EQ(actual.size(), 2U);
+        EXPECT_TRUE(actual.contains(sets.list[0]));
+        EXPECT_TRUE(actual.contains(sets.list[1]));
+      } else if (func->name == "initial") {
+        ASSERT_EQ(actual.size(), 2U);
+        EXPECT_TRUE(actual.contains(nullptr));
+        EXPECT_TRUE(actual.contains(sets.list[0]));
+      } else {
+        EXPECT_TRUE(actual.empty());
+      }
+    };
+
+    LocalGraph eager(func.get(), &wasm);
+    for (size_t i = 0; i < gets.size(); ++i) {
+      check(eager.getSets(gets[i]), i);
+    }
+    // Lazy queries must give the same answers regardless of query order,
+    // including when a query traverses a loop before earlier gets are queried.
+    for (bool reverse : {false, true}) {
+      LazyLocalGraph lazy(func.get(), &wasm);
+      for (size_t j = 0; j < gets.size(); ++j) {
+        size_t i = reverse ? gets.size() - 1 - j : j;
+        check(lazy.getSets(gets[i]), i);
+      }
+    }
+  }
+}
+
 TEST_F(LocalGraphTest, ObstacleBasics) {
   auto moduleText = R"wasm(
     (module