Improve management of instructions and basic blocks by having their inclusion
in a container automatically maintain their parent pointers, and change storage
from std::vector to the proper llvm::iplist type.

PiperOrigin-RevId: 202889037
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 067697c..c1c1d94 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -94,7 +94,7 @@
 private:
   const CFGFunction *function;
   raw_ostream &os;
-  DenseMap<BasicBlock*, unsigned> basicBlockIDs;
+  DenseMap<const BasicBlock*, unsigned> basicBlockIDs;
 };
 } // end anonymous namespace
 
@@ -103,8 +103,8 @@
 
   // Each basic block gets a unique ID per function.
   unsigned blockID = 0;
-  for (auto *block : function->blockList)
-    basicBlockIDs[block] = blockID++;
+  for (auto &block : *function)
+    basicBlockIDs[&block] = blockID++;
 }
 
 void CFGFunctionState::print() {
@@ -112,8 +112,8 @@
   printFunctionSignature(this->getFunction(), os);
   os << " {\n";
 
-  for (auto *block : function->blockList)
-    print(block);
+  for (auto &block : *function)
+    print(&block);
   os << "}\n\n";
 }
 
@@ -121,8 +121,8 @@
   os << "bb" << getBBID(block) << ":\n";
 
   // TODO Print arguments.
-  for (auto inst : block->instList)
-    print(inst);
+  for (auto &inst : block->getOperations())
+    print(&inst);
 
   print(block->getTerminator());
 }
diff --git a/lib/IR/BasicBlock.cpp b/lib/IR/BasicBlock.cpp
index 4cfe162..c2c865c 100644
--- a/lib/IR/BasicBlock.cpp
+++ b/lib/IR/BasicBlock.cpp
@@ -19,6 +19,68 @@
 #include "mlir/IR/CFGFunction.h"
 using namespace mlir;
 
