From eed337ade5f1caac6e1238b682e911e5b0ed9ae9 Mon Sep 17 00:00:00 2001
From: Laurent Perron <lperron@google.com>
Date: Fri, 18 Oct 2024 19:08:59 +0200
Subject: [PATCH] [CP-SAT] add support for element with affine expressions in
 C# + test

---
 ortools/sat/csharp/CpModel.cs        | 41 ++++++++++++++++------------
 ortools/sat/csharp/SatSolverTests.cs | 22 +++++++++++++++
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/ortools/sat/csharp/CpModel.cs b/ortools/sat/csharp/CpModel.cs
index d7c7d9f126..f7f92dc10d 100644
--- a/ortools/sat/csharp/CpModel.cs
+++ b/ortools/sat/csharp/CpModel.cs
@@ -233,19 +233,19 @@ public Constraint AddAllDifferent(IEnumerable<LinearExpr> exprs)
 
     /**
      * <summary>
-     * Adds the element constraint: <c>variables[index] == target</c>.
+     * Adds the element constraint: <c>exprs[index] == target</c>.
      * </summary>
      */
-    public Constraint AddElement(IntVar index, IEnumerable<IntVar> vars, IntVar target)
+    public Constraint AddElement(LinearExpr index, IEnumerable<LinearExpr> exprs, LinearExpr target)
     {
         ElementConstraintProto element = new ElementConstraintProto();
-        element.Index = index.Index;
-        element.Vars.TrySetCapacity(vars);
-        foreach (IntVar var in vars)
+        element.LinearIndex = GetLinearExpressionProto(index);
+        element.Exprs.TrySetCapacity(exprs);
+        foreach (LinearExpr expr in exprs)
         {
-            element.Vars.Add(var.Index);
+            element.Exprs.Add(GetLinearExpressionProto(expr));
         }
-        element.Target = target.Index;
+        element.LinearTarget = GetLinearExpressionProto(target);
 
         Constraint ct = new Constraint(model_);
         ct.Proto.Element = element;
@@ -257,16 +257,16 @@ public Constraint AddElement(IntVar index, IEnumerable<IntVar> vars, IntVar targ
      * Adds the element constraint: <c> values[index] == target</c>.
      * </summary>
      */
-    public Constraint AddElement(IntVar index, IEnumerable<long> values, IntVar target)
+    public Constraint AddElement(LinearExpr index, IEnumerable<long> values, LinearExpr target)
     {
         ElementConstraintProto element = new ElementConstraintProto();
-        element.Index = index.Index;
-        element.Vars.TrySetCapacity(values);
+        element.LinearIndex = GetLinearExpressionProto(index);
+        element.Exprs.TrySetCapacity(values);
         foreach (long value in values)
         {
-            element.Vars.Add(ConvertConstant(value));
+            element.Exprs.Add(GetLinearExpressionProto(value));
         }
-        element.Target = target.Index;
+        element.LinearTarget = GetLinearExpressionProto(target);
 
         Constraint ct = new Constraint(model_);
         ct.Proto.Element = element;
@@ -278,16 +278,16 @@ public Constraint AddElement(IntVar index, IEnumerable<long> values, IntVar targ
      * Adds the element constraint: <c> values[index] == target</c>.
      * </summary>
      */
-    public Constraint AddElement(IntVar index, IEnumerable<int> values, IntVar target)
+    public Constraint AddElement(LinearExpr index, IEnumerable<int> values, LinearExpr target)
     {
         ElementConstraintProto element = new ElementConstraintProto();
-        element.Index = index.Index;
-        element.Vars.TrySetCapacity(values);
+        element.LinearIndex = GetLinearExpressionProto(index);
+        element.Exprs.TrySetCapacity(values);
         foreach (int value in values)
         {
-            element.Vars.Add(ConvertConstant(value));
+            element.Exprs.Add(GetLinearExpressionProto(value));
         }
-        element.Target = target.Index;
+        element.LinearTarget = GetLinearExpressionProto(target);
 
         Constraint ct = new Constraint(model_);
         ct.Proto.Element = element;
@@ -1258,6 +1258,13 @@ internal LinearExpressionProto GetLinearExpressionProto(LinearExpr expr, bool ne
         return linear;
     }
 
+    internal LinearExpressionProto GetLinearExpressionProto(long value)
+    {
+        LinearExpressionProto linear = new LinearExpressionProto();
+        linear.Offset = value;
+        return linear;
+    }
+
     private CpModelProto model_;
     private Dictionary<long, int> constant_map_;
     private Dictionary<int, long> var_value_map_;
diff --git a/ortools/sat/csharp/SatSolverTests.cs b/ortools/sat/csharp/SatSolverTests.cs
index 7fa4d18666..d1bcf2d227 100644
--- a/ortools/sat/csharp/SatSolverTests.cs
+++ b/ortools/sat/csharp/SatSolverTests.cs
@@ -392,6 +392,28 @@ public void Modulo()
         // Console.WriteLine("response = " + response.ToString());
     }
 
+    [Fact]
+    public void ValueElement() 
+    {
+        CpModel model = new CpModel();
+        IntVar v1 = model.NewIntVar(1, 10, "v1");
+        IntVar v2 = model.NewIntVar(1, 10, "v2");
+        model.AddElement(v1 + 2, new int[] {1, 3, 5}, 5 - v2);
+        Assert.Equal(3, model.Model.Constraints[0].Element.Exprs.Count);
+    }
+
+    public void ExprElement() 
+    {
+        CpModel model = new CpModel();
+        IntVar v1 = model.NewIntVar(1, 10, "v1");
+        IntVar v2 = model.NewIntVar(1, 10, "v2");
+        IntVar x = model.NewIntVar(0, 5, "x");
+        IntVar y = model.NewIntVar(0, 5, "y");
+        IntVar z = model.NewIntVar(0, 5, "z");
+        model.AddElement(v1, new LinearExpr[] {x + 2, -y, LinearExpr.Constant(5), 2 * z}, 5 - v2);
+        Assert.Equal(4, model.Model.Constraints[0].Element.Exprs.Count);
+    }    
+
     [Fact]
     public void LargeWeightedSumLong()
     {