[IR]PATCH 2/2: Add MDNode::printTree and dumpTree

[myhsu]

Summary

This patch adds the functionalities to print MDNode in tree shape. For example, instead of printing a MDNode like this:

<0x5643e1166888> = !DILocalVariable(name: "foo", arg: 2, scope: <0x5643e11c9740>, file: <0x5643e11c6ec0>, line: 8, type: <0x5643e11ca8e0>, flags: DIFlagPublic | DIFlagFwdDecl, align: 8)

The printTree/dumpTree functions can give you:

<0x5643e1166888> = !DILocalVariable(name: "foo", arg: 2, scope: <0x5643e11c9740>, file: <0x5643e11c6ec0>, line: 8, type: <0x5643e11ca8e0>, flags: DIFlagPublic | DIFlagFwdDecl, align: 8)
  <0x5643e11c9740> = distinct !DISubprogram(scope: null, spFlags: 0)
  <0x5643e11c6ec0> = distinct !DIFile(filename: "file.c", directory: "/path/to/dir")
  <0x5643e11ca8e0> = distinct !DIDerivedType(tag: DW_TAG_pointer_type, baseType: <0x5643e11668d8>, size: 1, align: 2)
    <0x5643e11668d8> = !DIBasicType(tag: DW_TAG_unspecified_type, name: "basictype")

Which is useful when using it in debugger. Where sometimes printing the whole module to see all MDNodes is too expensive.

[aprantl]

I like this! Can you add a testcase where the MDNode graph has cycles?

[myhsu]

Add another unittest to test the case where there is a cycle in the MDNode to be printed by printTree.

[aprantl]

llvm/include/llvm/IR/Metadata.h

1047

