From 41d53194fb9d5aba3e4861233b1af9cb62cc999a Mon Sep 17 00:00:00 2001
From: Kirill Naumov <knaumov@azul.com>
Date: Sat, 6 Jun 2020 00:06:09 +0000
Subject: [PATCH] [BasicBlock] Added AnnotationWriter functionality to
 BasicBlock class

This functionality is very similar to Function compatibility with
AnnotationWriter. This change allows us to use AnnotationWriter with
BasicBlock through BB.print() method.

Reviewed-By: apilipenko
Differntial Revision: https://reviews.llvm.org/D81321
---
 llvm/include/llvm/IR/BasicBlock.h |  7 +++++++
 llvm/lib/IR/AsmWriter.cpp         | 11 +++++++++++
 2 files changed, 18 insertions(+)

diff --git a/llvm/include/llvm/IR/BasicBlock.h b/llvm/include/llvm/IR/BasicBlock.h
index 8eeafdbaa79..24d568a728c 100644
--- a/llvm/include/llvm/IR/BasicBlock.h
+++ b/llvm/include/llvm/IR/BasicBlock.h
@@ -31,6 +31,7 @@
 
 namespace llvm {
 
+class AssemblyAnnotationWriter;
 class CallInst;
 class Function;
 class LandingPadInst;
@@ -276,6 +277,12 @@ public:
                    static_cast<const BasicBlock *>(this)->getUniqueSuccessor());
   }
 
+  /// Print the basic block to an output stream with an optional
+  /// AssemblyAnnotationWriter.
+  void print(raw_ostream &OS, AssemblyAnnotationWriter *AAW = nullptr,
+             bool ShouldPreserveUseListOrder = false,
+             bool IsForDebug = false) const;
+
   //===--------------------------------------------------------------------===//
   /// Instruction iterator methods
   ///
diff --git a/llvm/lib/IR/AsmWriter.cpp b/llvm/lib/IR/AsmWriter.cpp
index efaeadbe7c3..53bfae23441 100644
--- a/llvm/lib/IR/AsmWriter.cpp
+++ b/llvm/lib/IR/AsmWriter.cpp
@@ -4360,6 +4360,17 @@ void Function::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
   W.printFunction(this);
 }
 
+void BasicBlock::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
+                     bool ShouldPreserveUseListOrder,
+                     bool IsForDebug) const {
+  SlotTracker SlotTable(this->getModule());
+  formatted_raw_ostream OS(ROS);
+  AssemblyWriter W(OS, SlotTable, this->getModule(), AAW,
+                   IsForDebug,
+                   ShouldPreserveUseListOrder);
+  W.printBasicBlock(this);
+}
+
 void Module::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
                    bool ShouldPreserveUseListOrder, bool IsForDebug) const {
   SlotTracker SlotTable(this);
-- 
2.11.0