//===--- DebugIR.cpp - Transform debug metadata to allow debugging IR -----===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// A Module transform pass that emits a succinct version of the IR and replaces
// the source file metadata to allow debuggers to step through the IR.
//
// The location where the IR file is emitted is the same as the directory
// operand of the !llvm.dbg.cu metadata node present in the input module. The
// file name is constructed from the original file name by stripping the
// extension and replacing it with "-debug.ll" or the Postfix string specified
// at construction.
//
// FIXME: instead of replacing debug metadata, additional metadata should be
// used to point capable debuggers to the IR file without destroying the
// mapping to the original source file.
//
// FIXME: this pass should not depend on the existance of debug metadata in
// the module as it does now. Instead, it should use DIBuilder to create the
// required metadata.
//
//===----------------------------------------------------------------------===//

#include <string>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/DebugInfo.h"
#include "llvm/DIBuilder.h"
#include "llvm/IR/AsmWriter.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Instrumentation.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

namespace {

/// Returns true if Node's name contains the string "llvm.dbg"
bool isDebugNamedMetadata(const NamedMDNode *Node) {
  return Node->getName().str().find("llvm.dbg") != std::string::npos;
}

/// Returns true if Inst is a call to llvm.dbg.value or llvm.dbg.declare
bool isDebugIntrinsic(const IntrinsicInst *Inst) {
  Intrinsic::ID id = Inst->getIntrinsicID();
  return id == Intrinsic::dbg_value || id == Intrinsic::dbg_declare;
}

/// An AssemblyWriter which generates a succinct representation of the module
/// (without debug intrinsics or metadata) suitable for debugging. As IR
/// instructions are printed, !dbg metadata nodes are added (or updated)
/// to point to the corresponding line in the generated IR file instead
/// of the original source file. The input module must have been constructed
/// with debug metadata (i.e. clang -g).
class IRDebugInfoHelper : public llvm::AssemblyWriter {
  /// Directory of debug metadata
  const DebugInfoFinder &Finder;

  /// Flags to control the verbosity of the generated IR file
  bool hideDebugIntrinsics;
  bool hideDebugMetadata;

  /// Set to track metadata nodes to be printed (used only when
  /// printDebugMetadata == false)
  SmallSet<const MDNode *, 4> NonDebugNodes;

public:
  IRDebugInfoHelper(
      formatted_raw_ostream &o, const Module *M,
      AssemblyAnnotationWriter *AAW, const DebugInfoFinder &Finder,
      bool hideDebugIntrinsics = true, bool hideDebugMetadata = true)
      : AssemblyWriter(o, M, AAW), Finder(Finder),
        hideDebugIntrinsics(hideDebugIntrinsics),
        hideDebugMetadata(hideDebugMetadata) {}

private:
  virtual void printInstruction(const Instruction &I) {
    DebugLoc Loc(I.getDebugLoc());

    if (hideDebugMetadata)
      removeDebugMetadata(const_cast<Instruction &>(I));

    AssemblyWriter::printInstruction(I);
    Out.flush();
    // Adjust line number by 1 because we have not yet printed the \n
    unsigned Line = Out.getLine() + 1;

    DebugLoc NewLoc;
    if (!Loc.isUnknown())
      // I had a previous debug location: re-use the DebugLoc
      NewLoc = DebugLoc::get(Line, /* FIXME: support columns */ 0,
                             Loc.getScope(I.getContext()),
                             Loc.getInlinedAt(I.getContext()));
    else if (MDNode *scope = findFunctionMD(I.getParent()->getParent()))
      // I had no previous debug location, but M has some debug information
      NewLoc = DebugLoc::get(Line, 0, scope, /*FIXME: inlined instructions*/ 0);
    else
      // Neither I nor M has any debug information -- nothing to do here.
      // FIXME: support debugging of undecorated IR (generated by clang without
      //        the -g option)
      return;

    if (hideDebugMetadata)
      saveNonDebugMetadata(I);

    addDebugLocation(const_cast<Instruction &>(I), NewLoc);
  }

