[CP-SAT] supports affine expression in automaton

ortools/flatzinc/cp_model_fz_solver.cc

@@ -749,7 +749,15 @@ void CpModelProtoWithMapping::FillConstraint(const fz::Constraint& fz_ct,
       arg->add_values(value);
   } else if (fz_ct.type == "ortools_regular") {
     auto* arg = ct->mutable_automaton();
-    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);
+      }
+    }
 
     int count = 0;
     const int num_states = fz_ct.arguments[1].Value();

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

@@ -509,38 +509,60 @@ public TableConstraint addForbiddenAssignments(Iterable<? extends LinearArgument
   /**
    * Adds an automaton constraint.
    *
-   * <p>An automaton constraint takes a list of variables (of size n), an initial state, a set of
-   * final states, and a set of transitions that will be added incrementally directly on the
+   * <p>An automaton constraint takes a list of affine expressions (of size n), an initial state, a
+   * set of final states, and a set of transitions that will be added incrementally directly on the
    * returned AutomatonConstraint instance. A transition is a triplet ('tail', 'transition',
    * 'head'), where 'tail' and 'head' are states, and 'transition' is the label of an arc from
-   * 'head' to 'tail', corresponding to the value of one variable in the list of variables.
+   * 'head' to 'tail', corresponding to the value of one expression in the list of expressions.
    *
    * <p>This automaton will be unrolled into a flow with n + 1 phases. Each phase contains the
    * possible states of the automaton. The first state contains the initial state. The last phase
    * contains the final states.
    *
    * <p>Between two consecutive phases i and i + 1, the automaton creates a set of arcs. For each
    * transition (tail, label, head), it will add an arc from the state 'tail' of phase i and the
-   * state 'head' of phase i + 1. This arc labeled by the value 'label' of the variables
-   * 'variables[i]'. That is, this arc can only be selected if 'variables[i]' is assigned the value
-   * 'label'.
+   * state 'head' of phase i + 1. This arc labeled by the value 'label' of the expression
+   * 'expressions[i]'. That is, this arc can only be selected if 'expressions[i]' is assigned the
+   * value 'label'.
    *
-   * <p>A feasible solution of this constraint is an assignment of variables such that, starting
-   * from the initial state in phase 0, there is a path labeled by the values of the variables that
-   * ends in one of the final states in the final phase.
+   * <p>A feasible solution of this constraint is an assignment of expressions such that, starting
+   * from the initial state in phase 0, there is a path labeled by the values of the expressions
+   * that ends in one of the final states in the final phase.
    *
-   * @param transitionVariables a non empty list of variables whose values correspond to the labels
-   *     of the arcs traversed by the automaton
+   * @param transitionExpressions a non empty list of affine expressions (a * var + b) whose values
+   *     correspond to the labels of the arcs traversed by the automaton
    * @param startingState the initial state of the automaton
    * @param finalStates a non empty list of admissible final states
    * @return an instance of the Constraint class
    */
   public AutomatonConstraint addAutomaton(
-      IntVar[] transitionVariables, long startingState, long[] finalStates) {
+      LinearArgument[] transitionExpressions, long startingState, long[] finalStates) {
     AutomatonConstraint ct = new AutomatonConstraint(modelBuilder);
     AutomatonConstraintProto.Builder automaton = ct.getBuilder().getAutomatonBuilder();
-    for (IntVar var : transitionVariables) {
-      automaton.addVars(var.getIndex());
+    for (LinearArgument expr : transitionExpressions) {
+      automaton.addExprs(
+          getLinearExpressionProtoBuilderFromLinearArgument(expr, /* negate= */ false));
+    }
+    automaton.setStartingState(startingState);
+    for (long c : finalStates) {
+      automaton.addFinalStates(c);
+    }
+    return ct;
+  }
+
+  /**
+   * Adds an automaton constraint.
+   *
+   * @see addAutomaton(LinearArgument[] transitionExpressions, long startingState, long[]
+   *     finalStates)
+   */
+  public AutomatonConstraint addAutomaton(Iterable<? extends LinearArgument> transitionExpressions,
+      long startingState, long[] finalStates) {
+    AutomatonConstraint ct = new AutomatonConstraint(modelBuilder);
+    AutomatonConstraintProto.Builder automaton = ct.getBuilder().getAutomatonBuilder();
+    for (LinearArgument expr : transitionExpressions) {
+      automaton.addExprs(
+          getLinearExpressionProtoBuilderFromLinearArgument(expr, /* negate= */ false));
     }
     automaton.setStartingState(startingState);
     for (long c : finalStates) {

ortools/sat/BUILD.bazel

@@ -757,6 +757,7 @@ cc_library(
     hdrs = ["presolve_context.h"],
     deps = [
         ":cp_model_cc_proto",
+        ":cp_model_checker",
         ":cp_model_loader",
         ":cp_model_mapping",
         ":cp_model_utils",
@@ -820,6 +821,7 @@ cc_library(
         ":presolve_context",
         "//ortools/base:stl_util",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/types:span",
@@ -980,6 +982,7 @@ cc_library(
         "//ortools/port:proto_utils",
         "//ortools/util:logging",
         "//ortools/util:sorted_interval_list",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",

ortools/sat/cp_model.cc

@@ -1056,12 +1056,42 @@ ReservoirConstraint CpModelBuilder::AddReservoirConstraint(int64_t min_level,
   return ReservoirConstraint(proto, this);
 }
 
+AutomatonConstraint CpModelBuilder::AddAutomaton(
+    absl::Span<const LinearExpr> transition_expressions, int starting_state,
+    absl::Span<const int> final_states) {
+  ConstraintProto* const proto = cp_model_.add_constraints();
+  for (const LinearExpr& expr : transition_expressions) {
+    *proto->mutable_automaton()->add_exprs() = LinearExprToProto(expr);
+  }
+  proto->mutable_automaton()->set_starting_state(starting_state);
+  for (const int final_state : final_states) {
+    proto->mutable_automaton()->add_final_states(final_state);
+  }
+  return AutomatonConstraint(proto);
+}
+
 AutomatonConstraint CpModelBuilder::AddAutomaton(
     absl::Span<const IntVar> transition_variables, int starting_state,
     absl::Span<const int> final_states) {
   ConstraintProto* const proto = cp_model_.add_constraints();
   for (const IntVar& var : transition_variables) {
-    proto->mutable_automaton()->add_vars(var.index_);
+    LinearExpressionProto* expr = proto->mutable_automaton()->add_exprs();
+    expr->add_vars(var.index_);
+    expr->add_coeffs(1);
+  }
+  proto->mutable_automaton()->set_starting_state(starting_state);
+  for (const int final_state : final_states) {
+    proto->mutable_automaton()->add_final_states(final_state);
+  }
+  return AutomatonConstraint(proto);
+}
+
+AutomatonConstraint CpModelBuilder::AddAutomaton(
+    std::initializer_list<LinearExpr> transition_expressions,
+    int starting_state, absl::Span<const int> final_states) {
+  ConstraintProto* const proto = cp_model_.add_constraints();
+  for (const LinearExpr& expr : transition_expressions) {
+    *proto->mutable_automaton()->add_exprs() = LinearExprToProto(expr);
   }
   proto->mutable_automaton()->set_starting_state(starting_state);
   for (const int final_state : final_states) {

ortools/sat/cp_model.h

@@ -1001,10 +1001,24 @@ class CpModelBuilder {
    * It returns an AutomatonConstraint that allows adding transition
    * incrementally after construction.
    */
+  AutomatonConstraint AddAutomaton(
+      absl::Span<const LinearExpr> transition_expressions, int starting_state,
+      absl::Span<const int> final_states);
+
+  /**
+   * An automaton constraint.
+   */
   AutomatonConstraint AddAutomaton(
       absl::Span<const IntVar> transition_variables, int starting_state,
       absl::Span<const int> final_states);
 
+  /**
+   * An automaton constraint.
+   */
+  AutomatonConstraint AddAutomaton(
+      std::initializer_list<LinearExpr> transition_expressions,
+      int starting_state, absl::Span<const int> final_states);
+
   /// Adds target == min(vars).
   Constraint AddMinEquality(const LinearExpr& target,
                             absl::Span<const IntVar> vars);

ortools/sat/cp_model.proto

@@ -268,7 +268,7 @@ message InverseConstraintProto {
   repeated int32 f_inverse = 2;
 }
 
-// This constraint forces a sequence of variables to be accepted by an
+// This constraint forces a sequence of expressions to be accepted by an
 // automaton.
 message AutomatonConstraintProto {
   // A state is identified by a non-negative number. It is preferable to keep
@@ -284,9 +284,12 @@ message AutomatonConstraintProto {
   repeated int64 transition_head = 5;
   repeated int64 transition_label = 6;
 
-  // The sequence of variables. The automaton is ran for vars_size() "steps" and
-  // the value of vars[i] corresponds to the transition label at step i.
+  // Legacy field.
   repeated int32 vars = 7;
+  // The sequence of affine expressions (a * var + b). The automaton is ran for
+  // exprs_size() "steps" and the value of exprs[i] corresponds to the
+  // transition label at step i.
+  repeated LinearExpressionProto exprs = 8;
 }
 
 // A list of variables, without any semantics.

ortools/sat/cp_model_checker.cc

@@ -498,18 +498,24 @@ std::string ValidateTableConstraint(const ConstraintProto& ct) {
 }
 
 std::string ValidateAutomatonConstraint(const ConstraintProto& ct) {
-  const int num_transistions = ct.automaton().transition_tail().size();
-  if (num_transistions != ct.automaton().transition_head().size() ||
-      num_transistions != ct.automaton().transition_label().size()) {
+  const AutomatonConstraintProto& automaton = ct.automaton();
+  if (!automaton.vars().empty() && !automaton.exprs().empty()) {
+    return absl::StrCat(
+        "Inconsistent automaton with both legacy and new format defined: ",
+        ProtobufShortDebugString(ct));
+  }
+  const int num_transistions = automaton.transition_tail().size();
+  if (num_transistions != automaton.transition_head().size() ||
+      num_transistions != automaton.transition_label().size()) {
     return absl::StrCat(
         "The transitions repeated fields must have the same size: ",
         ProtobufShortDebugString(ct));
   }
   absl::flat_hash_map<std::pair<int64_t, int64_t>, int64_t> tail_label_to_head;
   for (int i = 0; i < num_transistions; ++i) {
-    const int64_t tail = ct.automaton().transition_tail(i);
-    const int64_t head = ct.automaton().transition_head(i);
-    const int64_t label = ct.automaton().transition_label(i);
+    const int64_t tail = automaton.transition_tail(i);
+    const int64_t head = automaton.transition_head(i);
+    const int64_t label = automaton.transition_label(i);
     if (label <= std::numeric_limits<int64_t>::min() + 1 ||
         label == std::numeric_limits<int64_t>::max()) {
       return absl::StrCat("labels in the automaton constraint are too big: ",
@@ -1489,28 +1495,32 @@ class ConstraintChecker {
 
   bool AutomatonConstraintIsFeasible(const ConstraintProto& ct) {
     // Build the transition table {tail, label} -> head.
+    const AutomatonConstraintProto& automaton = ct.automaton();
     absl::flat_hash_map<std::pair<int64_t, int64_t>, int64_t> transition_map;
-    const int num_transitions = ct.automaton().transition_tail().size();
+    const int num_transitions = automaton.transition_tail().size();
     for (int i = 0; i < num_transitions; ++i) {
-      transition_map[{ct.automaton().transition_tail(i),
-                      ct.automaton().transition_label(i)}] =
-          ct.automaton().transition_head(i);
+      transition_map[{automaton.transition_tail(i),
+                      automaton.transition_label(i)}] =
+          automaton.transition_head(i);
     }
 
     // Walk the automaton.
-    int64_t current_state = ct.automaton().starting_state();
-    const int num_steps = ct.automaton().vars_size();
+    int64_t current_state = automaton.starting_state();
+    const int num_steps =
+        std::max(automaton.vars_size(), automaton.exprs_size());
     for (int i = 0; i < num_steps; ++i) {
-      const std::pair<int64_t, int64_t> key = {current_state,
-                                               Value(ct.automaton().vars(i))};
+      const std::pair<int64_t, int64_t> key = {
+          current_state, automaton.vars().empty()
+                             ? LinearExpressionValue(automaton.exprs(i))
+                             : Value(automaton.vars(i))};
       if (!transition_map.contains(key)) {
         return false;
       }
       current_state = transition_map[key];
     }
 
     // Check we are now in a final state.
-    for (const int64_t final : ct.automaton().final_states()) {
+    for (const int64_t final : automaton.final_states()) {
       if (current_state == final) return true;
     }
     return false;