1 files changed, 7 insertions, 32 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 e9be35b6f84..bcddfcbfc13 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
@@ -121,7 +121,7 @@ public class DimensionRenamer {
         List<IntermediateOperation> prioritizedOperations =
                 constraintsPerOperation.entrySet().stream()
                                                   .sorted(Comparator.comparingInt(entry -> - entry.getValue()))
-                                                  .map(entry -> entry.getKey())
+                                                  .map(Map.Entry::getKey)
                                                   .collect(Collectors.toList());
 
         List<RenameTarget> targets = new ArrayList<>();
@@ -131,8 +131,7 @@ public class DimensionRenamer {
                 if (inputType.isEmpty()) continue;
                 for (String dimensionName : inputType.get().dimensionNames()) {
                     RenameTarget target = new RenameTarget(operation, i, dimensionName, graph);
-                    if (target.rootKey != null) // TODO: Inserting renames under non-roots is not implemented
-                        targets.add(target);
+                    targets.add(target);
                 }
             }
         }
@@ -313,17 +312,10 @@ public class DimensionRenamer {
         final String dimensionName;
         final IntermediateGraph graph;
 
-        /**
-         * Returns the key of this operation in the root operations of the graph,
-         * or null if it is not a root operation
-         */
-        final String rootKey;
-
         public RenameTarget(IntermediateOperation operation, int inputNumber, String dimensionName, IntermediateGraph graph) {
             this.operation = operation;
             this.inputNumber = inputNumber;
             this.dimensionName = dimensionName;
-            this.rootKey = findRootKey(operation, graph);
             this.graph = graph;
         }
 
@@ -331,42 +323,25 @@ public class DimensionRenamer {
             return operation.inputs().get(inputNumber);
         }
 
-        private static String findRootKey(IntermediateOperation operation, IntermediateGraph graph) {
-            for (var entry : graph.operations().entrySet()) {
-                if (entry.getValue() == operation)
-                    return entry.getKey();
-            }
-            return null;
-        }
-
         /** Inserts a rename operation if possible. Returns whether an operation was inserted. */
         private boolean insertRename(DimensionRenamer renamer) {
             Rename rename = new Rename(operation.modelName(),
                                        dimensionName,
                                        renamer.dimensionPrefix + renamer.dimensions.size(),
                                        input());
-
-            List<IntermediateOperation> newInputs = new ArrayList<>(operation.inputs());
-            newInputs.set(inputNumber, rename);
-            IntermediateOperation newOperation = operation.withInputs(newInputs);
-            if (rootKey == null)
-                throw new IllegalStateException("Renaming non-roots is not implemented");
-            graph.put(rootKey, newOperation);
-
+            operation.insert(rename, inputNumber);
             removeConstraintsOf(operation, renamer);
             rename.addDimensionNameConstraints(renamer);
-            newOperation.addDimensionNameConstraints(renamer);
+            operation.addDimensionNameConstraints(renamer);
             return true;
         }
 
         /** Undo what insertRenameOperation has done: Set back the original operation and remove+add constraints */
         private void uninsertRename(DimensionRenamer renamer) {
-            IntermediateOperation newOperation = graph.operations().get(rootKey);
-            Rename rename = (Rename)newOperation.inputs().get(inputNumber);
-            graph.put(rootKey, operation);
-
+            Rename rename = (Rename)operation.inputs().get(inputNumber);
+            operation.uninsert(inputNumber);
             removeConstraintsOf(rename, renamer);
-            removeConstraintsOf(newOperation, renamer);
+            removeConstraintsOf(operation, renamer);
             operation.addDimensionNameConstraints(renamer);
         }