text: /*
   * Copyright (c) Facebook, Inc. and its affiliates.
   *
   * This source code is licensed under the MIT license found in the LICENSE
   * file in the root directory of this source tree.
   */
  #include "llvm/ADT/SmallVector.h"
  #include "llvm/ADT/StringRef.h"
  #include "llvm/Support/CommandLine.h"
  #include "llvm/Support/FileSystem.h"
  #include "llvm/Support/InitLLVM.h"
  #include "llvm/Support/MemoryBuffer.h"
  #include "llvm/Support/PrettyStackTrace.h"
  #include "llvm/Support/Signals.h"
  #include "llvm/Support/raw_ostream.h"
  
  #include "hermes/BCGen/HBC/BytecodeDisassembler.h"
  #include "hermes/BCGen/HBC/SerializedLiteralGenerator.h"
  #include "hermes/Public/Buffer.h"
  #include "hermes/Support/JSONEmitter.h"
  #include "hermes/Support/LEB128.h"
  #include "hermes/Support/MemoryBuffer.h"
  #include "hermes/Support/StringKind.h"
  
  #include <algorithm>
  #include <iostream>
  #include <iterator>
  #include <sstream>
  #include <string>
  #include <vector>
  
  /*
   * hbc-attribute attributes bundle size to each function. Additional tools
   * allow symbolicating and grouping the information into files. Data shared
   * entirely/partially between functions are correctly attributed.
   *
   * The basic principle is to emit a nested type, size and deduplication key
   * for each type of data. For example:
   *
   *     Function   Type                Key    Size
   *     Foo        data:string:entry   1      8
   *     Bar        data:string:entry   1      8
   *     Bar        data:string:entry   2      8
   *
   * The following is assumed to hold:
   *
   *  1. Records of different Type will not overlap.
   *  2. Records of the same Type but different Key will not overlap.
   *  3. Records of the same Type and same Key will completely overlap.
   *  4. Records of the same Type and same Key will have the same Size.
   *
   */
  
  using namespace hermes;
  using namespace hermes::hbc;
  using namespace hermes::inst;
  
  using llvm::MutableArrayRef;
  using llvm::raw_fd_ostream;
  using SLG = hermes::hbc::SerializedLiteralGenerator;
  
  /* This tool is highly dependent upon the current bytecode format.
   *
   * If you have added a simple instruction, make sure string parameters (if any)
   * are marked with the OPERAND_STRING_ID macro and bump the version below.
   *
   * If you have added or modified sections, make sure they're counted properly.
   */
  static_assert(
      BYTECODE_VERSION == 63,
      "Bytecode version changed. Please verify that hbc-attribute counts correctly..");
  
  static llvm::cl::opt<std::string> InputFilename(
      llvm::cl::Positional,
      llvm::cl::desc("Input bundle"),
      llvm::cl::init("-"));
  
  static llvm::cl::opt<std::string>
      OutputFilename("out", llvm::cl::desc("Output file"), llvm::cl::init("-"));
  
  namespace {
  template <typename T>
  unsigned byteSize(ArrayRef<T> ref) {
    return ref.size() * sizeof(T);
  }
  
  /// Walks the bytecode and outputs usage info.
  class UsageCounter : public BytecodeVisitor {
   protected:
    JSONEmitter &emitter_;
    llvm::DenseMap<unsigned, unsigned> virtualOffsets_;
    uintptr_t bundleStart_;
  
    unsigned currentFuncId_;
    uintptr_t opcodeStart_;
    uintptr_t opcodeEnd_;
    uintptr_t functionEnd_;
  
    llvm::DenseMap<std::pair<StringRef, unsigned>, unsigned> emitted_;
  
    /// Indices into the bytecode's string table corresponding to the (exclusive)
    /// end of each string kind entry.
    std::vector<uint32_t> stringKindEnds_;
  
    void appendRecord(llvm::StringRef type, unsigned dedupKey, unsigned size) {
      assert(size < (2 << 20) && "Abnormally large size!");
  
      if (size == 0) {
        // Don't bother counting anything of size 0.
        return;
      }
  
      // Do one pass of deduplication while emitting. This cuts output in half.
      std::pair<StringRef, unsigned> key = {type, dedupKey};
      if (emitted_.count(key)) {
        assert(emitted_[key] == size && "Expected deduped entry to be same size");
        return;
      }
      emitted_[key] = size;
  
      emitter_.openDict();
      emitter_.emitKeyValue("type", type);
      emitter_.emitKeyValue("dedupKey", dedupKey);
      emitter_.emitKeyValue("size", size);
      emitter_.closeDict();
    }
  
    /// Emits available location data we have, which is often just virtualOffset.
    void emitFunctionLocation() {
      auto debugInfo = bcProvider_->getDebugInfo();
      auto offsets = bcProvider_->getDebugOffsets(currentFuncId_);
  
      emitter_.emitKey("location");
      emitter_.openDict();
      if (offsets && offsets->sourceLocations != DebugOffsets::NO_OFFSET) {
        if (auto pos =
                debugInfo->getLocationForAddress(offsets->sourceLocations, 0)) {
          emitter_.emitKeyValue(
              "file", debugInfo->getFilenameByID(pos->filenameId));
          emitter_.emitKeyValue("line", pos->line);
          emitter_.emitKeyValue("column", pos->column);
        }
      }
      emitter_.emitKeyValue("virtualOffset", virtualOffsets_[currentFuncId_]);
      emitter_.emitKeyValue(
          "bytecodeSize",
          bcProvider_->getFunctionHeader(currentFuncId_).bytecodeSizeInBytes());
      emitter_.closeDict();
    }
  
    /// Emits per-bundle information.
    void emitGlobalInfo() {
      appendRecord("headers:global:bundle", 0, sizeof(BytecodeFileHeader));
      appendRecord("headers:global:debuginfo", 0, sizeof(DebugInfoHeader));
      // FIXME: Some padding is not included.
    }
  
    void beforeStart(unsigned funcId, const uint8_t *bytecodeStart) override {
      currentFuncId_ = funcId;
      emitted_.clear();
  
      opcodeStart_ = (uintptr_t)bytecodeStart;
      opcodeEnd_ = llvm::alignAddr(
          bytecodeStart +
              bcProvider_->getFunctionHeader(funcId).bytecodeSizeInBytes(),
          sizeof(uint32_t));
      functionEnd_ = opcodeEnd_;
  
      emitter_.emitKeyValue("functionId", funcId);
      emitFunctionLocation();
  
      emitter_.emitKey("usage");
      emitter_.openArray();
    }
  
    void countDebugInfo() {
      // FIXME: Avoid reimplementing this logic.
      auto *offsets = bcProvider_->getDebugOffsets(currentFuncId_);
      if (!offsets)
        return;
  
      if (offsets->sourceLocations != DebugOffsets::NO_OFFSET) {
        auto data = bcProvider_->getDebugInfo()->viewData().getData();
        auto offset = offsets->sourceLocations;
        int64_t n, trash;
        for (int i = 0; i < 3; i++) {
          offset += readSignedLEB128(data, offset, &n);
        }
        do {
          offset += readSignedLEB128(data, offset, &n);
          if (n == -1)
            break;
          offset += readSignedLEB128(data, offset, &n);
          if (n & 1)
            offset += readSignedLEB128(data, offset, &n);
          offset += readSignedLEB128(data, offset, &trash);
        } while (true);
        appendRecord(
            "debuginfo:sourcelocations",
            offsets->sourceLocations,
            offset - offsets->sourceLocations);
      }
  
      if (offsets->lexicalData &&
          offsets->lexicalData != DebugOffsets::NO_OFFSET) {
        auto data = bcProvider_->getDebugInfo()->viewData().getData();
        unsigned start = offsets->lexicalData +
            bcProvider_->getDebugInfo()->lexicalDataOffset();
        unsigned offset = start;
        int64_t trash;
  
        // Read parent id
        offset += readSignedLEB128(data, offset, &trash);
  
        // Read variable count
        int64_t count;
        offset += readSignedLEB128(data, offset, &count);
        // Read variables
        for (int64_t i = 0; i < count; i++) {
          int64_t stringLength;
          offset += readSignedLEB128(data, offset, &stringLength);
          offset += stringLength;
        }
        appendRecord(
            "debuginfo:lexicaldata", offsets->lexicalData, offset - start);
      }
    }
  
    void afterStart() override {
      auto header = bcProvider_->getFunctionHeader(currentFuncId_);
      auto exceptionTable = bcProvider_->getExceptionTable(currentFuncId_);
  
      // We always have a small header, and sometimes a large one too.
      appendRecord(
          "headers:function:small", currentFuncId_, sizeof(SmallFuncHeader));
      if (header.flags().overflowed) {
        appendRecord(
            "headers:function:large", currentFuncId_, sizeof(FunctionHeader));
      }
  
      countStringLiteral(header.functionName());
  
      if (header.flags().hasExceptionHandler) {
        // Exception tables are not deduplicated by function.
        appendRecord(
            "headers:exceptions",
            currentFuncId_,
            sizeof(ExceptionHandlerTableHeader));
        appendRecord(
            "bytecode:tables:exception",
            currentFuncId_,
            byteSize(exceptionTable));
      }
  
      appendRecord(
          "bytecode:instructions", header.offset(), opcodeEnd_ - opcodeStart_);
      appendRecord(
          "bytecode:tables:jump", header.offset(), functionEnd_ - opcodeEnd_);
  
      countDebugInfo();
  
      // Assign global headers to the global function
      if (bcProvider_->getGlobalFunctionIndex() == currentFuncId_) {
        emitGlobalInfo();
      }
      emitter_.closeArray();
    }
  
    void countStringKind(unsigned stringIndex) {
      // Map from string table index to kind index.
      auto it = std::upper_bound(
          stringKindEnds_.begin(), stringKindEnds_.end(), stringIndex);
  
      assert(it != stringKindEnds_.end() && "String index out of range");
      auto kindIndex = std::distance(stringKindEnds_.begin(), it);
      appendRecord("data:string:kind", kindIndex, sizeof(StringKind::Entry));
  
      StringKind::Kind kind = bcProvider_->getStringKinds()[kindIndex].kind();
      if (kind != StringKind::String) {
        // Strings whose kind are not "String" are Identifiers and have a
        // translation field.
        appendRecord(
            "data:string:identifier_translation", stringIndex, sizeof(uint32_t));
      }
    }
  
    void countStringLiteral(unsigned stringIndex) {
      countStringKind(stringIndex);
  
      auto entry = bcProvider_->getStringTableEntry(stringIndex);
      auto wasLarge = SmallStringTableEntry(entry, 0).isOverflowed();
  
      // Like functions headers, we have two kinds of string entries.
      appendRecord(
          "data:string:small_entry", stringIndex, sizeof(SmallStringTableEntry));
      if (wasLarge) {
        appendRecord(
            "data:string:overflow_entry",
            stringIndex,
            sizeof(OverflowStringTableEntry));
      }
  
      auto offset = entry.getOffset();
      auto length = entry.getLength() * (entry.isUTF16() ? 2 : 1);
  
      // Emit each byte separately for deduplication to allow overlap.
      // We could emit the entire length for non-overlapped strings, but those
      // make up maybe 10% of them so there's little to gain.
      for (unsigned i = 0; i < length; i++) {
        appendRecord("data:string:chars", offset + i, 1);
      }
    }
  
    void countRegex(uint32_t index) {
      if (index == 0xFFFFFFFF)
        return;
  
      auto regex = bcProvider_->getRegExpTable()[index];
      appendRecord("data:regex:entry", index, sizeof(RegExpTableEntry));
      appendRecord("data:regex:bytecode", index, regex.length);
    }
  
    void countSerializedLiteral(
        SLG::TagType tag,
        const unsigned char *buff,
        unsigned int *ind) {
      // FIXME: Avoid duplicating this logic.
  
      unsigned bundleOffset = (uintptr_t)(*ind + buff - bundleStart_);
      switch (tag) {
        case SLG::ByteStringTag: {
          uint8_t val = llvm::support::endian::read<uint8_t, 1>(
              buff + *ind, llvm::support::endianness::little);
          appendRecord("data:literalbuffer:bytestring", bundleOffset, 1);
          countStringLiteral(val);
          *ind += 1;
        } break;
        case SLG::ShortStringTag: {
          uint16_t val = llvm::support::endian::read<uint16_t, 1>(
              buff + *ind, llvm::support::endianness::little);
          appendRecord("data:literalbuffer:shortstring", bundleOffset, 2);
          countStringLiteral(val);
          *ind += 2;
        } break;
        case SLG::LongStringTag: {
          uint32_t val = llvm::support::endian::read<uint32_t, 1>(
              buff + *ind, llvm::support::endianness::little);
          appendRecord("data:literalbuffer:longstring", bundleOffset, 4);
          countStringLiteral(val);
          *ind += 4;
        } break;
        case SLG::NumberTag: {
          appendRecord("data:literalbuffer:double", bundleOffset, 8);
          *ind += 8;
        } break;
        case SLG::IntegerTag: {
          appendRecord("data:literalbuffer:int", bundleOffset, 4);
          *ind += 4;
        } break;
        case SLG::NullTag:
        case SLG::TrueTag:
        case SLG::FalseTag:
          break;
      }
    }
  
    void countSerializedLiterals(
        ArrayRef<unsigned char> array,
        unsigned offset,
        unsigned count) {
      const unsigned char *ptr = array.data();
      unsigned keyInd = offset;
      while (count > 0) {
        std::pair<int, SLG::TagType> keyTag = checkBufferTag(ptr + keyInd);
        auto tagLength = (keyTag.first > 0x0f ? 2 : 1);
        // This could conceivably overlap if the tag+data fits in the middle of
        // an existing 'double' literal, but we'll assume it's uncommon enough.
        appendRecord(
            "data:literalbuffer:tag",
            (unsigned)(uintptr_t)(keyInd + ptr - bundleStart_),
            tagLength);
        keyInd += tagLength;
        for (int i = 0; i < keyTag.first && count; i++) {
          countSerializedLiteral(keyTag.second, ptr, &keyInd);
          count--;
        }
      }
    }
  
    void visitSwitchImm(const inst::Inst *inst) {
      assert(inst->opCode == inst::OpCode::SwitchImm);
  
      const auto *curJmpTableView =
          reinterpret_cast<const uint32_t *>(llvm::alignAddr(
              (const uint8_t *)inst + inst->iSwitchImm.op2, sizeof(uint32_t)));
  
      unsigned start = inst->iSwitchImm.op4;
      unsigned end = inst->iSwitchImm.op5;
      assert(start < end && "Jump table spans negative range");
      unsigned count = end - start + 1;
  
      uintptr_t newEnd = (uintptr_t)&curJmpTableView[count];
      if (newEnd > functionEnd_)
        functionEnd_ = newEnd;
    }
  
    void preVisitInstruction(inst::OpCode opcode, const uint8_t *ip, int length)
        override {
      auto inst = (inst::Inst const *)ip;
  
  // Count all strings
  #define OPERAND_STRING_ID(OP, N)           \
    if (opcode == OpCode::OP) {              \
      countStringLiteral(inst->i##OP.op##N); \
    }
  #include "hermes/BCGen/HBC/BytecodeList.def"
  
      // Count non-string misc references
      switch (opcode) {
        case OpCode::SwitchImm:
          visitSwitchImm(inst);
          break;
        case OpCode::NewObjectWithBuffer:
          countSerializedLiterals(
              bcProvider_->getObjectKeyBuffer(),
              inst->iNewObjectWithBuffer.op4,
              inst->iNewObjectWithBuffer.op3);
          countSerializedLiterals(
              bcProvider_->getObjectValueBuffer(),
              inst->iNewObjectWithBuffer.op5,
              inst->iNewObjectWithBuffer.op3);
          break;
        case OpCode::NewObjectWithBufferLong:
          countSerializedLiterals(
              bcProvider_->getObjectKeyBuffer(),
              inst->iNewObjectWithBufferLong.op4,
              inst->iNewObjectWithBufferLong.op3);
          countSerializedLiterals(
              bcProvider_->getObjectValueBuffer(),
              inst->iNewObjectWithBufferLong.op5,
              inst->iNewObjectWithBufferLong.op3);
          break;
        case OpCode::NewArrayWithBuffer:
          countSerializedLiterals(
              bcProvider_->getArrayBuffer(),
              inst->iNewArrayWithBuffer.op4,
              inst->iNewArrayWithBuffer.op3);
          break;
        case OpCode::NewArrayWithBufferLong:
          countSerializedLiterals(
              bcProvider_->getArrayBuffer(),
              inst->iNewArrayWithBufferLong.op4,
              inst->iNewArrayWithBufferLong.op3);
          break;
        case OpCode::CreateRegExp:
          countRegex(inst->iCreateRegExp.op4);
        default:
          break;
      }
    }
  
   public:
    UsageCounter(
        std::shared_ptr<BCProvider> bc,
        JSONEmitter &emitter,
        llvm::DenseMap<unsigned, unsigned> offsets,
        uintptr_t bundleStart)
        : BytecodeVisitor(bc),
          emitter_(emitter),
          virtualOffsets_(offsets),
          bundleStart_(bundleStart) {
      unsigned end = 0;
      for (auto entry : bc->getStringKinds()) {
        end += entry.count();
        stringKindEnds_.push_back(end);
      }
    }
  };
  
  // Getting all virtual offsets is O(N^2) unless we do them in a single pass.
  llvm::DenseMap<unsigned, unsigned> getVirtualOffsets(
      std::shared_ptr<BCProvider> bc) {
    llvm::DenseMap<unsigned, unsigned> map(bc->getFunctionCount());
  
    unsigned virtualOffset = 0;
    for (unsigned i = 0, e = bc->getFunctionCount(); i < e; i++) {
      auto header = bc->getFunctionHeader(i);
      map[i] = virtualOffset;
      virtualOffset += header.bytecodeSizeInBytes();
    }
    return map;
  }
  
  bool attribute(
      std::unique_ptr<llvm::MemoryBuffer> input,
      JSONEmitter &emitter) {
    assert(
        llvm::alignAddr(input->getBufferStart(), sizeof(uint32_t)) ==
        (uintptr_t)input->getBufferStart());
  
    uintptr_t bundleStart = (uintptr_t)input->getBuffer().data();
    auto hermesBuffer = llvm::make_unique<hermes::MemoryBuffer>(input.get());
    auto ret = hbc::BCProviderFromBuffer::createBCProviderFromBuffer(
        std::move(hermesBuffer));
    if (!ret.first) {
      llvm::errs() << "Can't deserialize: " << ret.second << "\n";
      return false;
    }
    std::shared_ptr<BCProvider> bc = std::move(ret.first);
  
    // TODO: Add records for the bytecode header and similar.
    UsageCounter counter(bc, emitter, getVirtualOffsets(bc), bundleStart);
    for (int i = 0, e = bc->getFunctionCount(); i < e; i++) {
      emitter.openDict();
      counter.visitInstructionsInFunction(i);
      emitter.closeDict();
      emitter.endJSONL();
    }
    return true;
  }
  
  } // namespace
  
  int main(int argc, char **argv) {
    // Normalize the arg vector.
    llvm::InitLLVM initLLVM(argc, argv);
    llvm::sys::PrintStackTraceOnErrorSignal("hbc-attribute");
    llvm::PrettyStackTraceProgram X(argc, argv);
    llvm::llvm_shutdown_obj Y;
    llvm::cl::ParseCommandLineOptions(
        argc, argv, "Hermes bytecode size attribution tool\n");
  
    llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileBufOrErr =
        llvm::MemoryBuffer::getFileOrSTDIN(InputFilename);
    if (!fileBufOrErr) {
      llvm::errs() << "Error: fail to open file: " << InputFilename << ": "
                   << fileBufOrErr.getError().message() << "\n";
      return 1;
    }
  
    llvm::Optional<raw_fd_ostream> fileOS;
    if (!OutputFilename.empty()) {
      std::error_code EC;
      fileOS.emplace(OutputFilename.data(), EC, llvm::sys::fs::F_Text);
      if (EC) {
        llvm::errs() << "Error: fail to open file " << OutputFilename << ": "
                     << EC.message() << '\n';
        return 1;
      }
    }
  
    auto &output = fileOS ? *fileOS : llvm::outs();
    JSONEmitter emitter(output);
    if (!attribute(std::move(fileBufOrErr.get()), emitter)) {
      return 3;
    }
    output.flush();
  
    return 0;
  }
  