StructDump.cpp

#include "BoilerPlate.hpp"
#include "DwarfTag.hpp"
#include <cstdlib>
#include <llvm/BinaryFormat/Dwarf.h>
#include <llvm/DebugInfo/DWARF/DWARFContext.h>
#include <llvm/DebugInfo/DWARF/DWARFDie.h>
#include <llvm/DebugInfo/DWARF/DWARFObject.h>
#include <llvm/Support/Format.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/Support/InitLLVM.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Support/raw_ostream.h>
#include <string>
using namespace llvm;
using namespace llvm::dwarf;
using namespace llvm::object;

using Variable = DWARFDie;
using Type = DWARFDie;

static void PrintIndentLevel(raw_ostream &os, unsigned indentLevel) {
  for (unsigned lv = 0; lv < indentLevel; lv++) {
    os << '-';
  }
}
static bool IsStructType(Type type) {
  return type.getTag() == DW_TAG_structure_type;
}
static Type FindVariableType(Variable variable, StringRef variableName,
                             raw_ostream &os) {
  auto typeDie = GetDW_AT_type(variable);
  assert(typeDie.isValid());
  auto typeName = typeDie.find(DW_AT_name);
  if (typeName.has_value()) { // typedef有名，struct tag没名
    os << formatv("type of {0} is {1}\n", variableName,
                  typeName->getAsCString().get());
  }
  return ResolveTypedef(typeDie);
}
static bool VariableNameMatch(DWARFDie &die, StringRef expectName) {
  if (die.getTag() != DW_TAG_variable)
    return false;
  auto name = die.find(dwarf::DW_AT_name);
  // 有的variable没有name
  if (!name.has_value())
    return false;
  auto actualName = name.value().getAsCString();
  assert(actualName);
  return actualName.get() == expectName;
}
static Variable FindVariable(DWARFContext &DICtx, StringRef name) {
  DWARFContext::unit_iterator_range Units = DICtx.info_section_units();
  for (const auto &U : Units) {
    DWARFDie rootDie = U->getUnitDIE(false);
    for (auto child = rootDie.getFirstChild(); child;
         child = child.getSibling()) {
      if (VariableNameMatch(child, name)) {
        return child;
      }
    }
  }
  return {};
}
static void ProcessMember(DWARFDie die, raw_ostream &os, uint64_t baseOffset,
                          unsigned childLv, Twine prefix);
