[mlir] Add basic block arguments

This patch adds support for basic block arguments including parsing and printing.

In doing so noticed that `ssa-id-and-type` is undefined in the MLIR spec; suggested an implementation in the spec doc.

PiperOrigin-RevId: 205593369
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index d811223..0d5610f 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -665,7 +665,9 @@
 
 /// Number all of the SSA values in the specified basic block.
 void CFGFunctionPrinter::numberValuesInBlock(const BasicBlock *block) {
-  // TODO: basic block arguments.
+  for (auto *arg : block->getArguments()) {
+    numberValueID(arg);
+  }
   for (auto &op : *block) {
     // We number instruction that have results, and we only number the first
     // result.
@@ -686,16 +688,26 @@
 }
 
 void CFGFunctionPrinter::print(const BasicBlock *block) {
-  os << "bb" << getBBID(block) << ":\n";
+  os << "bb" << getBBID(block);
 
-  // TODO Print arguments.
+  if (!block->args_empty()) {
+    os << '(';
+    interleaveComma(block->getArguments(), [&](const BBArgument *arg) {
+      printValueID(arg);
+      os << ": ";
+      ModulePrinter::print(arg->getType());
+    });
+    os << ')';
+  }
+  os << ":\n";
+
   for (auto &inst : block->getOperations()) {
     print(&inst);
-    os << "\n";
+    os << '\n';
   }
 
   print(block->getTerminator());
-  os << "\n";
+  os << '\n';
 }
 
 void CFGFunctionPrinter::print(const Instruction *inst) {
diff --git a/lib/IR/BasicBlock.cpp b/lib/IR/BasicBlock.cpp
index c2c865c..1aad430 100644
--- a/lib/IR/BasicBlock.cpp
+++ b/lib/IR/BasicBlock.cpp
@@ -19,12 +19,14 @@
 #include "mlir/IR/CFGFunction.h"
 using namespace mlir;
 
-BasicBlock::BasicBlock() {
-}
+BasicBlock::BasicBlock() {}
 
 BasicBlock::~BasicBlock() {
   if (terminator)
     terminator->eraseFromBlock();
+  for (BBArgument *arg : arguments)
+    delete arg;
+  arguments.clear();
 }
 
 /// Unlink this BasicBlock from its CFGFunction and delete it.
@@ -84,3 +86,17 @@
   for (; first != last; ++first)
     first->function = curParent;
 }
+
+BBArgument *BasicBlock::addArgument(Type *type) {
+  arguments.push_back(new BBArgument(type, this));
+  return arguments.back();
+}
+
+llvm::iterator_range<BasicBlock::BBArgListType::iterator>
+BasicBlock::addArguments(ArrayRef<Type *> types) {
+  auto initial_size = arguments.size();
+  for (auto *type : types) {
+    addArgument(type);
+  }
+  return {arguments.data() + initial_size, arguments.data() + arguments.size()};
+}
diff --git a/lib/IR/Verifier.cpp b/lib/IR/Verifier.cpp
index 3c3bb16..d1bb2ac 100644
--- a/lib/IR/Verifier.cpp
+++ b/lib/IR/Verifier.cpp
@@ -106,6 +106,26 @@
 bool CFGFuncVerifier::verify() {
   // TODO: Lots to be done here, including verifying dominance information when
   // we have uses and defs.
+  // TODO: Verify the first block has no predecessors.
+
+  if (fn.empty())
+    return failure("cfgfunc must have at least one basic block", fn);
+
+  // Verify that the argument list of the function and the arg list of the first
+  // block line up.
+  auto *firstBB = &fn.front();
+  auto fnInputTypes = fn.getType()->getInputs();
+  if (fnInputTypes.size() != firstBB->getNumArguments())
+    return failure("first block of cfgfunc must have " +
+                       Twine(fnInputTypes.size()) +
+                       " arguments to match function signature",
+                   fn);
+  for (unsigned i = 0, e = firstBB->getNumArguments(); i != e; ++i)
+    if (fnInputTypes[i] != firstBB->getArgument(i)->getType())
+      return failure(
+          "type of argument #" + Twine(i) +
+              " must match corresponding argument in function signature",
+          fn);
 
   for (auto &block : fn) {
     if (verifyBlock(block))
@@ -121,6 +141,11 @@
   if (verifyTerminator(*block.getTerminator()))
     return true;
 
+  for (auto *arg : block.getArguments()) {
+    if (arg->getOwner() != &block)
+      return failure("basic block argument not owned by block", block);
+  }
+
   for (auto &inst : block) {
     if (verifyOperation(inst))
       return true;