-BasicBlock::BasicBlock(CFGFunction *function) : function(function) {
-  function->blockList.push_back(this);
+BasicBlock::BasicBlock() {
+}
+
+BasicBlock::~BasicBlock() {
+  if (terminator)
+    terminator->eraseFromBlock();
+}
+
+/// Unlink this BasicBlock from its CFGFunction and delete it.
+void BasicBlock::eraseFromFunction() {
+  assert(getFunction() && "BasicBlock has no parent");
+  getFunction()->getBlocks().erase(this);
+}
+
+void BasicBlock::setTerminator(TerminatorInst *inst) {
+  // If we already had a terminator, abandon it.
+  if (terminator)
+    terminator->block = nullptr;
+
+  // Reset our terminator to the new instruction.
+  terminator = inst;
+  if (inst)
+    inst->block = this;
+}
+
+mlir::CFGFunction *
+llvm::ilist_traits<::mlir::BasicBlock>::getContainingFunction() {
+  size_t Offset(
+    size_t(&((CFGFunction *)nullptr->*CFGFunction::getSublistAccess(nullptr))));
+  iplist<BasicBlock> *Anchor(static_cast<iplist<BasicBlock> *>(this));
+  return reinterpret_cast<CFGFunction *>(reinterpret_cast<char *>(Anchor) -
+                                           Offset);
+}
+
+/// This is a trait method invoked when a basic block is added to a function.
+/// We keep the function pointer up to date.
+void llvm::ilist_traits<::mlir::BasicBlock>::
+addNodeToList(BasicBlock *block) {
+  assert(!block->function && "already in a function!");
+  block->function = getContainingFunction();
+}
+
+/// This is a trait method invoked when an instruction is removed from a
+/// function.  We keep the function pointer up to date.
+void llvm::ilist_traits<::mlir::BasicBlock>::
+removeNodeFromList(BasicBlock *block) {
+  assert(block->function && "not already in a function!");
+  block->function = nullptr;
+}
+
+/// This is a trait method invoked when an instruction is moved from one block
+/// to another.  We keep the block pointer up to date.
+void llvm::ilist_traits<::mlir::BasicBlock>::
+transferNodesFromList(ilist_traits<BasicBlock> &otherList,
+                      block_iterator first, block_iterator last) {
+  // If we are transferring instructions within the same function, the parent
+  // pointer doesn't need to be updated.
+  CFGFunction *curParent = getContainingFunction();
+  if (curParent == otherList.getContainingFunction())
+    return;
+
+  // Update the 'function' member of each BasicBlock.
+  for (; first != last; ++first)
+    first->function = curParent;
 }
diff --git a/lib/IR/Instructions.cpp b/lib/IR/Instructions.cpp
index 2222a12..729936a 100644
--- a/lib/IR/Instructions.cpp
+++ b/lib/IR/Instructions.cpp
@@ -23,6 +23,27 @@
 // Instruction
 //===----------------------------------------------------------------------===//
 
+// Instructions are deleted through the destroy() member because we don't have
+// a virtual destructor.
+Instruction::~Instruction() {
+  assert(block == nullptr && "instruction destroyed but still in a block");
+}
+
+/// Destroy this instruction or one of its subclasses.
+void Instruction::destroy(Instruction *inst) {
+  switch (inst->getKind()) {
+  case Kind::Operation:
+    delete cast<OperationInst>(inst);
+    break;
+  case Kind::Branch:
+    delete cast<BranchInst>(inst);
+    break;
+  case Kind::Return:
+    delete cast<ReturnInst>(inst);
+    break;
+  }
+}
+
 CFGFunction *Instruction::getFunction() const {
   return getBlock()->getFunction();
 }
@@ -31,21 +52,64 @@
 // OperationInst
 //===----------------------------------------------------------------------===//
 
-OperationInst::OperationInst(Identifier name, BasicBlock *block) :
-  Instruction(Kind::Operation, block), name(name) {
-  getBlock()->instList.push_back(this);
+mlir::BasicBlock *
+llvm::ilist_traits<::mlir::OperationInst>::getContainingBlock() {
+  size_t Offset(
+      size_t(&((BasicBlock *)nullptr->*BasicBlock::getSublistAccess(nullptr))));
+  iplist<OperationInst> *Anchor(static_cast<iplist<OperationInst> *>(this));
+  return reinterpret_cast<BasicBlock *>(reinterpret_cast<char *>(Anchor) -
+                                           Offset);
 }
 
+/// This is a trait method invoked when an instruction is added to a block.  We
+/// keep the block pointer up to date.
+void llvm::ilist_traits<::mlir::OperationInst>::
+addNodeToList(OperationInst *inst) {
+  assert(!inst->getBlock() && "already in a basic block!");
+  inst->block = getContainingBlock();
+}
+
+/// This is a trait method invoked when an instruction is removed from a block.
+/// We keep the block pointer up to date.
+void llvm::ilist_traits<::mlir::OperationInst>::
+removeNodeFromList(OperationInst *inst) {
+  assert(inst->block && "not already in a basic block!");
+  inst->block = nullptr;
+}
+
+/// This is a trait method invoked when an instruction is moved from one block
+/// to another.  We keep the block pointer up to date.
+void llvm::ilist_traits<::mlir::OperationInst>::
+transferNodesFromList(ilist_traits<OperationInst> &otherList,
+                      instr_iterator first, instr_iterator last) {
+  // If we are transferring instructions within the same basic block, the block
+  // pointer doesn't need to be updated.
+  BasicBlock *curParent = getContainingBlock();
+  if (curParent == otherList.getContainingBlock())
+    return;
+
+  // Update the 'block' member of each instruction.
+  for (; first != last; ++first)
+    first->block = curParent;
+}
+
+/// Unlink this instruction from its BasicBlock and delete it.
+void OperationInst::eraseFromBlock() {
+  assert(getBlock() && "Instruction has no parent");
+  getBlock()->getOperations().erase(this);
+}
+
+
+
 //===----------------------------------------------------------------------===//
 // Terminators
 //===----------------------------------------------------------------------===//
 
-ReturnInst::ReturnInst(BasicBlock *parent)
-  : TerminatorInst(Kind::Return, parent) {
-  getBlock()->setTerminator(this);
+/// Remove this terminator from its BasicBlock and delete it.
+void TerminatorInst::eraseFromBlock() {
+  assert(getBlock() && "Instruction has no parent");
+  getBlock()->setTerminator(nullptr);
+  TerminatorInst::destroy(this);
 }
 
-BranchInst::BranchInst(BasicBlock *dest, BasicBlock *parent)
-  : TerminatorInst(Kind::Branch, parent), dest(dest) {
-  getBlock()->setTerminator(this);
-}
+