[CP-SAT] change element and table protos, API to support affine expre…

…ssions

ortools/flatzinc/cp_model_fz_solver.cc

@@ -424,7 +424,11 @@ void CpModelProtoWithMapping::BuildTableFromDomainIntLinEq(
   const std::vector<int> vars = LookupVars(fz_ct.arguments[1]);
   const int rhs = fz_ct.arguments[2].Value();
   CHECK_EQ(coeffs.back(), -1);
-  for (const int var : vars) ct->mutable_table()->add_vars(var);
+  for (const int var : vars) {
+    LinearExpressionProto* expr = ct->mutable_table()->add_exprs();
+    expr->add_vars(var);
+    expr->add_coeffs(1);
+  }
 
   switch (vars.size()) {
     case 3: {
@@ -732,7 +736,15 @@ void CpModelProtoWithMapping::FillConstraint(const fz::Constraint& fz_ct,
     }
   } else if (fz_ct.type == "ortools_table_int") {
     auto* arg = ct->mutable_table();
-    for (const int var : LookupVars(fz_ct.arguments[0])) arg->add_vars(var);
+    for (const VarOrValue v : LookupVarsOrValues(fz_ct.arguments[0])) {
+      LinearExpressionProto* expr = arg->add_exprs();
+      if (v.var != kNoVar) {
+        expr->add_vars(v.var);
+        expr->add_coeffs(1);
+      } else {
+        expr->set_offset(v.value);
+      }
+    }
     for (const int64_t value : fz_ct.arguments[1].values)
       arg->add_values(value);
   } else if (fz_ct.type == "ortools_regular") {

ortools/flatzinc/fz.cc

@@ -32,6 +32,7 @@
 #include "absl/strings/match.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/time/time.h"
 #include "ortools/base/commandlineflags.h"
 #include "ortools/base/logging.h"
 #include "ortools/base/path.h"
@@ -127,7 +128,7 @@ std::vector<char*> FixAndParseParameters(int* argc, char*** argv) {
 }
 
 Model ParseFlatzincModel(const std::string& input, bool input_is_filename,
-                         SolverLogger* logger) {
+                         SolverLogger* logger, absl::Duration* parse_duration) {
   WallTimer timer;
   timer.Start();
 
@@ -147,8 +148,9 @@ Model ParseFlatzincModel(const std::string& input, bool input_is_filename,
     CHECK(ParseFlatzincString(input, &model));
   }
 
+  *parse_duration = timer.GetDuration();
   SOLVER_LOG(logger, "File ", (input_is_filename ? input : "stdin"),
-             " parsed in ", timer.GetInMs(), " ms");
+             " parsed in ", absl::ToInt64Milliseconds(*parse_duration), " ms");
   SOLVER_LOG(logger, "");
 
   // Print statistics.
@@ -205,9 +207,11 @@ int main(int argc, char** argv) {
     logger.SetLogToStdOut(true);
   }
 
+  absl::Duration parse_duration;
   operations_research::fz::Model model =
       operations_research::fz::ParseFlatzincModel(
-          input, !absl::GetFlag(FLAGS_read_from_stdin), &logger);
+          input, !absl::GetFlag(FLAGS_read_from_stdin), &logger,
+          &parse_duration);
   operations_research::sat::ProcessFloatingPointOVariablesAndObjective(&model);
 
   operations_research::fz::FlatzincSatParameters parameters;
@@ -219,7 +223,8 @@ int main(int argc, char** argv) {
   parameters.random_seed = absl::GetFlag(FLAGS_fz_seed);
   parameters.display_statistics = absl::GetFlag(FLAGS_statistics);
   parameters.number_of_threads = absl::GetFlag(FLAGS_threads);
-  parameters.max_time_in_seconds = absl::GetFlag(FLAGS_time_limit);
+  parameters.max_time_in_seconds =
+      absl::GetFlag(FLAGS_time_limit) - absl::ToInt64Seconds(parse_duration);
   parameters.ortools_mode = absl::GetFlag(FLAGS_ortools_mode);
 
   operations_research::SolverLogger solution_logger;

ortools/java/com/google/ortools/sat/CpModel.java

@@ -448,59 +448,59 @@ public MultipleCircuitConstraint addMultipleCircuit() {
   }
 
   /**
-   * Adds {@code AllowedAssignments(variables)}.
+   * Adds {@code AllowedAssignments(expressions)}.
    *
-   * <p>An AllowedAssignments constraint is a constraint on an array of variables that forces, when
-   * all variables are fixed to a single value, that the corresponding list of values is equal to
-   * one of the tuples of the tupleList.
+   * <p>An AllowedAssignments constraint is a constraint on an array of affine expressions that
+   * forces, when all expressions are fixed to a single value, that the corresponding list of values
+   * is equal to one of the tuples of the tupleList.
    *
-   * @param variables a list of variables
+   * @param expressions a list of affine expressions (a * var + b)
    * @return an instance of the TableConstraint class without any tuples. Tuples can be added
    *     directly to the table constraint.
    */
-  public TableConstraint addAllowedAssignments(IntVar[] variables) {
-    return addAllowedAssignments(Arrays.asList(variables));
+  public TableConstraint addAllowedAssignments(LinearArgument[] expressions) {
+    return addAllowedAssignments(Arrays.asList(expressions));
   }
 
   /**
-   * Adds {@code AllowedAssignments(variables)}.
+   * Adds {@code AllowedAssignments(expressions)}.
    *
-   * @see addAllowedAssignments(IntVar[])
+   * @see addAllowedAssignments(LinearArgument[] expressions)
    */
-  public TableConstraint addAllowedAssignments(Iterable<IntVar> variables) {
+  public TableConstraint addAllowedAssignments(Iterable<? extends LinearArgument> expressions) {
     TableConstraint ct = new TableConstraint(modelBuilder);
     TableConstraintProto.Builder table = ct.getBuilder().getTableBuilder();
-    for (IntVar var : variables) {
-      table.addVars(var.getIndex());
+    for (LinearArgument expr : expressions) {
+      table.addExprs(getLinearExpressionProtoBuilderFromLinearArgument(expr, /* negate= */ false));
     }
     table.setNegated(false);
     return ct;
   }
 
   /**
-   * Adds {@code ForbiddenAssignments(variables)}.
+   * Adds {@code ForbiddenAssignments(expressions)}.
    *
-   * <p>A ForbiddenAssignments constraint is a constraint on an array of variables where the list of
-   * impossible combinations is provided in the tuples list.
+   * <p>A ForbiddenAssignments constraint is a constraint on an array of affine expressions where
+   * the list of impossible combinations is provided in the tuples list.
    *
-   * @param variables a list of variables
+   * @param expressions a list of affine expressions (a * var + b)
    * @return an instance of the TableConstraint class without any tuples. Tuples can be added
    *     directly to the table constraint.
    */
-  public TableConstraint addForbiddenAssignments(IntVar[] variables) {
-    return addForbiddenAssignments(Arrays.asList(variables));
+  public TableConstraint addForbiddenAssignments(LinearArgument[] expressions) {
+    return addForbiddenAssignments(Arrays.asList(expressions));
   }
 
   /**
-   * Adds {@code ForbiddenAssignments(variables)}.
+   * Adds {@code ForbiddenAssignments(expressions)}.
    *
-   * @see addForbiddenAssignments(IntVar[])
+   * @see addForbiddenAssignments(LinearArgument[] expressions)
    */
-  public TableConstraint addForbiddenAssignments(Iterable<IntVar> variables) {
+  public TableConstraint addForbiddenAssignments(Iterable<? extends LinearArgument> expressions) {
     TableConstraint ct = new TableConstraint(modelBuilder);
     TableConstraintProto.Builder table = ct.getBuilder().getTableBuilder();
-    for (IntVar var : variables) {
-      table.addVars(var.getIndex());
+    for (LinearArgument expr : expressions) {
+      table.addExprs(getLinearExpressionProtoBuilderFromLinearArgument(expr, /* negate= */ false));
     }
     table.setNegated(true);
     return ct;

ortools/java/com/google/ortools/sat/TableConstraint.java

@@ -32,13 +32,13 @@ public TableConstraint(CpModelProto.Builder builder) {
    *
    * @param tuple the tuple to add to the constraint.
    * @throws CpModel.WrongLength if the tuple does not have the same length as the array of
-   *     variables of the constraint.
+   *     expressions of the constraint.
    */
   public TableConstraint addTuple(int[] tuple) {
     TableConstraintProto.Builder table = getBuilder().getTableBuilder();
-    if (tuple.length != table.getVarsCount()) {
+    if (tuple.length != table.getExprsCount()) {
       throw new CpModel.WrongLength(
-          "addTuple", "tuple does not have the same length as the variables");
+          "addTuple", "tuple does not have the same length as the expressions");
     }
     for (int value : tuple) {
       table.addValues(value);
@@ -51,13 +51,13 @@ public TableConstraint addTuple(int[] tuple) {
    *
    * @param tuple the tuple to add to the constraint.
    * @throws CpModel.WrongLength if the tuple does not have the same length as the array of
-   *     variables of the constraint.
+   *     expressions of the constraint.
    */
   public TableConstraint addTuple(long[] tuple) {
     TableConstraintProto.Builder table = getBuilder().getTableBuilder();
-    if (tuple.length != table.getVarsCount()) {
+    if (tuple.length != table.getExprsCount()) {
       throw new CpModel.WrongLength(
-          "addTuple", "tuple does not have the same length as the variables");
+          "addTuple", "tuple does not have the same length as the expressions");
     }
     for (long value : tuple) {
       table.addValues(value);
@@ -75,9 +75,9 @@ public TableConstraint addTuple(long[] tuple) {
   public TableConstraint addTuples(int[][] tuples) {
     TableConstraintProto.Builder table = getBuilder().getTableBuilder();
     for (int[] tuple : tuples) {
-      if (tuple.length != table.getVarsCount()) {
+      if (tuple.length != table.getExprsCount()) {
         throw new CpModel.WrongLength(
-            "addTuples", "a tuple does not have the same length as the variables");
+            "addTuples", "a tuple does not have the same length as the expressions");
       }
       for (int value : tuple) {
         table.addValues(value);
@@ -91,12 +91,12 @@ public TableConstraint addTuples(int[][] tuples) {
    *
    * @param tuples the list of tuples to add to the constraint.
    * @throws CpModel.WrongLength if one tuple does not have the same length as the array of
-   *     variables of the constraint.
+   *     expressions of the constraint.
    */
   public TableConstraint addTuples(long[][] tuples) {
     TableConstraintProto.Builder table = getBuilder().getTableBuilder();
     for (long[] tuple : tuples) {
-      if (tuple.length != table.getVarsCount()) {
+      if (tuple.length != table.getExprsCount()) {
         throw new CpModel.WrongLength(
             "addTuples", "a tuple does not have the same length as the variables");
       }

ortools/sat/BUILD.bazel

@@ -743,7 +743,7 @@ cc_library(
         ":util",
         "//ortools/algorithms:sparse_permutation",
         "//ortools/base",
-        "//ortools/base:mathutil",
+        "//ortools/base:stl_util",
         "//ortools/port:proto_utils",
         "//ortools/util:affine_relation",
         "//ortools/util:bitset",

ortools/sat/cp_model.cc

@@ -524,7 +524,7 @@ void MultipleCircuitConstraint::AddArc(int tail, int head, BoolVar literal) {
 }
 
 void TableConstraint::AddTuple(absl::Span<const int64_t> tuple) {
-  CHECK_EQ(tuple.size(), proto_->table().vars_size());
+  CHECK_EQ(tuple.size(), proto_->table().exprs_size());
   for (const int64_t t : tuple) {
     proto_->mutable_table()->add_values(t);
   }
@@ -986,24 +986,55 @@ MultipleCircuitConstraint CpModelBuilder::AddMultipleCircuitConstraint() {
 }
 
 TableConstraint CpModelBuilder::AddAllowedAssignments(
-    absl::Span<const IntVar> vars) {
+    absl::Span<const LinearExpr> expressions) {
   ConstraintProto* const proto = cp_model_.add_constraints();
-  for (const IntVar& var : vars) {
-    proto->mutable_table()->add_vars(var.index_);
+  for (const LinearExpr& expr : expressions) {
+    *proto->mutable_table()->add_exprs() = LinearExprToProto(expr);
   }
   return TableConstraint(proto);
 }
 
-TableConstraint CpModelBuilder::AddForbiddenAssignments(
-    absl::Span<const IntVar> vars) {
+TableConstraint CpModelBuilder::AddAllowedAssignments(
+    absl::Span<const IntVar> variables) {
   ConstraintProto* const proto = cp_model_.add_constraints();
-  for (const IntVar& var : vars) {
-    proto->mutable_table()->add_vars(var.index_);
+  for (const IntVar var : variables) {
+    LinearExpressionProto* expr = proto->mutable_table()->add_exprs();
+    expr->add_vars(var.index_);
+    expr->add_coeffs(1);
   }
-  proto->mutable_table()->set_negated(true);
   return TableConstraint(proto);
 }
 
+TableConstraint CpModelBuilder::AddAllowedAssignments(
+    std::initializer_list<LinearExpr> expressions) {
+  ConstraintProto* const proto = cp_model_.add_constraints();
+  for (const LinearExpr& expr : expressions) {
+    *proto->mutable_table()->add_exprs() = LinearExprToProto(expr);
+  }
+  return TableConstraint(proto);
+}
+
+TableConstraint CpModelBuilder::AddForbiddenAssignments(
+    absl::Span<const LinearExpr> expressions) {
+  TableConstraint ct = AddAllowedAssignments(expressions);
+  ct.MutableProto()->mutable_table()->set_negated(true);
+  return ct;
+}
+
+TableConstraint CpModelBuilder::AddForbiddenAssignments(
+    absl::Span<const IntVar> variables) {
+  TableConstraint ct = AddAllowedAssignments(variables);
+  ct.MutableProto()->mutable_table()->set_negated(true);
+  return ct;
+}
+
+TableConstraint CpModelBuilder::AddForbiddenAssignments(
+    std::initializer_list<LinearExpr> expressions) {
+  TableConstraint ct = AddAllowedAssignments(expressions);
+  ct.MutableProto()->mutable_table()->set_negated(true);
+  return ct;
+}
+
 Constraint CpModelBuilder::AddInverseConstraint(
     absl::Span<const IntVar> variables,
     absl::Span<const IntVar> inverse_variables) {

ortools/sat/cp_model.h

@@ -899,27 +899,51 @@ class CpModelBuilder {
   /**
    * Adds an allowed assignments constraint.
    *
-   * An AllowedAssignments constraint is a constraint on an array of variables
-   * that forces, when all variables are fixed to a single value, that the
-   * corresponding list of values is equal to one of the tuples added to the
-   * constraint.
+   * An AllowedAssignments constraint is a constraint on an array of affine
+   * expressions (a * var + b) that forces, when all expressions are fixed to a
+   * single value, that the corresponding list of values is equal to one of the
+   * tuples added to the constraint.
    *
    * It returns a table constraint that allows adding tuples incrementally after
    * construction.
    */
-  TableConstraint AddAllowedAssignments(absl::Span<const IntVar> vars);
+  TableConstraint AddAllowedAssignments(
+      absl::Span<const LinearExpr> expressions);
+
+  /**
+   * Adds an allowed assignments constraint.
+   */
+  TableConstraint AddAllowedAssignments(absl::Span<const IntVar> variables);
+
+  /**
+   * Adds an allowed assignments constraint.
+   */
+  TableConstraint AddAllowedAssignments(
+      std::initializer_list<LinearExpr> expressions);
 
   /**
    * Adds an forbidden assignments constraint.
    *
-   * A ForbiddenAssignments constraint is a constraint on an array of variables
-   * where the list of impossible combinations is provided in the tuples added
-   * to the constraint.
+   * A ForbiddenAssignments constraint is a constraint on an array of affine
+   * expressions (a * var + b) where the list of impossible combinations is
+   * provided in the tuples added to the constraint.
    *
    * It returns a table constraint that allows adding tuples incrementally after
    * construction.
    */
-  TableConstraint AddForbiddenAssignments(absl::Span<const IntVar> vars);
+  TableConstraint AddForbiddenAssignments(
+      absl::Span<const LinearExpr> expression);
+
+  /**
+   * Adds an forbidden assignments constraint.
+   */
+  TableConstraint AddForbiddenAssignments(absl::Span<const IntVar> variables);
+
+  /**
+   * Adds an forbidden assignments constraint.
+   */
+  TableConstraint AddForbiddenAssignments(
+      std::initializer_list<LinearExpr> expressions);
 
   /** An inverse constraint.
    *

ortools/sat/cp_model.proto

@@ -247,12 +247,14 @@ message RoutesConstraintProto {
   int64 capacity = 5;
 }
 
-// The values of the n-tuple formed by the given variables can only be one of
+// The values of the n-tuple formed by the given expression can only be one of
 // the listed n-tuples in values. The n-tuples are encoded in a flattened way:
 //     [tuple0_v0, tuple0_v1, ..., tuple0_v{n-1}, tuple1_v0, ...].
+// Expressions must be affine (a * var + b).
 message TableConstraintProto {
-  repeated int32 vars = 1;
+  repeated int32 vars = 1;  // Legacy field.
   repeated int64 values = 2;
+  repeated LinearExpressionProto exprs = 4;
 
   // If true, the meaning is "negated", that is we forbid any of the given
   // tuple from a feasible assignment.