  virtual void printInstructionLine(const Instruction &I) {
    if (hideDebugIntrinsics)
      if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I))
        if (isDebugIntrinsic(II))
          return;
    AssemblyWriter::printInstructionLine(I);
  }

  virtual void writeMDNode(unsigned Slot, const MDNode *Node) {
    if (hideDebugMetadata == false || isDebugMetadata(Node) == false)
      AssemblyWriter::writeMDNode(Slot, Node);
  }

  virtual void printNamedMDNode(const NamedMDNode *NMD) {
    if (hideDebugMetadata == false || isDebugNamedMetadata(NMD) == false)
      AssemblyWriter::printNamedMDNode(NMD);
  }

  /// Returns the MDNode that corresponds with F
  MDNode *findFunctionMD(const Function *F) {
    for (DebugInfoFinder::iterator i = Finder.subprogram_begin(),
                                   e = Finder.subprogram_end();
         i != e; ++i) {
      DISubprogram S(*i);
      if (S.getFunction() == F)
        return *i;
    }
    // cannot find F -- likely means there is no debug information
    return 0;
  }

  /// Saves all non-debug metadata attached to I
  void saveNonDebugMetadata(const Instruction &I) {
    typedef SmallVector<std::pair<unsigned, MDNode *>, 4> MDNodeVector;
    MDNodeVector Others;
    I.getAllMetadataOtherThanDebugLoc(Others);
    for (MDNodeVector::iterator i = Others.begin(), e = Others.end(); i != e;
         ++i)
      NonDebugNodes.insert(i->second);
  }

  /// Returns true if Node was not saved as non-debug metadata with
  /// saveNonDebugMetadata(), false otherwise.
  bool isDebugMetadata(const MDNode *Node) {
    return NonDebugNodes.count(Node) == 0;
  }

  void removeDebugMetadata(Instruction &I) {
    if (I.getMetadata(LLVMContext::MD_dbg))
      I.setMetadata(LLVMContext::MD_dbg, 0);
  }

  void addDebugLocation(Instruction &I, DebugLoc Loc) {
    MDNode *MD = Loc.getAsMDNode(I.getContext());
    I.setMetadata(LLVMContext::MD_dbg, MD);
  }
};

class DebugIR : public ModulePass {
  std::string Postfix;
  std::string Filename;
  DebugInfoFinder Finder;

public:
  static char ID;

  DebugIR() : ModulePass(ID), Postfix("-debug.ll") {}

  /// Customize the postfix string used to replace the extension of the
  /// original filename that appears in the !llvm.dbg.cu metadata node.
  DebugIR(StringRef postfix) : ModulePass(ID), Postfix(postfix) {}

private:
  // Modify the filename embedded in the Compilation-Unit debug information of M
  bool replaceFilename(Module &M) {
    bool changed = false;

    // Sanity check -- if llvm.dbg.cu node exists, the DebugInfoFinder
    // better have found at least one CU!
    if (M.getNamedMetadata("llvm.dbg.cu"))
      assert(Finder.compile_unit_count() > 0 &&
             "Found no compile units but llvm.dbg.cu node exists");

    for (DebugInfoFinder::iterator i = Finder.compile_unit_begin(),
                                   e = Finder.compile_unit_end();
         i != e; ++i) {
      DICompileUnit CU(*i);
      Filename = CU.getFilename();

      // Replace extension with postfix
      size_t dot = Filename.find_last_of(".");
      if (dot != std::string::npos)
        Filename.erase(dot);
      Filename += Postfix;

      CU.setFilename(Filename, M.getContext());
      changed = true;
    }
    return changed;
  }

  void writeAndUpdateDebugIRFile(Module *M) {
    std::string error;
    tool_output_file OutFile(Filename.c_str(), error);
    OutFile.keep();
    formatted_raw_ostream OS;
    OS.setStream(OutFile.os(), false);

    IRDebugInfoHelper W(OS, M, 0, Finder);
    W.printModule(M);
  }

  bool runOnModule(Module &M) {
    Finder.processModule(M);
    bool changed = replaceFilename(M);
    if (changed)
      writeAndUpdateDebugIRFile(&M);
    return changed;
  }
};

} // anonymous namespace

char DebugIR::ID = 0;
INITIALIZE_PASS(DebugIR, "debug-ir", "Enable debugging IR", false, false)
    ModulePass *llvm::createDebugIRPass(StringRef FilenamePostfix) {
  return new DebugIR(FilenamePostfix);
}