fix dce bug with not updating the parent when turning a node unreachable (#1198)

2 files changed, 36 insertions, 23 deletions

diff --git a/src/ast/type-updating.h b/src/ast/type-updating.h
index 3cf2f2afe..8063b3e59 100644
--- a/src/ast/type-updating.h
+++ b/src/ast/type-updating.h
@@ -75,10 +75,16 @@ struct TypeUpdater : public ExpressionStackWalker<TypeUpdater, UnifiedExpression
 
   // note the replacement of one node with another. this should be called
   // after performing the replacement.
-  // this does *not* look into the node. you should do so yourself if necessary
-  void noteReplacement(Expression* from, Expression* to) {
+  // this does *not* look into the node by default. see noteReplacementWithRecursiveRemoval
+  // (we don't support recursive addition because in practice we do not create
+  // new trees in the passes that use this, they just move around children)
+  void noteReplacement(Expression* from, Expression* to, bool recursivelyRemove=false) {
     auto parent = parents[from];
-    noteRemoval(from);
+    if (recursivelyRemove) {
+      noteRecursiveRemoval(from);
+    } else {
+      noteRemoval(from);
+    }
     // if we are replacing with a child, i.e. a node that was already present
     // in the ast, then we just have a type and parent to update
     if (parents.find(to) != parents.end()) {
@@ -91,6 +97,10 @@ struct TypeUpdater : public ExpressionStackWalker<TypeUpdater, UnifiedExpression
     }
   }
 
+  void noteReplacementWithRecursiveRemoval(Expression* from, Expression* to) {
+    noteReplacement(from, to, true);
+  }
+
   // note the removal of a node
   void noteRemoval(Expression* curr) {
     noteRemovalOrAddition(curr, nullptr);
diff --git a/src/passes/DeadCodeElimination.cpp b/src/passes/DeadCodeElimination.cpp
index 5017569aa..bca898e25 100644
--- a/src/passes/DeadCodeElimination.cpp
+++ b/src/passes/DeadCodeElimination.cpp
@@ -48,6 +48,7 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
 
   Expression* replaceCurrent(Expression* expression) {
     auto* old = getCurrent();
+    if (old == expression) return expression;
     WalkerPass<PostWalker<DeadCodeElimination>>::replaceCurrent(expression);
     // also update the type updater
     typeUpdater.noteReplacement(old, expression);
@@ -219,14 +220,17 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
   }
 
   static void scan(DeadCodeElimination* self, Expression** currp) {
+    auto* curr = *currp;
     if (!self->reachable) {
-      // convert to an unreachable. do this without UB, even though we have no destructors on AST nodes
+      // convert to an unreachable safely
       #define DELEGATE(CLASS_TO_VISIT) { \
-        self->typeUpdater.noteRecursiveRemoval(*currp); \
-        ExpressionManipulator::convert<CLASS_TO_VISIT, Unreachable>(static_cast<CLASS_TO_VISIT*>(*currp)); \
+        auto* parent = self->typeUpdater.parents[curr]; \
+        self->typeUpdater.noteRecursiveRemoval(curr); \
+        ExpressionManipulator::convert<CLASS_TO_VISIT, Unreachable>(static_cast<CLASS_TO_VISIT*>(curr)); \
+        self->typeUpdater.noteAddition(curr, parent); \
         break; \
       }
-      switch ((*currp)->_id) {
+      switch (curr->_id) {
         case Expression::Id::BlockId: DELEGATE(Block);
         case Expression::Id::IfId: DELEGATE(If);
         case Expression::Id::LoopId: DELEGATE(Loop);
@@ -256,7 +260,6 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
       #undef DELEGATE
       return;
     }
-    auto* curr =* currp;
     if (curr->is<If>()) {
       self->pushTask(DeadCodeElimination::doVisitIf, currp);
       if (curr->cast<If>()->ifFalse) {
@@ -328,18 +331,18 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
     for (size_t i = 0; i < list.size(); ++i) {
       auto* elem = list[i];
       if (isUnreachable(elem)) {
-	auto* replacement = elem;
-	if (i > 0) {
-	  auto* block = getModule()->allocator.alloc<Block>();
-	  for (size_t j = 0; j < i; ++j) {
-	    block->list.push_back(drop(list[j]));
-	  }
-	  block->list.push_back(list[i]);
-	  block->finalize(type);
-	  replacement = block;
-	}
-	replaceCurrent(replacement);
-	return;
+        auto* replacement = elem;
+        if (i > 0) {
+          auto* block = getModule()->allocator.alloc<Block>();
+          for (size_t j = 0; j < i; ++j) {
+            block->list.push_back(drop(list[j]));
+          }
+          block->list.push_back(list[i]);
+          block->finalize(type);
+          replacement = block;
+        }
+        replaceCurrent(replacement);
+        return;
       }
     }
   }
@@ -349,7 +352,7 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
   }
 
   void visitLoad(Load* curr) {
-    blockifyReachableOperands({ curr->ptr}, curr->type);
+    blockifyReachableOperands({ curr->ptr }, curr->type);
   }
 
   void visitStore(Store* curr) {
@@ -369,11 +372,11 @@ struct DeadCodeElimination : public WalkerPass<PostWalker<DeadCodeElimination>>
   }
 
   void visitBinary(Binary* curr) {
-    blockifyReachableOperands({ curr->left, curr->right}, curr->type);
+    blockifyReachableOperands({ curr->left, curr->right }, curr->type);
   }
 
   void visitSelect(Select* curr) {
-    blockifyReachableOperands({ curr->ifTrue, curr->ifFalse, curr->condition}, curr->type);
+    blockifyReachableOperands({ curr->ifTrue, curr->ifFalse, curr->condition }, curr->type);
   }
 
   void visitDrop(Drop* curr) {