Compatibility of AnnotationWriter with BasicBlock Class

[knaumov]

Summary

In our downstream, we need to annotate only individual BasicBlocks and not the full Function. This change allows us to use AnnotationWriter with BasicBlock through BB.print() method. This implementation is identical to the existing one for Function.

[apilipenko]

Look good.

[knaumov]

Closed due to commit: https://github.com/llvm/llvm-project/commit/41d53194fb9d5aba3e4861233b1af9cb62cc999a

llvm/include/llvm/IR/BasicBlock.h

 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Compiler.h"
 #include <cassert>
 #include <cstddef>
 #include <iterator>
 
 namespace llvm {
 
+class AssemblyAnnotationWriter;
 class CallInst;
 class Function;
 class LandingPadInst;
 class LLVMContext;
 class Module;
 class PHINode;
 class ValueSymbolTable;
 
public:
   ///
   /// This method is analogous to getUniquePredecessor above.
   const BasicBlock *getUniqueSuccessor() const;
   BasicBlock *getUniqueSuccessor() {
     return const_cast<BasicBlock *>(
                    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
   ///
   inline iterator                begin()       { return InstList.begin(); }
   inline const_iterator          begin() const { return InstList.begin(); }
   inline iterator                end  ()       { return InstList.end();   }
   inline const_iterator          end  () const { return InstList.end();   }
 

llvm/lib/IR/AsmWriter.cpp

void Function::print(raw_ostream &ROS, AssemblyAnnotationWriter *AAW,
   SlotTracker SlotTable(this->getParent());
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this->getParent(), AAW,
                    IsForDebug,
                    ShouldPreserveUseListOrder);
   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);
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this, AAW, IsForDebug,
                    ShouldPreserveUseListOrder);
   W.printModule(this);
 }