Nit, I think this can be \{. Looks more TeXy.

llvm/lib/IR/AsmWriter.cpp

4670

///

[myhsu]

Minor formatting updates.

llvm/include/llvm/IR/Metadata.h

public:
   /// Create a distinct version of \c N -- in place, if possible -- and return
   /// it.  Takes ownership of the temporary node.
   template <class T>
   static std::enable_if_t<std::is_base_of<MDNode, T>::value, T *>
   replaceWithDistinct(std::unique_ptr<T, TempMDNodeDeleter> N) {
     return cast<T>(N.release()->replaceWithDistinctImpl());
   }
 
+  /// Print in tree shape.
+  ///
+  /// Prints definition of \c this in tree shape.
+  ///
+  /// If \c M is provided, metadata nodes will be numbered canonically;
+  /// otherwise, pointer addresses are substituted.
+  /// @{
+  void printTree(raw_ostream &OS, const Module *M = nullptr) const;
+  void printTree(raw_ostream &OS, ModuleSlotTracker &MST,
+                 const Module *M = nullptr) const;
+  /// @}
+
+  /// User-friendly dump in tree shape.
+  ///
+  /// If \c M is provided, metadata nodes will be numbered canonically;
+  /// otherwise, pointer addresses are substituted.
+  ///
+  /// Note: this uses an explicit overload instead of default arguments so that
+  /// the nullptr version is easy to call from a debugger.
+  ///
+  /// @{
+  void dumpTree() const;
+  void dumpTree(const Module *M) const;
+  /// @}
+
 private:
   MDNode *replaceWithPermanentImpl();
   MDNode *replaceWithUniquedImpl();
   MDNode *replaceWithDistinctImpl();
 
 protected:
   /// Set an operand.
   ///

llvm/lib/IR/AsmWriter.cpp

 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/Argument.h"
struct AsmWriterContext {
 
   AsmWriterContext(TypePrinting *TP, SlotTracker *ST, const Module *M = nullptr)
       : TypePrinter(TP), Machine(ST), Context(M) {}
 
   static AsmWriterContext &getEmpty() {
     static AsmWriterContext EmptyCtx(nullptr, nullptr);
     return EmptyCtx;
   }
+
+  /// A callback that will be triggered when the underlying printer
+  /// prints a Metadata as operand.
+  virtual void onWriteMetadataAsOperand(const Metadata *) {}
+
+  virtual ~AsmWriterContext() {}
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // AsmWriter Implementation
 //===----------------------------------------------------------------------===//
 
 static void WriteAsOperandInternal(raw_ostream &Out, const Value *V,
if (!MD)
       Out << "null";
     else if (auto *MDV = dyn_cast<ValueAsMetadata>(MD)) {
       Value *V = MDV->getValue();
       WriterCtx.TypePrinter->print(V->getType(), Out);
       Out << ' ';
       WriteAsOperandInternal(Out, V, WriterCtx);
     } else {
       WriteAsOperandInternal(Out, MD, WriterCtx);
+      WriterCtx.onWriteMetadataAsOperand(MD);
     }
     if (mi + 1 != me)
       Out << ", ";
   }
 
   Out << "}";
 }
 

 
 static void writeMetadataAsOperand(raw_ostream &Out, const Metadata *MD,
                                    AsmWriterContext &WriterCtx) {
   if (!MD) {
     Out << "null";
     return;
   }
   WriteAsOperandInternal(Out, MD, WriterCtx);
+  WriterCtx.onWriteMetadataAsOperand(MD);
 }
 
 void MDFieldPrinter::printMetadata(StringRef Name, const Metadata *MD,
                                    bool ShouldSkipNull) {
   if (ShouldSkipNull && !MD)
     return;
 
   Out << FS << Name << ": ";
void Value::printAsOperand(raw_ostream &O, bool PrintType,
                            ModuleSlotTracker &MST) const {
   if (!PrintType)
     if (printWithoutType(*this, O, MST.getMachine(), MST.getModule()))
       return;
 
   printAsOperandImpl(*this, O, PrintType, MST);
 }
 
+/// Recursive version of printMetadataImpl.
+static void printMetadataImplRec(raw_ostream &ROS, const Metadata &MD,
+                                 AsmWriterContext &WriterCtx) {
+  formatted_raw_ostream OS(ROS);
+  WriteAsOperandInternal(OS, &MD, WriterCtx, /* FromValue */ true);
+
+  auto *N = dyn_cast<MDNode>(&MD);
+  if (!N || isa<DIExpression>(MD) || isa<DIArgList>(MD))
+    return;
+
+  OS << " = ";
+  WriteMDNodeBodyInternal(OS, N, WriterCtx);
+}
+
+namespace {
+struct MDTreeAsmWriterContext : public AsmWriterContext {
+  unsigned Level;
+  // {Level, Printed string}
+  using EntryTy = std::pair<unsigned, std::string>;
+  SmallVector<EntryTy, 4> Buffer;
+
+  // Used to break the cycle in case there is any.
+  SmallPtrSet<const Metadata *, 4> Visited;
+
+  raw_ostream &MainOS;
+
+  MDTreeAsmWriterContext(TypePrinting *TP, SlotTracker *ST, const Module *M,
+                         raw_ostream &OS, const Metadata *InitMD)
+      : AsmWriterContext(TP, ST, M), Level(0U), Visited({InitMD}), MainOS(OS) {}
+
+  void onWriteMetadataAsOperand(const Metadata *MD) override {
+    if (Visited.count(MD))
+      return;
+    Visited.insert(MD);
+
+    std::string Str;
+    raw_string_ostream SS(Str);
+    ++Level;
+    // A placeholder entry to memorize the correct
+    // position in buffer.
+    Buffer.emplace_back(std::make_pair(Level, ""));
+    unsigned InsertIdx = Buffer.size() - 1;
+
+    printMetadataImplRec(SS, *MD, *this);
+    Buffer[InsertIdx].second = std::move(SS.str());
+    --Level;
+  }
+
+  ~MDTreeAsmWriterContext() {
+    for (const auto &Entry : Buffer) {
+      MainOS << "\n";
+      unsigned NumIndent = Entry.first * 2U;
+      MainOS.indent(NumIndent) << Entry.second;
+    }
+  }
+};
+} // end anonymous namespace
+
 static void printMetadataImpl(raw_ostream &ROS, const Metadata &MD,
                               ModuleSlotTracker &MST, const Module *M,
-                              bool OnlyAsOperand) {
+                              bool OnlyAsOperand, bool PrintAsTree = false) {
   formatted_raw_ostream OS(ROS);
 
   TypePrinting TypePrinter(M);
 
-  AsmWriterContext WriterCtx(&TypePrinter, MST.getMachine(), M);
-  WriteAsOperandInternal(OS, &MD, WriterCtx, /* FromValue */ true);
+  std::unique_ptr<AsmWriterContext> WriterCtx;
+  if (PrintAsTree && !OnlyAsOperand)
+    WriterCtx = std::make_unique<MDTreeAsmWriterContext>(
+        &TypePrinter, MST.getMachine(), M, OS, &MD);
+  else
+    WriterCtx =
+        std::make_unique<AsmWriterContext>(&TypePrinter, MST.getMachine(), M);
+
+  WriteAsOperandInternal(OS, &MD, *WriterCtx, /* FromValue */ true);
 
   auto *N = dyn_cast<MDNode>(&MD);
   if (OnlyAsOperand || !N || isa<DIExpression>(MD) || isa<DIArgList>(MD))
     return;
 
   OS << " = ";
-  WriteMDNodeBodyInternal(OS, N, WriterCtx);
+  WriteMDNodeBodyInternal(OS, N, *WriterCtx);
 }
 
 void Metadata::printAsOperand(raw_ostream &OS, const Module *M) const {
   ModuleSlotTracker MST(M, isa<MDNode>(this));
   printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ true);
 }
 
 void Metadata::printAsOperand(raw_ostream &OS, ModuleSlotTracker &MST,
                               const Module *M) const {
   printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ true);
 }
 
 void Metadata::print(raw_ostream &OS, const Module *M,
                      bool /*IsForDebug*/) const {
   ModuleSlotTracker MST(M, isa<MDNode>(this));
   printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ false);
 }
 
 void Metadata::print(raw_ostream &OS, ModuleSlotTracker &MST,
                      const Module *M, bool /*IsForDebug*/) const {
   printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ false);
 }
 
