Remove duplicate let values (#554)

dsharlet · Jan 15, 2025 · 2cba9a7 · 2cba9a7
1 parent 72d14c1
commit 2cba9a7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/builder/simplify.cc b/builder/simplify.cc
@@ -1053,10 +1053,20 @@ class simplifier : public node_mutator {
     std::vector<sv_type> scoped_values;
     scoped_values.reserve(op->lets.size());
 
+    std::map<expr, var, node_less> reverse_lets;
+
     bool values_changed = false;
     for (const auto& s : op->lets) {
       expr_info value_info;
       expr value = mutate(s.second, &value_info);
+
+      var& name = reverse_lets[value];
+      if (name.defined()) {
+        // This value is the same as another let value, use that variable instead.
+        value = expr(name);
+      } else {
+        name = s.first;
+      }
       if (should_substitute(value)) {
         value_info = expr_info::substitution(std::move(value));
         values_changed = true;

diff --git a/builder/test/simplify/simplify.cc b/builder/test/simplify/simplify.cc
@@ -234,6 +234,9 @@ TEST(simplify, let) {
   ASSERT_THAT(simplify(let::make({{x, y * 2}, {z, y}}, z)), matches(y));
   ASSERT_THAT(simplify(let::make({{x, y}, {z, (x + 1) / x}}, (z + 1) / z)),
       matches(let::make({{z, (y + 1) / y}}, (z + 1) / z)));
+
+  // Duplicate lets
+  ASSERT_THAT(simplify(let::make({{x, y * 2}, {z, y * 2}}, x + z)), matches(let::make(x, y * 2, x * 2)));
 }
 
 TEST(simplify, loop) {