#include "llvm/ADT/STLExtras.h"

#include "src/Accelerators/PIM/Conversion/SpatialToPim/Cleanup.hpp"

using namespace mlir;

namespace onnx_mlir {

LogicalResult erasePendingOps(SmallVectorImpl<Operation*>& pendingOps, IRRewriter& rewriter) {
  while (!pendingOps.empty()) {
    bool erasedAnyOp = false;
    for (auto it = pendingOps.begin(); it != pendingOps.end();) {
      Operation* opToRemove = *it;
      if (!opToRemove->use_empty()) {
        ++it;
        continue;
      }

      rewriter.eraseOp(opToRemove);
      it = pendingOps.erase(it);
      erasedAnyOp = true;
    }

    if (erasedAnyOp)
      continue;

    for (Operation* opToRemove : pendingOps) {
      InFlightDiagnostic diag = opToRemove->emitError("pending Spatial-to-PIM cleanup could not erase operation");
      diag << "; op has " << llvm::range_size(opToRemove->getUsers()) << " remaining user(s)";
      for (Operation* user : opToRemove->getUsers()) {
        bool userPendingRemoval = llvm::is_contained(pendingOps, user);
        opToRemove->emitRemark() << "remaining user `" << user->getName() << "`"
                                 << (userPendingRemoval ? " is also pending removal" : " is not pending removal");
      }
    }
    return failure();
  }

  return success();
}

} // namespace onnx_mlir