If an instruction contains blocks, IfInst/ForInst, make sure to drop references held by those blocks when dropping references for the instruction.
PiperOrigin-RevId: 229278667
diff --git a/lib/IR/Instruction.cpp b/lib/IR/Instruction.cpp
index 89c23e5..a2e01f1 100644
--- a/lib/IR/Instruction.cpp
+++ b/lib/IR/Instruction.cpp
@@ -260,9 +260,26 @@
for (auto &op : getInstOperands())
op.drop();
- if (isTerminator())
- for (auto &dest : cast<OperationInst>(this)->getBlockOperands())
- dest.drop();
+ switch (getKind()) {
+ case Kind::For:
+ // Make sure to drop references held by instructions within the body.
+ cast<ForInst>(this)->getBody()->dropAllReferences();
+ break;
+ case Kind::If: {
+ // Make sure to drop references held by instructions within the 'then' and
+ // 'else' blocks.
+ auto *ifInst = cast<IfInst>(this);
+ ifInst->getThen()->dropAllReferences();
+ if (auto *elseBlock = ifInst->getElse())
+ elseBlock->dropAllReferences();
+ break;
+ }
+ case Kind::OperationInst:
+ if (isTerminator())
+ for (auto &dest : cast<OperationInst>(this)->getBlockOperands())
+ dest.drop();
+ break;
+ }
}
//===----------------------------------------------------------------------===//
