Merge pull request #30009 from vespa-engine/balder/unify-with-compute-if-absent

- Replace usage of putIfAbsent and manual optimizes get-if-not-put wi…

5 files changed, 8 insertions, 16 deletions

diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json
index 6eb0aaf340e..df75a6f6d1f 100644
--- a/vespajlib/abi-spec.json
+++ b/vespajlib/abi-spec.json
@@ -1487,6 +1487,7 @@
       "public java.util.Set dimensionNames()",
       "public java.util.Optional dimension(java.lang.String)",
       "public java.util.Optional indexOfDimension(java.lang.String)",
+      "public int indexOfDimensionAsInt(java.lang.String)",
       "public java.util.Optional sizeOfDimension(java.lang.String)",
       "public boolean isAssignableTo(com.yahoo.tensor.TensorType)",
       "public boolean isConvertibleTo(com.yahoo.tensor.TensorType)",
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
index dd6838bfee3..95d1d70118a 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MixedTensor.java
@@ -9,7 +9,6 @@ import com.yahoo.tensor.impl.StringTensorAddress;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -306,12 +305,7 @@ public class MixedTensor implements Tensor {
         }
 
         private double[] denseSubspace(TensorAddress sparseAddress) {
-            double [] values = denseSubspaceMap.get(sparseAddress);
-            if (values == null) {
-                values = new double[(int)denseSubspaceSize()];
-                denseSubspaceMap.put(sparseAddress, values);
-            }
-            return values;
+            return denseSubspaceMap.computeIfAbsent(sparseAddress, (key) -> new double[(int)denseSubspaceSize()]);
         }
 
         public IndexedTensor.DirectIndexBuilder denseSubspaceBuilder(TensorAddress sparseAddress) {
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
index 4c65977fdc9..dcfee88d599 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
@@ -200,7 +200,7 @@ public class TensorType {
         return Optional.empty();
     }
     /** Returns the 0-base index of this dimension, or empty if it is not present */
-    int indexOfDimensionAsInt(String dimension) {
+    public int indexOfDimensionAsInt(String dimension) {
         for (int i = 0; i < dimensions.size(); i++)
             if (dimensions.get(i).name().equals(dimension))
                 return i;
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
index fa9dcda9179..e0ac549651c 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
@@ -335,12 +335,11 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
         int[] bIndexesInJoined = mapIndexes(b.type(), joinedType);
 
         // Iterate once through the smaller tensor and construct a hash map for common dimensions
-        Map<TensorAddress, List<Tensor.Cell>> aCellsByCommonAddress = new HashMap<>();
+        Map<TensorAddress, List<Tensor.Cell>> aCellsByCommonAddress = new HashMap<>(a.sizeAsInt());
         for (Iterator<Tensor.Cell> cellIterator = a.cellIterator(); cellIterator.hasNext(); ) {
             Tensor.Cell aCell = cellIterator.next();
             TensorAddress partialCommonAddress = partialCommonAddress(aCell, aIndexesInCommon);
-            aCellsByCommonAddress.putIfAbsent(partialCommonAddress, new ArrayList<>());
-            aCellsByCommonAddress.get(partialCommonAddress).add(aCell);
+            aCellsByCommonAddress.computeIfAbsent(partialCommonAddress, (key) -> new ArrayList<>()).add(aCell);
         }
 
         // Iterate once through the larger tensor and use the hash map to find joinable cells
@@ -363,7 +362,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
     }
 
     /**
-     * Returns the an array having one entry in order for each dimension of fromType
+     * Returns an array having one entry in order for each dimension of fromType
      * containing the index at which toType contains the same dimension name.
      * That is, if the returned array contains n at index i then
      * fromType.dimensions().get(i).name.equals(toType.dimensions().get(n).name())
@@ -372,7 +371,7 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP
     static int[] mapIndexes(TensorType fromType, TensorType toType) {
         int[] toIndexes = new int[fromType.dimensions().size()];
         for (int i = 0; i < fromType.dimensions().size(); i++)
-            toIndexes[i] = toType.indexOfDimension(fromType.dimensions().get(i).name()).orElse(-1);
+            toIndexes[i] = toType.indexOfDimensionAsInt(fromType.dimensions().get(i).name());
         return toIndexes;
     }
 
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
index 3ba57b29ebc..77e82b818a7 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
@@ -137,9 +137,7 @@ public class Reduce<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMET
         for (Iterator<Tensor.Cell> i = argument.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> cell = i.next();
             TensorAddress reducedAddress = reduceDimensions(indexesToKeep, cell.getKey());
-            ValueAggregator aggr = aggregatingCells.putIfAbsent(reducedAddress, ValueAggregator.ofType(aggregator));
-            if (aggr == null)
-                aggr = aggregatingCells.get(reducedAddress);
+            ValueAggregator aggr = aggregatingCells.computeIfAbsent(reducedAddress, (key) ->ValueAggregator.ofType(aggregator));
             aggr.aggregate(cell.getValue());
         }
         Tensor.Builder reducedBuilder = Tensor.Builder.of(reducedType);