//===- OpStats.cpp - Prints stats of operations in module -----------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Transforms/Passes.h"

#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "llvm/Support/Format.h"
#include "llvm/Support/raw_ostream.h"

namespace mlir {
#define GEN_PASS_DEF_PRINTOPSTATSPASS
#include "mlir/Transforms/Passes.h.inc"
} // namespace mlir

using namespace mlir;

namespace {
struct PrintOpStatsPass : public impl::PrintOpStatsPassBase<PrintOpStatsPass> {
  using impl::PrintOpStatsPassBase<PrintOpStatsPass>::PrintOpStatsPassBase;

  explicit PrintOpStatsPass(raw_ostream &os) : os(&os) {}

  explicit PrintOpStatsPass(raw_ostream &os, bool printAsJSON) : os(&os) {
    this->printAsJSON = printAsJSON;
  }

  // Prints the resultant operation statistics post iterating over the module.
  void runOnOperation() override;

  // Print summary of op stats.
  void printSummary();

  // Print symmary of op stats in JSON.
  void printSummaryInJSON();

private:
  llvm::StringMap<int64_t> opCount;
  raw_ostream *os = &llvm::errs();
};
} // namespace

void PrintOpStatsPass::runOnOperation() {
  opCount.clear();

  // Compute the operation statistics for the currently visited operation.
  getOperation()->walk(
      [&](Operation *op) { ++opCount[op->getName().getStringRef()]; });
  if (printAsJSON)
    printSummaryInJSON();
  else
    printSummary();
  markAllAnalysesPreserved();
}

void PrintOpStatsPass::printSummary() {
  *os << "Operations encountered:\n";
  *os << "-----------------------\n";
  SmallVector<StringRef, 64> sorted(opCount.keys());
  llvm::sort(sorted);

  // Split an operation name from its dialect prefix.
  auto splitOperationName = [](StringRef opName) {
    auto splitName = opName.split('.');
    return splitName.second.empty() ? std::make_pair("", splitName.first)
                                    : splitName;
  };

  // Compute the largest dialect and operation name.
  size_t maxLenOpName = 0, maxLenDialect = 0;
  for (const auto &key : sorted) {
    auto [dialectName, opName] = splitOperationName(key);
    maxLenDialect = std::max(maxLenDialect, dialectName.size());
    maxLenOpName = std::max(maxLenOpName, opName.size());
  }

  for (const auto &key : sorted) {
    auto [dialectName, opName] = splitOperationName(key);

    // Left-align the names (aligning on the dialect) and right-align the count
    // below. The alignment is for readability and does not affect CSV/FileCheck
    // parsing.
    if (dialectName.empty())
      os->indent(maxLenDialect + 3);
    else
      *os << llvm::right_justify(dialectName, maxLenDialect + 2) << '.';

    // Left justify the operation name.
    *os << llvm::left_justify(opName, maxLenOpName) << " , " << opCount[key]
        << '\n';
  }
}

void PrintOpStatsPass::printSummaryInJSON() {
  SmallVector<StringRef, 64> sorted(opCount.keys());
  llvm::sort(sorted);

  *os << "{\n";

  for (unsigned i = 0, e = sorted.size(); i != e; ++i) {
    const auto &key = sorted[i];
    *os << "  \"" << key << "\" : " << opCount[key];
    if (i != e - 1)
      *os << ",\n";
    else
      *os << "\n";
  }
  *os << "}\n";
}

std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os) {
  return std::make_unique<PrintOpStatsPass>(os);
}

std::unique_ptr<Pass> mlir::createPrintOpStatsPass(raw_ostream &os,
                                                   bool printAsJSON) {
  return std::make_unique<PrintOpStatsPass>(os, printAsJSON);
}