+void MDNode::printTree(raw_ostream &OS, const Module *M) const {
+  ModuleSlotTracker MST(M, true);
+  printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ false,
+                    /*PrintAsTree=*/true);
+}
+
+void MDNode::printTree(raw_ostream &OS, ModuleSlotTracker &MST,
+                       const Module *M) const {
+  printMetadataImpl(OS, *this, MST, M, /* OnlyAsOperand */ false,
+                    /*PrintAsTree=*/true);
+}
+
 void ModuleSummaryIndex::print(raw_ostream &ROS, bool IsForDebug) const {
   SlotTracker SlotTable(this);
   formatted_raw_ostream OS(ROS);
   AssemblyWriter W(OS, SlotTable, this, IsForDebug);
   W.printModuleSummaryIndex();
 }
 
 void ModuleSlotTracker::collectMDNodes(MachineMDNodeListType &L, unsigned LB,

 void Metadata::dump() const { dump(nullptr); }
 
 LLVM_DUMP_METHOD
 void Metadata::dump(const Module *M) const {
   print(dbgs(), M, /*IsForDebug=*/true);
   dbgs() << '\n';
 }
 
+LLVM_DUMP_METHOD
+void MDNode::dumpTree() const { dumpTree(nullptr); }
+
+LLVM_DUMP_METHOD
+void MDNode::dumpTree(const Module *M) const {
+  printTree(dbgs(), M);
+  dbgs() << '\n';
+}
+
 // Allow printing of ModuleSummaryIndex from the debugger.
 LLVM_DUMP_METHOD
 void ModuleSummaryIndex::dump() const { print(dbgs(), /*IsForDebug=*/true); }
 #endif

llvm/unittests/IR/MetadataTest.cpp

TEST_F(MDNodeTest, PrintWithDroppedCallOperand) {
   auto *N0 = MDNode::getDistinct(Context, None);
   R0->setMetadata("md", N0);
 
   // Printing the metadata node would previously result in a failed assertion
   // due to the call instruction's dropped function operand.
   ModuleSlotTracker MST(&M);
   EXPECT_PRINTER_EQ("!0 = distinct !{}", N0->print(OS, MST));
 }
+
+TEST_F(MDNodeTest, PrintTree) {
+  DILocalScope *Scope = getSubprogram();
+  DIFile *File = getFile();
+  DINode::DIFlags Flags = static_cast<DINode::DIFlags>(7);
+  {
+    DIType *Type = getDerivedType();
+    auto *Var = DILocalVariable::get(Context, Scope, "foo", File,
+                                     /*LineNo=*/8, Type, /*ArgNo=*/2, Flags,
+                                     /*Align=*/8, nullptr);
+    std::string Expected;
+    {
+      raw_string_ostream SS(Expected);
+      Var->print(SS);
+      // indent level 1
+      Scope->print((SS << "\n").indent(2));
+      File->print((SS << "\n").indent(2));
+      Type->print((SS << "\n").indent(2));
+      // indent level 2
+      auto *BaseType = cast<DIDerivedType>(Type)->getBaseType();
+      BaseType->print((SS << "\n").indent(4));
+    }
+
+    EXPECT_PRINTER_EQ(Expected, Var->printTree(OS));
+  }
+
+  {
+    // Test if printTree works correctly when there is
+    // a cycle in the MDNode and its dependencies.
+    //
+    // We're trying to create type like this:
+    // struct LinkedList {
+    //   LinkedList *Head;
+    // };
+    auto *StructTy = cast<DICompositeType>(getCompositeType());
+    DIType *PointerTy = DIDerivedType::getDistinct(
+        Context, dwarf::DW_TAG_pointer_type, "", nullptr, 0, nullptr, StructTy,
+        1, 2, 0, None, DINode::FlagZero);
+    StructTy->replaceElements(MDTuple::get(Context, PointerTy));
+
+    auto *Var = DILocalVariable::get(Context, Scope, "foo", File,
+                                     /*LineNo=*/8, StructTy, /*ArgNo=*/2, Flags,
+                                     /*Align=*/8, nullptr);
+    std::string Expected;
+    {
+      raw_string_ostream SS(Expected);
+      Var->print(SS);
+      // indent level 1
+      Scope->print((SS << "\n").indent(2));
+      File->print((SS << "\n").indent(2));
+      StructTy->print((SS << "\n").indent(2));
+      // indent level 2
+      StructTy->getRawElements()->print((SS << "\n").indent(4));
+      // indent level 3
+      auto Elements = StructTy->getElements();
+      Elements[0]->print((SS << "\n").indent(6));
+    }
+
+    EXPECT_PRINTER_EQ(Expected, Var->printTree(OS));
+  }
+}
 #undef EXPECT_PRINTER_EQ
 
 TEST_F(MDNodeTest, NullOperand) {
   // metadata !{}
   MDNode *Empty = MDNode::get(Context, None);
 
   // metadata !{metadata !{}}
   Metadata *Ops[] = {Empty};