/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include <gtest/gtest.h>
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/runtime/logging.h>
#include <tvm/te/operation.h>

TEST(Simplify, MinMax) {
  tvm::arith::Analyzer ana;
  auto x = tvm::te::var("x");
  auto e1 = (tvm::max(x, 1) - tvm::max(x, 1));
  auto e1s = ana->canonical_simplify(e1);
  TVM_FFI_ICHECK(tvm::tirx::is_zero(e1s));

  auto e2 = (x * tvm::min(x, 1)) - (x * tvm::min(x, 1));
  auto e2s = ana->canonical_simplify(e2);
  TVM_FFI_ICHECK(tvm::tirx::is_zero(e2s));
}

TEST(Simplify, Mul) {
  tvm::arith::Analyzer ana;
  auto x = tvm::te::var("x");
  auto e = (x * x) - (x * x);
  auto es = ana->canonical_simplify(e);
  TVM_FFI_ICHECK(tvm::tirx::is_zero(es));
}

TEST(Simplify, Mod) {
  tvm::arith::Analyzer ana;
  auto x = tvm::IntImm::Int32(10);
  auto y = tvm::IntImm::Int32(12);
  // Mod::make is used instead of % to avoid constant folding during
  // calling operator%(x,y). Mod::make doesn't try constant folding,
  // and therefore, the constant folding will be attempted in CanonicalSimplify
  auto mod = ana->canonical_simplify(tvm::tirx::Mod(x, y));
  auto es = ana->canonical_simplify(mod - x);
  TVM_FFI_ICHECK(tvm::tirx::is_zero(es));
}

TEST(AnalyzerObjectRef, CopySharesMutableState) {
  tvm::arith::Analyzer analyzer;
  tvm::arith::Analyzer copy = analyzer;
  auto x = tvm::te::var("x");

  copy->Bind(x, tvm::Range::FromMinExtent(0, 8));

  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
}

TEST(AnalyzerObjectRef, ConstHandleRefCanMutateAnalyzerState) {
  tvm::arith::Analyzer analyzer;
  const tvm::arith::Analyzer& analyzer_ref = analyzer;
  auto x = tvm::te::var("x");

  analyzer_ref->Bind(x, tvm::Range::FromMinExtent(0, 8));

  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
}

TEST(AnalyzerObjectRef, CloneIsIndependent) {
  tvm::arith::Analyzer analyzer;
  auto x = tvm::te::var("x");
  auto y = tvm::te::var("y");

  analyzer->Bind(x, tvm::Range::FromMinExtent(0, 8));
  analyzer->modular_set.Update(x, tvm::arith::ModularSet(4, 0));

  tvm::arith::Analyzer clone = analyzer->Clone();
  TVM_FFI_ICHECK(clone->CanProve(x < 8));
  TVM_FFI_ICHECK(clone->modular_set(x)->coeff == 4);

  clone->Bind(y, tvm::Range::FromMinExtent(0, 4));
  clone->modular_set.Update(x, tvm::arith::ModularSet(8, 0), true);
  TVM_FFI_ICHECK(clone->CanProve(y < 4));
  TVM_FFI_ICHECK(!analyzer->CanProve(y < 4));
  TVM_FFI_ICHECK(analyzer->CanProve(x < 8));
  TVM_FFI_ICHECK(analyzer->modular_set(x)->coeff == 4);
  TVM_FFI_ICHECK(clone->modular_set(x)->coeff == 8);
}

TEST(ConstantFold, Broadcast) {
  tvm::ffi::StructuralEqual checker;
  auto i32x4 = tvm::tirx::Broadcast(tvm::IntImm::Int32(10), 4);
  auto i64x4 = tvm::cast(i32x4.ty().WithBits(64), i32x4);
  auto i64x4_expected = tvm::tirx::Broadcast(tvm::IntImm::Int64(10), 4);
  ASSERT_TRUE(checker(i64x4, i64x4_expected));
}

TEST(ConstantFold, Ramp) {
  tvm::ffi::StructuralEqual checker;
  auto i32x4 = tvm::tirx::Ramp(tvm::IntImm::Int32(10), tvm::IntImm::Int32(1), 4);
  auto i64x4 = tvm::cast(i32x4.ty().WithBits(64), i32x4);
  auto i64x4_expected = tvm::tirx::Ramp(tvm::IntImm::Int64(10), tvm::IntImm::Int64(1), 4);
  ASSERT_TRUE(checker(i64x4, i64x4_expected));

  auto f32x4 = tvm::cast(tvm::PrimType::Float(32, 4), i32x4);
  auto f32x4_expected = tvm::tirx::Cast(tvm::PrimType::Float(32, 4), i32x4);
  ASSERT_TRUE(checker(f32x4, f32x4_expected));
}