CompilerPipeline.cpp

/*
 * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
 * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
 * All rights reserved.
 *
 * SPDX-License-Identifier: MIT
 *
 * Licensed under the MIT License
 */

#include "mlir/Compiler/CompilerPipeline.h"

#include "mlir/Conversion/QCOToQC/QCOToQC.h"
#include "mlir/Conversion/QCToQCO/QCToQCO.h"
#include "mlir/Conversion/QCToQIR/QIRAdaptive/QCToQIRAdaptive.h"
#include "mlir/Conversion/QCToQIR/QIRBase/QCToQIRBase.h"
#include "mlir/Dialect/QCO/Transforms/Passes.h"
#include "mlir/Support/Passes.h"
#include "mlir/Support/PrettyPrinting.h"

#include <llvm/Support/raw_ostream.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/Pass/PassManager.h>
#include <mlir/Support/LLVM.h>

#include <string>

namespace mlir {

/**
 * @brief Pretty print IR with ASCII art borders and stage identifier
 *
 * @param module The module to print
 * @param stageName Name of the compilation stage
 * @param stageNumber Current stage number
 * @param totalStages Total number of stages (for progress indication)
 */
static void prettyPrintStage(ModuleOp module, const StringRef stageName,
                             const int stageNumber, const int totalStages) {
  llvm::errs() << "\n";

  // Build the stage header
  const std::string stageHeader = "Stage " + std::to_string(stageNumber) + "/" +
                                  std::to_string(totalStages) + ": " +
                                  stageName.str();
  printProgram(module, stageHeader, llvm::errs());
}

void QuantumCompilerPipeline::configurePassManager(PassManager& pm) const {
  // Enable timing statistics if requested
  if (config_.enableTiming) {
    pm.enableTiming();
  }

  // Enable pass statistics if requested
  if (config_.enableStatistics) {
    pm.enableStatistics();
  }
}

LogicalResult
QuantumCompilerPipeline::runPipeline(ModuleOp module,
                                     CompilationRecord* record) const {
  if (config_.convertToQIRBase && config_.convertToQIRAdaptive) {
    llvm::errs()
        << "convertToQIRBase and convertToQIRAdaptive are mutually "
           "exclusive; only one QIR profile can be targeted at a time.\n";
    return failure();
  }
  const auto convertToQIR =
      config_.convertToQIRAdaptive || config_.convertToQIRBase;

  // Ensure printIRAfterAllStages implies recordIntermediates
  if (config_.printIRAfterAllStages &&
      (!config_.recordIntermediates || record == nullptr)) {
    llvm::errs() << "printIRAfterAllStages requires recordIntermediates to be "
                    "enabled and the record pointer to be non-null.\n";
    return failure();
  }

  auto runStage = [&](auto&& populatePasses) -> LogicalResult {
    PassManager pm(module.getContext());
    configurePassManager(pm);
    populatePasses(pm);
    return pm.run(module);
  };

  // Determine total number of stages for progress indication
  // 1. QC import
  // 2. QC cleanup
  // 3. QC-to-QCO conversion
  // 4. QCO cleanup
  // 5. Optimization passes
  // 6. QCO cleanup
  // 7. QCO-to-QC conversion
  // 8. QC cleanup
  // 9. QC-to-QIR conversion (optional)
  // 10. QIR cleanup (optional)
  auto totalStages = 8;
  if (convertToQIR) {
    totalStages += 2;
  }
  auto currentStage = 0;

  // Stage 1: QC import
  if (record != nullptr && config_.recordIntermediates) {
    record->afterQCImport = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "QC Import", ++currentStage, totalStages);
    }
  }

  // Stage 2: QC cleanup
  if (failed(
          runStage([&](PassManager& pm) { populateQCCleanupPipeline(pm); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterInitialCanon = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "Initial QC Cleanup", ++currentStage,
                       totalStages);
    }
  }
  // Stage 3: QC-to-QCO conversion
  if (failed(runStage([&](PassManager& pm) { pm.addPass(createQCToQCO()); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterQCOConversion = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "QC → QCO Conversion", ++currentStage,
                       totalStages);
    }
  }
  // Stage 4: QCO cleanup
  if (failed(
          runStage([&](PassManager& pm) { populateQCOCleanupPipeline(pm); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterQCOCanon = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "Initial QCO Cleanup", ++currentStage,
                       totalStages);
    }
  }
  // Stage 5: Optimization passes
  if (failed(runStage([&](PassManager& pm) {
        if (!config_.disableMergeSingleQubitRotationGates) {
          pm.addPass(qco::createMergeSingleQubitRotationGates());
        }
        if (config_.enableHadamardLifting) {
          pm.addPass(qco::createHadamardLifting());
        }
      }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterOptimization = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "Optimization Passes", ++currentStage,
                       totalStages);
    }
  }
  // Stage 6: QCO cleanup
  if (failed(
          runStage([&](PassManager& pm) { populateQCOCleanupPipeline(pm); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterOptimizationCanon = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "Post-Optimization QCO Cleanup", ++currentStage,
                       totalStages);
    }
  }
  // Stage 7: Transpilation passes (optional)
  if (config_.device != nullptr) {
    if (failed(runStage([&](PassManager& pm) {
          /// TODO:
          // Individual passes use the device handle to query properties.
          // if (device.hasCouplingMap())
          //     pm.addPass(createMappingPass(config_.device))
          // }
          // pm.addPass(createNativeGateDecompositionPass(config_.device))
        }))) {
      return failure();
    }
    if (record != nullptr && config_.recordIntermediates) {
      record->afterTranspilation = captureIR(module);
      if (config_.printIRAfterAllStages) {
        prettyPrintStage(module, "Transpilation Passes", ++currentStage,
                         totalStages);
      }
    }
    // Stage 8: QCO cleanup (optional)
    if (failed(runStage(
            [&](PassManager& pm) { populateQCOCleanupPipeline(pm); }))) {
      return failure();
    }
    if (record != nullptr && config_.recordIntermediates) {
      record->afterTranspilationCanon = captureIR(module);
      if (config_.printIRAfterAllStages) {
        prettyPrintStage(module, "Post-Transpilation QCO Cleanup",
                         ++currentStage, totalStages);
      }
    }
  }

  // Stage 9: QCO-to-QC conversion
  if (failed(runStage([&](PassManager& pm) { pm.addPass(createQCOToQC()); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterQCConversion = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "QCO → QC Conversion", ++currentStage,
                       totalStages);
    }
  }
  // Stage 10: QC cleanup
  if (failed(
          runStage([&](PassManager& pm) { populateQCCleanupPipeline(pm); }))) {
    return failure();
  }
  if (record != nullptr && config_.recordIntermediates) {
    record->afterQCCanon = captureIR(module);
    if (config_.printIRAfterAllStages) {
      prettyPrintStage(module, "Final QC Cleanup", ++currentStage, totalStages);
    }
  }
  // Stage 9: QC-to-QIR conversion (optional)
  if (convertToQIR) {
    auto addConversionPass = [&](PassManager& pm) {
      if (config_.convertToQIRBase) {
        pm.addPass(createQCToQIRBase());
      } else {
        pm.addPass(createQCToQIRAdaptive());
      }
    };

    if (failed(runStage(addConversionPass))) {
      return failure();
    }

    if (record != nullptr && config_.recordIntermediates) {
      record->afterQIRConversion = captureIR(module);
      if (config_.printIRAfterAllStages) {
        prettyPrintStage(module, "QC → QIR Conversion", ++currentStage,
                         totalStages);
      }
    }
    // Stage 12: QIR cleanup (optional)
    if (failed(runStage(
            [&](PassManager& pm) { populateQIRCleanupPipeline(pm); }))) {
      return failure();
    }
    if (record != nullptr && config_.recordIntermediates) {
      record->afterQIRCanon = captureIR(module);
      if (config_.printIRAfterAllStages) {
        prettyPrintStage(module, "QIR Cleanup", ++currentStage, totalStages);
      }
    }
  }

  // Print compilation summary
  if (config_.printIRAfterAllStages) {
    llvm::errs() << "\n";
    printBoxTop();

    printBoxLine("✓ Compilation Complete");

    const std::string summaryLine =
        "Successfully processed " + std::to_string(currentStage) + " stages";
    printBoxLine(summaryLine, 1); // Indent by 1 space

    printBoxBottom();
    llvm::errs() << "\n";
    llvm::errs().flush();
  }
  return success();
}

std::string captureIR(ModuleOp module) {
  std::string result;
  llvm::raw_string_ostream os(result);
  module.print(os);
  return result;
}

} // namespace mlir