/******************************************************************************
 * Top contributors (to current version):
 *   Morgan Deters, Tim King, Aina Niemetz
 *
 * This file is part of the cvc5 project.
 *
 * Copyright (c) 2009-2025 by the authors listed in the file AUTHORS
 * in the top-level source directory and their institutional affiliations.
 * All rights reserved.  See the file COPYING in the top-level source
 * directory for licensing information.
 * ****************************************************************************
 *
 * Enumerators for rationals and integers.
 */

#include "cvc5_private.h"

#ifndef CVC5__THEORY__ARITH__TYPE_ENUMERATOR_H
#define CVC5__THEORY__ARITH__TYPE_ENUMERATOR_H

#include "expr/kind.h"
#include "expr/type_node.h"
#include "theory/type_enumerator.h"
#include "util/integer.h"
#include "util/rational.h"

namespace cvc5::internal {
namespace theory {
namespace arith {

class RationalEnumerator : public TypeEnumeratorBase<RationalEnumerator> {
  Rational d_rat;

 public:
  RationalEnumerator(TypeNode type, TypeEnumeratorProperties* tep = nullptr)
      : TypeEnumeratorBase<RationalEnumerator>(type), d_rat(0)
  {
    Assert(type.getKind() == Kind::TYPE_CONSTANT
           && type.getConst<TypeConstant>() == REAL_TYPE);
  }

  Node operator*() override
  {
    return getType().getNodeManager()->mkConstReal(d_rat);
  }
  RationalEnumerator& operator++() override
  {
    // sequence is 0, then diagonal with negatives interleaved
    // ( 0, 1/1, -1/1, 2/1, -2/1, 1/2, -1/2, 3/1, -3/1, 1/3, -1/3,
    // 4/1, -4/1, 3/2, -3/2, 2/3, -2/3, 1/4, -1/4, ...)
    if(d_rat == 0) {
      d_rat = 1;
    } else if(d_rat < 0) {
      d_rat = -d_rat;
      Integer num = d_rat.getNumerator();
      Integer den = d_rat.getDenominator();
      do {
        num -= 1;
        den += 1;
        if(num == 0) {
          num = den;
          den = 1;
        }
        d_rat = Rational(num, den);
      } while(d_rat.getNumerator() != num);
    } else {
      d_rat = -d_rat;
    }
    return *this;
  }

  bool isFinished() override { return false; }
};/* class RationalEnumerator */

class IntegerEnumerator : public TypeEnumeratorBase<IntegerEnumerator> {
  Integer d_int;

 public:
  IntegerEnumerator(TypeNode type, TypeEnumeratorProperties* tep = nullptr)
      : TypeEnumeratorBase<IntegerEnumerator>(type), d_int(0)
  {
    Assert(type.getKind() == Kind::TYPE_CONSTANT
           && type.getConst<TypeConstant>() == INTEGER_TYPE);
  }

  Node operator*() override
  {
    return getType().getNodeManager()->mkConstInt(Rational(d_int));
  }

  IntegerEnumerator& operator++() override
  {
    // sequence is 0, 1, -1, 2, -2, 3, -3, ...
    if(d_int <= 0) {
      d_int = -d_int + 1;
    } else {
      d_int = -d_int;
    }
    return *this;
  }

  bool isFinished() override { return false; }
};/* class IntegerEnumerator */

}  // namespace arith
}  // namespace theory
}  // namespace cvc5::internal

#endif /* CVC5__THEORY__ARITH__TYPE_ENUMERATOR_H */
