//===- SimpliyMath.cpp - Simplify Mathematical operations  ------------------ //
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower custom ops generated by the Enzyme AD
// procedure to the MemRef dialect.
//===----------------------------------------------------------------------===//

#include "Dialect/Ops.h"
#include "Interfaces/AutoDiffOpInterface.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"

#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

using namespace mlir;
using namespace mlir::enzyme;

namespace mlir {
namespace enzyme {
using namespace mlir::enzyme;
#define GEN_PASS_DEF_MATHEMATICSIMPLIFICATIONPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

using llvm::errs;
namespace {

struct ApplySimplificationPattern
    : public OpInterfaceRewritePattern<enzyme::MathSimplifyInterface> {
  using OpInterfaceRewritePattern<
      enzyme::MathSimplifyInterface>::OpInterfaceRewritePattern;

  LogicalResult matchAndRewrite(enzyme::MathSimplifyInterface op,
                                PatternRewriter &rewriter) const override {
    return op.simplifyMath(rewriter);
  }
};

struct MathematicSimplification
    : public enzyme::impl::MathematicSimplificationPassBase<
          MathematicSimplification> {
  void runOnOperation() override {

    RewritePatternSet patterns(&getContext());
    patterns.insert<ApplySimplificationPattern>(&getContext());

    GreedyRewriteConfig config;
    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
                                       config);
  };
};
} // end anonymous namespace
