//===- AffineStructuresTest.cpp - Tests for AffineStructures ----*- C++ -*-===//

//

// 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/Analysis/AffineStructures.h"

#include <gmock/gmock.h>

#include <gtest/gtest.h>

#include <numeric>

namespace mlir {

/// Evaluate the value of the given affine expression at the specified point.

/// The expression is a list of coefficients for the dimensions followed by the

/// constant term.

int64_t valueAt(ArrayRef<int64_t> expr, ArrayRef<int64_t> point) {

  assert(expr.size() == 1 + point.size());

  int64_t value = expr.back();

  for (unsigned i = 0; i < point.size(); ++i)

    value += expr[i] * point[i];

  return value;

}

/// If hasValue is true, check that findIntegerSample returns a valid sample

/// for the FlatAffineConstraints fac.

///

/// If hasValue is false, check that findIntegerSample does not return None.

void checkSample(bool hasValue, const FlatAffineConstraints &fac) {

  Optional<SmallVector<int64_t, 8>> maybeSample = fac.findIntegerSample();

  if (!hasValue) {

    EXPECT_FALSE(maybeSample.hasValue());

    if (maybeSample.hasValue()) {

      for (auto x : *maybeSample)

        llvm::errs() << x << ' ';

      llvm::errs() << '\n';

    }

  } else {

    ASSERT_TRUE(maybeSample.hasValue());

    for (unsigned i = 0; i < fac.getNumEqualities(); ++i)

      EXPECT_EQ(valueAt(fac.getEquality(i), *maybeSample), 0);

    for (unsigned i = 0; i < fac.getNumInequalities(); ++i)

      EXPECT_GE(valueAt(fac.getInequality(i), *maybeSample), 0);

  }

}

/// Construct a FlatAffineConstraints from a set of inequality and

/// equality constraints.

FlatAffineConstraints

makeFACFromConstraints(unsigned dims,

                       SmallVector<SmallVector<int64_t, 4>, 4> ineqs,

                       SmallVector<SmallVector<int64_t, 4>, 4> eqs) {

  FlatAffineConstraints fac(ineqs.size(), eqs.size(), dims + 1, dims);

  for (const auto &eq : eqs)

    fac.addEquality(eq);

  for (const auto &ineq : ineqs)

    fac.addInequality(ineq);

  return fac;

}

/// Check sampling for all the permutations of the dimensions for the given

/// constraint set. Since the GBR algorithm progresses dimension-wise, different

/// orderings may cause the algorithm to proceed differently. At lesast some of

///.these permutations should make it past the heuristics and test the

/// implementation of the GBR algorithm itself.

void checkPermutationsSample(bool hasValue, unsigned nDim,

                             SmallVector<SmallVector<int64_t, 4>, 4> ineqs,

                             SmallVector<SmallVector<int64_t, 4>, 4> eqs) {

  SmallVector<unsigned, 4> perm(nDim);

  std::iota(perm.begin(), perm.end(), 0);

  auto permute = [&perm](ArrayRef<int64_t> coeffs) {

    SmallVector<int64_t, 4> permuted;

    for (unsigned id : perm)

      permuted.push_back(coeffs[id]);

    permuted.push_back(coeffs.back());

    return permuted;

  };

  do {

    SmallVector<SmallVector<int64_t, 4>, 4> permutedIneqs, permutedEqs;

    for (const auto &ineq : ineqs)

      permutedIneqs.push_back(permute(ineq));

    for (const auto &eq : eqs)

      permutedEqs.push_back(permute(eq));

    checkSample(hasValue,

                makeFACFromConstraints(nDim, permutedIneqs, permutedEqs));

  } while (std::next_permutation(perm.begin(), perm.end()));

}

TEST(FlatAffineConstraintsTest, FindSampleTest) {

  // Bounded sets with only inequalities.

  // 0 <= 7x <= 5

  checkSample(true, makeFACFromConstraints(1, {{7, 0}, {-7, 5}}, {}));

  // 1 <= 5x and 5x <= 4 (no solution).

  checkSample(false, makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}));

  // 1 <= 5x and 5x <= 9 (solution: x = 1).

  checkSample(true, makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}));

  // Bounded sets with equalities.

  // x >= 8 and 40 >= y and x = y.

  checkSample(

      true, makeFACFromConstraints(2, {{1, 0, -8}, {0, -1, 40}}, {{1, -1, 0}}));

  // x <= 10 and y <= 10 and 10 <= z and x + 2y = 3z.

  // solution: x = y = z = 10.

  checkSample(true, makeFACFromConstraints(

                        3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -10}},

                        {{1, 2, -3, 0}}));

  // x <= 10 and y <= 10 and 11 <= z and x + 2y = 3z.

  // This implies x + 2y >= 33 and x + 2y <= 30, which has no solution.

  checkSample(false, makeFACFromConstraints(

                         3, {{-1, 0, 0, 10}, {0, -1, 0, 10}, {0, 0, 1, -11}},

                         {{1, 2, -3, 0}}));

  // 0 <= r and r <= 3 and 4q + r = 7.

  // Solution: q = 1, r = 3.

  checkSample(true,

              makeFACFromConstraints(2, {{0, 1, 0}, {0, -1, 3}}, {{4, 1, -7}}));

  // 4q + r = 7 and r = 0.

  // Solution: q = 1, r = 3.

  checkSample(false, makeFACFromConstraints(2, {}, {{4, 1, -7}, {0, 1, 0}}));

  // The next two sets are large sets that should take a long time to sample

  // with a naive branch and bound algorithm but can be sampled efficiently with

  // the GBR algroithm.

  //

  // This is a triangle with vertices at (1/3, 0), (2/3, 0) and (10000, 10000).

  checkSample(

      true,

      makeFACFromConstraints(

          2, {{0, 1, 0}, {300000, -299999, -100000}, {-300000, 299998, 200000}},

          {}));

  // This is a tetrahedron with vertices at

  // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 10000), and (10000, 10000, 10000).

  // The first three points form a triangular base on the xz plane with the

  // apex at the fourth point, which is the only integer point.

  checkPermutationsSample(

      true, 3,

      {

          {0, 1, 0, 0},  // y >= 0

          {0, -1, 1, 0}, // z >= y

          {300000, -299998, -1,

           -100000},                    // -300000x + 299998y + 100000 + z <= 0.

          {-150000, 149999, 0, 100000}, // -150000x + 149999y + 100000 >= 0.

      },

      {});

  // Same thing with some spurious extra dimensions equated to constants.

  checkSample(true,

              makeFACFromConstraints(

                  5,

                  {

                      {0, 1, 0, 1, -1, 0},

                      {0, -1, 1, -1, 1, 0},

                      {300000, -299998, -1, -9, 21, -112000},

                      {-150000, 149999, 0, -15, 47, 68000},

                  },

                  {{0, 0, 0, 1, -1, 0},       // p = q.

                   {0, 0, 0, 1, 1, -2000}})); // p + q = 20000 => p = q = 10000.

  // This is a tetrahedron with vertices at

  // (1/3, 0, 0), (2/3, 0, 0), (2/3, 0, 100), (100, 100 - 1/3, 100).

  checkPermutationsSample(false, 3,

                          {

                              {0, 1, 0, 0},

                              {0, -300, 299, 0},

                              {300 * 299, -89400, -299, -100 * 299},

                              {-897, 894, 0, 598},

                          },

                          {});

  // Two tests involving equalities that are integer empty but not rational

  // empty.

  // This is a line segment from (0, 1/3) to (100, 100 + 1/3).

  checkSample(false, makeFACFromConstraints(

                         2,

                         {

                             {1, 0, 0},   // x >= 0.

                             {-1, 0, 100} // -x + 100 >= 0, i.e., x <= 100.

                         },

                         {

                             {3, -3, 1} // 3x - 3y + 1 = 0, i.e., y = x + 1/3.

                         }));

  // A thin parallelogram. 0 <= x <= 100 and x + 1/3 <= y <= x + 2/3.

  checkSample(false, makeFACFromConstraints(2,

                                            {

                                                {1, 0, 0},    // x >= 0.

                                                {-1, 0, 100}, // x <= 100.

                                                {3, -3, 2},   // 3x - 3y >= -2.

                                                {-3, 3, -1},  // 3x - 3y <= -1.

                                            },

                                            {}));

  checkSample(true, makeFACFromConstraints(2,

                                           {

                                               {2, 0, 0},   // 2x >= 1.

                                               {-2, 0, 99}, // 2x <= 99.

                                               {0, 2, 0},   // 2y >= 0.

                                               {0, -2, 99}, // 2y <= 99.

                                           },

                                           {}));

}

