Add inequality constraints

7 files changed, 74 insertions, 40 deletions

diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/DimensionRenamer.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/DimensionRenamer.java
index d2d7367585c..fc54dcac39a 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/DimensionRenamer.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/DimensionRenamer.java
@@ -18,7 +18,7 @@ import java.util.logging.Logger;
 import java.util.stream.Collectors;
 
 /**
- * A constraint satisfier to find suitable dimension names to reduce the
+ * A constraint solver which finds suitable dimension names to reduce the
  * amount of necessary renaming during evaluation of an imported model.
  *
  * @author lesters
@@ -29,7 +29,7 @@ public class DimensionRenamer {
 
     private final String dimensionPrefix;
     private final ListMap<String, Integer> variables = new ListMap<>();
-    private final Map<Arc, Constraint> constraints = new HashMap<>();
+    private final ListMap<Arc, Constraint> constraints = new ListMap<>();
 
     /** The solution to this, or null if no solution is found (yet) */
     private Map<String, Integer> renames = null;
@@ -91,14 +91,22 @@ public class DimensionRenamer {
         boolean solved = trySolve(variables, constraints, maxIterations, renames);
         if ( ! solved) {
             renames.clear();
-            Map<Arc, Constraint> hardConstraints = new HashMap<>();
-            constraints.entrySet().stream().filter(e -> !e.getValue().isSoft())
-                       .forEach(e -> hardConstraints.put(e.getKey(), e.getValue()));
-            if (hardConstraints.size() < constraints.size())
+            ListMap<Arc, Constraint> hardConstraints = new ListMap<>();
+            boolean relaxed = false;
+            for (var entry : constraints.entrySet()) {
+                Arc arc = entry.getKey();
+                for (Constraint constraint : entry.getValue()) {
+                    if ( ! constraint.isSoft())
+                        hardConstraints.put(arc, constraint);
+                    else
+                        relaxed = true;
+                }
+            }
+            if (relaxed)
                 solved = trySolve(variables, hardConstraints, maxIterations, renames);
             if ( ! solved) {
                 throw new IllegalArgumentException("Could not find a dimension naming solution " +
-                                                   " given constraints\n" + constraintsToString(hardConstraints));
+                                                   "given constraints\n" + constraintsToString(hardConstraints));
             }
         }
 
