Eliminate renames reuse

2 files changed, 28 insertions, 30 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 10d39a43c61..e8291470054 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
@@ -7,7 +7,6 @@ import ai.vespa.rankingexpression.importer.operations.Rename;
 import com.yahoo.collections.ListMap;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -35,7 +34,7 @@ public class DimensionRenamer {
     private final ListMap<String, Integer> variables = new ListMap<>();
     private final ListMap<Arc, Constraint> constraints = new ListMap<>();
 
-    /** The solution to this, or null if no solution is found (yet) */
+    /** The solution to this, or null if no solution is found yet */
     private Map<String, Integer> renames = null;
 
     public DimensionRenamer(IntermediateGraph graph) {
@@ -89,11 +88,8 @@ public class DimensionRenamer {
      * @return the solution in the form of the renames to perform
      */
     private Map<String, Integer> solve(int maxIterations) {
-        Map<String, Integer> renames = new HashMap<>();
-
-        // Todo: evaluate possible improved efficiency by using a heuristic such as min-conflicts
-        boolean solved = NamingConstraintSolver.solve(variables, constraints, maxIterations, renames);
-        if ( ! solved) {
+        Map<String, Integer> solution = NamingConstraintSolver.solve(variables, constraints, maxIterations);
+        if ( solution == null) {
             IntermediateOperation operation = graph.operations().get("dense_out/MatMul");
             if (operation != null && operation instanceof MatMul) {
                 IntermediateOperation arg0 = operation.inputs().get(0);
@@ -109,17 +105,15 @@ public class DimensionRenamer {
                 addDimension("renamed_0");
                 newOperation.addDimensionNameConstraints(this);
 
-                renames.clear();
-                solved = NamingConstraintSolver.solve(variables, constraints, maxIterations, renames);
+                solution = NamingConstraintSolver.solve(variables, constraints, maxIterations);
             }
         }
-        if ( ! solved) {
-            renames.clear();
+        if ( solution == null) {
             ListMap<Arc, Constraint> hardConstraints = new ListMap<>();
             boolean anyRemoved = copyHard(constraints, hardConstraints);
             if (anyRemoved)
-                solved = NamingConstraintSolver.solve(variables, hardConstraints, maxIterations, renames);
-            if ( ! solved) {
+                solution = NamingConstraintSolver.solve(variables, hardConstraints, maxIterations);
+            if ( solution == null) {
                 throw new IllegalArgumentException("Could not find a dimension naming solution " +
                                                    "given constraints\n" + constraintsToString(hardConstraints));
             }
@@ -130,7 +124,7 @@ public class DimensionRenamer {
         // with the most remaining constraints, and inject a rename operation.
         // Then run this algorithm again.
 
-        return renames;
+        return solution;
     }
 
     /** Removes soft constraints and returns whether something was removed */
diff --git a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/NamingConstraintSolver.java b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/NamingConstraintSolver.java
index b059bb96d91..54f97a3427c 100644
--- a/model-integration/src/main/java/ai/vespa/rankingexpression/importer/NamingConstraintSolver.java
+++ b/model-integration/src/main/java/ai/vespa/rankingexpression/importer/NamingConstraintSolver.java
@@ -6,6 +6,7 @@ import com.yahoo.lang.MutableInteger;
 
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -24,33 +25,31 @@ class NamingConstraintSolver {
     private int iterations = 0;
     private final int maxIterations;
 
-    /** The solution to this rename problem */
-    private Map<String, Integer> renames;
-
     private NamingConstraintSolver(ListMap<String, Integer> inputVariables,
                                    ListMap<DimensionRenamer.Arc, DimensionRenamer.Constraint> constraints,
-                                   int maxIterations,
-                                   Map<String, Integer> renames) {
+                                   int maxIterations) {
         this.variables = new ListMap<>(inputVariables);
         initialize(variables);
         this.constraints = constraints;
         this.maxIterations = maxIterations;
-        this.renames = renames;
     }
 
     /** Try the solve the constraint problem given in the arguments, and put the result in renames */
-    private boolean trySolve() {
+    private Map<String, Integer> trySolve() {
+        // TODO: Evaluate possible improved efficiency by using a heuristic such as min-conflicts
+
+        Map<String, Integer> solution = new HashMap<>();
         for (String dimension : variables.keySet()) {
             List<Integer> values = variables.get(dimension);
             if (values.size() > 1) {
-                if ( ! ac3()) return false;
+                if ( ! ac3()) return null;
                 values.sort(Integer::compare);
                 variables.replace(dimension, values.get(0));
             }
-            renames.put(dimension, variables.get(dimension).get(0));
-            if (iterations > maxIterations) return false;
+            solution.put(dimension, variables.get(dimension).get(0));
+            if (iterations > maxIterations) return null;
         }
-        return true;
+        return solution;
     }
 
     private static void initialize(ListMap<String, Integer> variables) {
@@ -101,11 +100,16 @@ class NamingConstraintSolver {
         return revised;
     }
 
-    public static boolean solve(ListMap<String, Integer> inputVariables,
-                                ListMap<DimensionRenamer.Arc, DimensionRenamer.Constraint> constraints,
-                                int maxIterations,
-                                Map<String, Integer> renames) {
-        return new NamingConstraintSolver(inputVariables, constraints, maxIterations, renames).trySolve();
+    /**
+     * Attempts to solve the given naming problem. The input maps are never modified.
+     *
+     * @return the solution as a map from existing names to name ids represented as integers, or NULL
+     *         if no solution could be found
+     */
+    public static Map<String, Integer> solve(ListMap<String, Integer> inputVariables,
+                                             ListMap<DimensionRenamer.Arc, DimensionRenamer.Constraint> constraints,
+                                             int maxIterations) {
+        return new NamingConstraintSolver(inputVariables, constraints, maxIterations).trySolve();
     }
 
 }