TEST(FlatAffineConstraintsTest, IsIntegerEmptyTest) {

  // 1 <= 5x and 5x <= 4 (no solution).

  EXPECT_TRUE(

      makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}).isIntegerEmpty());

  // 1 <= 5x and 5x <= 9 (solution: x = 1).

  EXPECT_FALSE(

      makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}).isIntegerEmpty());

  // An unbounded set, which isIntegerEmpty should detect as unbounded and

  // return without calling findIntegerSample.

  EXPECT_FALSE(makeFACFromConstraints(3,

                                      {

                                          {2, 0, 0, -1},

                                          {-2, 0, 0, 1},

                                          {0, 2, 0, -1},

                                          {0, -2, 0, 1},

                                          {0, 0, 2, -1},

                                      },

                                      {})

                   .isIntegerEmpty());

  // FlatAffineConstraints::isEmpty() does not detect the following sets to be

  // empty.

  // 3x + 7y = 1 and 0 <= x, y <= 10.

  // Since x and y are non-negative, 3x + 7y can never be 1.

  EXPECT_TRUE(

      makeFACFromConstraints(

          2, {{1, 0, 0}, {-1, 0, 10}, {0, 1, 0}, {0, -1, 10}}, {{3, 7, -1}})

          .isIntegerEmpty());

  // 2x = 3y and y = x - 1 and x + y = 6z + 2 and 0 <= x, y <= 100.

  // Substituting y = x - 1 in 3y = 2x, we obtain x = 3 and hence y = 2.

  // Since x + y = 5 cannot be equal to 6z + 2 for any z, the set is empty.

  EXPECT_TRUE(

      makeFACFromConstraints(3,

                             {

                                 {1, 0, 0, 0},

                                 {-1, 0, 0, 100},

                                 {0, 1, 0, 0},

                                 {0, -1, 0, 100},

                             },

                             {{2, -3, 0, 0}, {1, -1, 0, -1}, {1, 1, -6, -2}})

          .isIntegerEmpty());

  // 2x = 3y and y = x - 1 + 6z and x + y = 6q + 2 and 0 <= x, y <= 100.

  // 2x = 3y implies x is a multiple of 3 and y is even.

  // Now y = x - 1 + 6z implies y = 2 mod 3. In fact, since y is even, we have

  // y = 2 mod 6. Then since x = y + 1 + 6z, we have x = 3 mod 6, implying

  // x + y = 5 mod 6, which contradicts x + y = 6q + 2, so the set is empty.

  EXPECT_TRUE(makeFACFromConstraints(

                  4,

                  {

                      {1, 0, 0, 0, 0},

                      {-1, 0, 0, 0, 100},

                      {0, 1, 0, 0, 0},

                      {0, -1, 0, 0, 100},

                  },

                  {{2, -3, 0, 0, 0}, {1, -1, 6, 0, -1}, {1, 1, 0, -6, -2}})

                  .isIntegerEmpty());

}

} // namespace mlir