@@ -112,7 +120,7 @@ public class DimensionRenamer {
 
     /** Try the solve the constraint problem given in the arguments, and put the result in renames */
     private static boolean trySolve(ListMap<String, Integer> inputVariables,
-                                    Map<Arc, Constraint> constraints,
+                                    ListMap<Arc, Constraint> constraints,
                                     int maxIterations,
                                     Map<String, Integer> renames) {
         var variables = new ListMap<>(inputVariables);
@@ -123,7 +131,7 @@ public class DimensionRenamer {
             if (values.size() > 1) {
                 if ( ! ac3(iterations, variables, constraints)) return false;
                 values.sort(Integer::compare);
-                variables.put(dimension, values.get(0));
+                variables.replace(dimension, values.get(0));
             }
             renames.put(dimension, variables.get(dimension).get(0));
             if (iterations.get() > maxIterations) return false;
@@ -133,10 +141,8 @@ public class DimensionRenamer {
 
     void solve() {
         log.log(Level.FINE, () -> "Rename problem:\n" + constraintsToString(constraints));
-        System.out.println("Rename problem:\n" + constraintsToString(constraints));
         renames = solve(100000);
         log.log(Level.FINE, () -> "Rename solution:\n" + renamesToString(renames));
-        System.out.println("Rename solution:\n" + renamesToString(renames));
     }
 
     private static String renamesToString(Map<String, Integer> renames) {
@@ -156,7 +162,7 @@ public class DimensionRenamer {
 
     private static boolean ac3(MutableInteger iterations,
                                ListMap<String, Integer> variables,
-                               Map<Arc, Constraint> constraints) {
+                               ListMap<Arc, Constraint> constraints) {
         Deque<Arc> workList = new ArrayDeque<>(constraints.keySet());
         while ( ! workList.isEmpty()) {
             Arc arc = workList.pop();
@@ -177,16 +183,15 @@ public class DimensionRenamer {
 
     private static boolean revise(Arc arc,
                                   ListMap<String, Integer> variables,
-                                  Map<Arc, Constraint> constraints) {
+                                  ListMap<Arc, Constraint> constraints) {
         boolean revised = false;
         for (Iterator<Integer> fromIterator = variables.get(arc.from).iterator(); fromIterator.hasNext(); ) {
             Integer from = fromIterator.next();
             boolean satisfied = false;
             for (Iterator<Integer> toIterator = variables.get(arc.to).iterator(); toIterator.hasNext(); ) {
                 Integer to = toIterator.next();
-                if (constraints.get(arc).test(from, to)) {
+                if (constraints.get(arc).stream().allMatch(constraint -> constraint.test(from, to)))
                     satisfied = true;
-                }
             }
             if ( ! satisfied) {
                 fromIterator.remove();
@@ -196,11 +201,20 @@ public class DimensionRenamer {
         return revised;
     }
 
-    private static String constraintsToString(Map<Arc, Constraint> constraints) {
-        return constraints.entrySet().stream()
-                                     .filter(e -> ! e.getValue().isOpposite())
-                                     .map(e -> "  "  + e.getKey().from + " " + e.getValue() + " " + e.getKey().to + " (origin: " + e.getKey().operation + ")")
-                                     .collect(Collectors.joining("\n"));
+    private static String constraintsToString(ListMap<Arc, Constraint> constraints) {
+        StringBuilder b = new StringBuilder();
+        for (var entry : constraints.entrySet()) {
+            Arc arc = entry.getKey();
+            for (Constraint constraint : entry.getValue()) {
+                if (constraint.isOpposite()) continue; // noise
+                b.append("  ");
+                if (constraint.isSoft())
+                    b.append("(soft) ");
+                b.append(arc.from).append(" ").append(constraint).append(" ").append(arc.to);
+                b.append("  (origin: ").append(arc.operation).append(")\n");
+            }
+        }
+        return b.toString();
     }
 
     private static class Arc {
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/OrderedTensorType.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/OrderedTensorType.java
index 8ea9c9a258d..1cb8f3a2951 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/OrderedTensorType.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/OrderedTensorType.java
@@ -7,8 +7,11 @@ import com.yahoo.tensor.TensorTypeParser;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
@@ -130,15 +133,18 @@ public class OrderedTensorType {
     }
 
     public OrderedTensorType rename(DimensionRenamer renamer) {
-        System.out.println("Renaming " + this);
         List<TensorType.Dimension> renamedDimensions = new ArrayList<>(dimensions.size());
+        Map<String, String> new2Old = new HashMap<>(); // Just to create meaningful error messages
         for (TensorType.Dimension dimension : dimensions) {
             String oldName = dimension.name();
             Optional<String> newName = renamer.dimensionNameOf(oldName);
-            if (!newName.isPresent())
-                return this; // presumably, already renamed
-            if ( ! oldName.equals(newName.get()))
-                System.out.println("  Renaming " + oldName + " to " + newName.get());
+            if ( newName.isEmpty()) return this; // presumably already renamed
+
+            if (new2Old.containsKey(newName.get()))
+                throw new IllegalArgumentException("Can not rename '" + oldName + "' to '" + newName.get() + "' in " + this +
+                                                   " as '" + new2Old.get(newName.get()) + "' should also be renamed to it");
+            new2Old.put(newName.get(), oldName);
+
             TensorType.Dimension.Type dimensionType = dimension.type();
             if (dimensionType == TensorType.Dimension.Type.indexedBound) {
                 renamedDimensions.add(TensorType.Dimension.indexed(newName.get(), dimension.size().get()));
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
index 2d746bf338c..8fe70f7eefb 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Argument.java
@@ -39,15 +39,7 @@ public class Argument extends IntermediateOperation {
 
     @Override
     public void addDimensionNameConstraints(DimensionRenamer renamer) {
-        for (int i = 0; i < type.dimensions().size(); i++) {
-            renamer.addDimension(type.dimensions().get(i).name());
-
-            // Each dimension is distinct:
-            for (int j = i + 1; j < type.dimensions().size(); j++)
-                renamer.addConstraint(type.dimensions().get(i).name(), type.dimensions().get(j).name(),
-                                      DimensionRenamer.Constraint.notEqual(false),
-                                      this);
-        }
+        addConstraintsFrom(type, renamer);
     }
 
     @Override
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
index e7c4b355217..c48a99c2716 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Const.java
@@ -62,9 +62,7 @@ public class Const extends IntermediateOperation {
 
     @Override
     public void addDimensionNameConstraints(DimensionRenamer renamer) {
-        for (TensorType.Dimension dimension : type.type().dimensions()) {
-            renamer.addDimension(dimension.name());
-        }
+        addConstraintsFrom(type, renamer);
     }
 
     @Override
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Constant.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Constant.java
index a1cc83296b0..8c6e69584e0 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Constant.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/Constant.java
@@ -48,9 +48,7 @@ public class Constant extends IntermediateOperation {
 
     @Override
     public void addDimensionNameConstraints(DimensionRenamer renamer) {
-        for (TensorType.Dimension dimension : type.type().dimensions()) {
-            renamer.addDimension(dimension.name());
-        }
+        addConstraintsFrom(type, renamer);
     }
 
     @Override
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/IntermediateOperation.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/IntermediateOperation.java
index 78bed31f5b0..878fa1ca1b1 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/IntermediateOperation.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/operations/IntermediateOperation.java
@@ -99,6 +99,20 @@ public abstract class IntermediateOperation {
     /** Add dimension name constraints for this operation */
     public void addDimensionNameConstraints(DimensionRenamer renamer) { }
 
+    /** Conveinence method to adds dimensions and constraints of the given tensor type */
+    protected void addConstraintsFrom(OrderedTensorType type, DimensionRenamer renamer) {
+        for (int i = 0; i < type.dimensions().size(); i++) {
+            renamer.addDimension(type.dimensions().get(i).name());
+
+            // Each dimension is distinct:
+            for (int j = i + 1; j < type.dimensions().size(); j++) {
+                renamer.addConstraint(type.dimensions().get(i).name(), type.dimensions().get(j).name(),
+                                      DimensionRenamer.Constraint.notEqual(false),
+                                      this);
+            }
+        }
+    }
+
     /** Performs dimension rename for this operation */
     public void renameDimensions(DimensionRenamer renamer) { type = type.rename(renamer); }
 
diff --git a/vespajlib/src/main/java/com/yahoo/collections/ListMap.java b/vespajlib/src/main/java/com/yahoo/collections/ListMap.java
index 052ea55d6fe..479850beb1a 100644
--- a/vespajlib/src/main/java/com/yahoo/collections/ListMap.java
+++ b/vespajlib/src/main/java/com/yahoo/collections/ListMap.java
@@ -60,6 +60,18 @@ public class ListMap<K, V> {
         }
     }
 
+    /** Put this map in the state where it has just the given value of the given key */
+    public void replace(K key, V value) {
+        List<V> list = map.get(key);
+        if (list == null) {
+            put(key);
+        }
+        else {
+            list.clear();
+            list.add(value);
+        }
+    }
+
     public void removeAll(K key) {
         map.remove(key);
     }