Loop unrolling update.

- deal with non-operation stmt's (if/for stmt's) in loops being unrolled
  (unrolling of non-innermost loops works).
- update uses in unrolled bodies to use results of new operations that may be
  introduced in the unrolled bodies.

Unrolling now works for all kinds of loop nests - perfect nests, imperfect
nests, loops at any depth, and with any kind of operation in the body. (IfStmt
support not done, hence untested there).

Added missing dump/print method for StmtBlock.

TODO: add test case for outer loop unrolling.
PiperOrigin-RevId: 207314286
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index d71fb92..6875120 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -1223,6 +1223,15 @@
 
 void Statement::dump() const { print(llvm::errs()); }
 
+void StmtBlock::print(raw_ostream &os) const {
+  MLFunction *function = findFunction();
+  ModuleState state(function->getContext());
+  ModulePrinter modulePrinter(os, state);
+  MLFunctionPrinter(function, modulePrinter).print(this);
+}
+
+void StmtBlock::dump() const { print(llvm::errs()); }
+
 void Function::print(raw_ostream &os) const {
   ModuleState state(getContext());
   ModulePrinter(os, state).print(this);
diff --git a/lib/IR/Statement.cpp b/lib/IR/Statement.cpp
index 978137b..44e44c8 100644
--- a/lib/IR/Statement.cpp
+++ b/lib/IR/Statement.cpp
@@ -77,6 +77,19 @@
   return nlc.numNestedLoops == 1;
 }
 
+Statement *Statement::clone() const {
+  switch (kind) {
+  case Kind::Operation:
+    return cast<OperationStmt>(this)->clone();
+  case Kind::If:
+    llvm_unreachable("cloning for if's not implemented yet");
+    return cast<IfStmt>(this)->clone();
+  case Kind::For:
+    llvm_unreachable("cloning for loops not implemented yet");
+    return cast<ForStmt>(this)->clone();
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // ilist_traits for Statement
 //===----------------------------------------------------------------------===//
@@ -227,6 +240,15 @@
       StmtBlock(StmtBlockKind::For), lowerBound(lowerBound),
       upperBound(upperBound), step(step) {}
 
+ForStmt *ForStmt::clone() const {
+  auto *stmt = new ForStmt(getLowerBound(), getUpperBound(), getStep(),
+                           Statement::findFunction()->getContext());
+  for (auto &s : getStatements()) {
+    stmt->getStatements().push_back(s.clone());
+  }
+  return stmt;
+}
+
 //===----------------------------------------------------------------------===//
 // IfStmt
 //===----------------------------------------------------------------------===//
@@ -236,3 +258,8 @@
   if (elseClause)
     delete elseClause;
 }
+
+IfStmt *IfStmt::clone() const {
+  llvm_unreachable("cloning for if's not implemented yet");
+  return nullptr;
+}