static void ProcessStruct(Type type, raw_ostream &os, uint64_t baseOffset,
                          unsigned childLv, Twine prefix = "") {
  PrintIndentLevel(os, childLv);
  DwarfTagStructureType st{type};
  os << formatv("struct {0} size {1}\n", st.TagName(), st.ByteSize());
  // iterate until DW_TAG_null
  for (auto child = type.getFirstChild();
       child && child.getTag() != DW_TAG_null; child = child.getSibling()) {
    ProcessMember(child, os, baseOffset, childLv + 1, prefix);
  }
}
static DWARFDie resolveReferencedType(DWARFDie D, DWARFFormValue F) {
  return D.getAttributeValueAsReferencedDie(F).resolveTypeUnitReference();
}
static void PrintDW_AT_type(DWARFDie Die, DWARFFormValue FormValue,
                            raw_ostream &OS) {
  DWARFDie D = resolveReferencedType(Die, FormValue);
  dumpTypeQualifiedName(D, OS); // 如果DW_AT_type是数组，可以打印出数组长度
}
// if(PRINT_PREFIX) print member1.member2.member3
// else print member3
constexpr bool PRINT_PREFIX = true;
static void PrintPrefix(raw_ostream &os, Twine prefix,
                        const DwarfTagMember &member) {
  if constexpr (PRINT_PREFIX) {
    if (prefix.isTriviallyEmpty()) {
      os << formatv("{0}: ", member.Name());
    } else {
      os << formatv("{0}.{1}: ", prefix, member.Name());
    }
  } else {
    os << formatv("{0}: ", member.Name());
  }
}
#include <format>
static std::string GetTypeName(DWARFDie type) {
  /*
  case typedef -> DW_AT_name
  case base type -> DW_AT_name
  case array -> type[length] how
  case struct -> struct tag
  case enum -> enum tag
  Note typedef struct{}, typedef enum is in case typedef
  */
  switch (type.getTag()) {
  case DW_TAG_typedef:
    return DwarfTagTypedef{type}.Name();
  case DW_TAG_base_type:
    return DwarfTagBaseType{type}.Name();
  case DW_TAG_array_type: {
    DwarfTagArrayType arr{type};
    return std::format("{}[{}]", arr.ElementType().getShortName(),
                       arr.Length());
  }
  case DW_TAG_structure_type: {
    DwarfTagStructureType st{type};
    return std::format("struct {}", st.TagName());
  }
  case DW_TAG_enumeration_type:
    return std::format("enum {}", DwarfTagEnumerationType{type}.TagName());
  default:
    return "(Unknown)"; // add brackets to avoid being conflict with type name
                        // "Unknown"
  }
}
static std::string OffsetStr(uint64_t baseOffset, uint64_t memberOffset) {
  if (baseOffset == 0) {
    return std::to_string(memberOffset);
  }
  return std::format("{}({}+{})", baseOffset + memberOffset, baseOffset,
                     memberOffset);
}
static uint64_t GetTypeSize(DWARFDie type) {
  /*
  case typedef -> recursively resolve to DW_AT_type
  case base type -> DW_AT_byte_size
  case array -> type * length
  case struct -> DW_AT_byte_size
  case enum -> DW_AT_byte_size
  */
  switch (type.getTag()) {
  case DW_TAG_base_type:
  case DW_TAG_enumeration_type:
  case DW_TAG_structure_type:
    return GetDW_AT_byte_size(type);
  case DW_TAG_typedef:
    return GetTypeSize(DwarfTagTypedef{type}.ResolvedType());
    break;
  case DW_TAG_array_type: {
    DwarfTagArrayType arr{type};
    return GetTypeSize(ResolveTypedef(arr.ElementType())) * arr.Length();
  } break;
  default:
    assert(0 && "Unknown tag");
  }
}
static std::string TypeSizeStr(DWARFDie type) {
  switch (type.getTag()) {
  case DW_TAG_base_type:
  case DW_TAG_enumeration_type:
  case DW_TAG_structure_type:
    return std::to_string(GetDW_AT_byte_size(type));
  case DW_TAG_typedef:
    // if is array, print one size
    return std::to_string(GetTypeSize(DwarfTagTypedef{type}.ResolvedType()));
  case DW_TAG_array_type: {
    DwarfTagArrayType arr{type};
    auto elemSize = GetTypeSize(ResolveTypedef(arr.ElementType()));
    return std::format("{0}({1}*{2})", elemSize * arr.Length(), elemSize,
                       arr.Length());
  } break;
  default:
    return "Unknown tag";
  }
}
static void ProcessMember(DWARFDie die, raw_ostream &os, uint64_t baseOffset,
                          unsigned childLv, Twine prefix) {
  const DwarfTagMember member{die};
  const auto type = member.Type();
  PrintIndentLevel(os, childLv);
  PrintPrefix(os, prefix, member);
  os << formatv("{0} offset {1} size {2}\n", GetTypeName(type),
                OffsetStr(baseOffset, member.MemberOffset()),
                TypeSizeStr(type));
  // TODO 整理重复代码
  const DWARFDie resolvedType = ResolveTypedef(type);

  // 避免临时变量自动析构
  Twine dot{'.'};
  Twine memberName{member.Name()};
  Twine tmp{prefix.concat(dot)};
  Twine nextPrefix =
      prefix.isTriviallyEmpty() ? memberName : tmp.concat(memberName);
  if (resolvedType.getTag() == DW_TAG_structure_type) {
    for (auto child = resolvedType.getFirstChild();
         child && child.getTag() != DW_TAG_null; child = child.getSibling()) {
      ProcessMember(child, os, baseOffset + member.MemberOffset(), childLv + 1,
                    nextPrefix);
    }
  } else if (resolvedType.getTag() == DW_TAG_array_type) {
    auto elementType =
        ResolveTypedef(DwarfTagArrayType{resolvedType}.ElementType());
    if (elementType.getTag() == DW_TAG_structure_type) {
      // elementType.dump();
      for (auto child = elementType.getFirstChild();
           child && child.getTag() != DW_TAG_null; child = child.getSibling()) {
        ProcessMember(child, os, baseOffset + member.MemberOffset(),
                      childLv + 1, nextPrefix);
      }
    }
  }
}
static bool dumpObjectFile(ObjectFile &Obj, DWARFContext &DICtx,
                           const Twine &FileName, StringRef VariableName,
                           raw_ostream &OS) {
  (void)Obj;
  (void)FileName;
  auto variable = FindVariable(DICtx, VariableName);
  if (!variable.isValid()) {
    OS << "variable not found\n";
    exit(EXIT_FAILURE);
  }
  auto type = FindVariableType(variable, VariableName, OS);
  if (!IsStructType(type)) {
    OS << "variable type is not a struct\n";
    exit(EXIT_FAILURE);
  }
  ProcessStruct(type, OS, 0, 0, "");
  return false;
}
int main(int argc, char **argv) {
  if (argc != 3) {
    llvm::outs() << "usage: StructDump <elf> <variable>\n";
    exit(EXIT_FAILURE);
  }
  llvm::InitLLVM X(argc, argv); // catch SIGABRT to print stacktrace
  errs().tie(&outs());
  StringRef FileName{argv[1]};
  StringRef VariableName{argv[2]};
  handleFile(FileName, dumpObjectFile, VariableName, llvm::outs());
}