Use Tensor.Cell in iteration interface

8 files changed, 75 insertions, 55 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
index b89185b5131..509877f48f6 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/IndexedTensor.java
@@ -48,7 +48,7 @@ public class IndexedTensor implements Tensor {
      * indexes of later dimensions in the dimension type before earlier.
      */
     @Override
-    public Iterator<Map.Entry<TensorAddress, Double>> cellIterator() {
+    public Iterator<Cell> cellIterator() {
         return new CellIterator();
     }
 
@@ -424,7 +424,7 @@ public class IndexedTensor implements Tensor {
 
     }
     
-    private final class CellIterator implements Iterator<Map.Entry<TensorAddress, Double>> {
+    private final class CellIterator implements Iterator<Cell> {
 
         private int count = 0;
         private final Indexes indexes = Indexes.of(dimensionSizes, dimensionSizes, values.length);
@@ -435,7 +435,7 @@ public class IndexedTensor implements Tensor {
         }
 
         @Override
-        public Map.Entry<TensorAddress, Double> next() {
+        public Cell next() {
             if ( ! hasNext()) throw new NoSuchElementException("No cell at " + indexes);
             count++;
             indexes.next();
@@ -446,44 +446,6 @@ public class IndexedTensor implements Tensor {
         
     }
 
-    private class Cell implements Map.Entry<TensorAddress, Double> {
-
-        private final TensorAddress address;
-        private final Double value;
-
-        private Cell(TensorAddress address, Double value) {
-            this.address = address;
-            this.value = value;
-        }
-
-        @Override
-        public TensorAddress getKey() { return address; }
-
-        @Override
-        public Double getValue() { return value; }
-
-        @Override
-        public Double setValue(Double value) {
-            throw new UnsupportedOperationException("A tensor cannot be modified");
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (o == this) return true;
-            if ( ! ( o instanceof Map.Entry)) return false;
-            Map.Entry<?,?> other = (Map.Entry)o;
-            if ( ! this.getValue().equals(other.getValue())) return false;
-            if ( ! this.getKey().equals(other.getKey())) return false;
-            return true;
-        }
-
-        @Override
-        public int hashCode() {
-            return getKey().hashCode() ^ getValue().hashCode(); // by Map.Entry spec
-        }
-
-    }
-
     private final class ValueIterator implements Iterator<Double> {
 
         private int count = 0;
@@ -612,7 +574,7 @@ public class IndexedTensor implements Tensor {
             indexes.next();
             int valueIndex = indexes.toValueIndex();
             TensorAddress address = indexes.toAddress(valueIndex);
-            return new Cell(address, get(valueIndex)); // TODO: Change type to Cell, then change Cell to work with indexes + valueIndex instead of creating an address?
+            return new Cell(address, get(valueIndex)); // TODO: Change type to Cell, then change Cell to work with indexes + valueIndex instead of creating an address
         }
 
     }
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
index 8d72e860473..034bf714e48 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/MappedTensor.java
@@ -35,7 +35,7 @@ public class MappedTensor implements Tensor {
     public double get(TensorAddress address) { return cells.getOrDefault(address, Double.NaN); }
 
     @Override
-    public Iterator<Map.Entry<TensorAddress, Double>> cellIterator() { return cells.entrySet().iterator(); }
+    public Iterator<Cell> cellIterator() { return new CellIteratorAdaptor(cells.entrySet().iterator()); }
 
     @Override
     public Iterator<Double> valueIterator() { return cells.values().iterator(); }
@@ -91,4 +91,24 @@ public class MappedTensor implements Tensor {
         }
     
     }
+
+    private static class CellIteratorAdaptor implements Iterator<Cell> {
+
+        private final Iterator<Map.Entry<TensorAddress, Double>> adaptedIterator;
+        
+        private CellIteratorAdaptor(Iterator<Map.Entry<TensorAddress, Double>> adaptedIterator) {
+            this.adaptedIterator = adaptedIterator;
+        }
+        
+        @Override
+        public boolean hasNext() { return adaptedIterator.hasNext(); }
+
+        @Override
+        public Cell next() {
+            Map.Entry<TensorAddress, Double> entry = adaptedIterator.next();
+            return new Cell(entry.getKey(), entry.getValue()); 
+        }
+
+    }
+
 }
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java b/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
index 4a3302d7a71..a898ea5ad0c 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/Tensor.java
@@ -65,7 +65,7 @@ public interface Tensor {
     double get(TensorAddress address);
 
     /** Returns the cell of this in some undefined order */
-    Iterator<Map.Entry<TensorAddress, Double>> cellIterator();
+    Iterator<Cell> cellIterator();
 
     /** Returns the values of this in some undefined order */
     Iterator<Double> valueIterator();
@@ -248,8 +248,8 @@ public interface Tensor {
         if (a == b) return true;
         if ( ! a.type().mathematicallyEquals(b.type())) return false;
         if ( a.size() != b.size()) return false;
-        for (Iterator<Map.Entry<TensorAddress, Double>> aIterator = a.cellIterator(); aIterator.hasNext(); ) {
-            Map.Entry<TensorAddress, Double> aCell = aIterator.next();
+        for (Iterator<Cell> aIterator = a.cellIterator(); aIterator.hasNext(); ) {
+            Cell aCell = aIterator.next();
             if ( ! aCell.getValue().equals(b.get(aCell.getKey()))) return false;
         }
         return true;
@@ -285,7 +285,45 @@ public interface Tensor {
     static Tensor from(String tensorString) {
         return TensorParser.tensorFrom(tensorString, Optional.empty());
     }
-    
+
+    class Cell implements Map.Entry<TensorAddress, Double> {
+
+        private final TensorAddress address;
+        private final Double value;
+
+        Cell(TensorAddress address, Double value) {
+            this.address = address;
+            this.value = value;
+        }
+
+        @Override
+        public TensorAddress getKey() { return address; }
+
+        @Override
+        public Double getValue() { return value; }
+
+        @Override
+        public Double setValue(Double value) {
+            throw new UnsupportedOperationException("A tensor cannot be modified");
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (o == this) return true;
+            if ( ! ( o instanceof Map.Entry)) return false;
+            Map.Entry<?,?> other = (Map.Entry)o;
+            if ( ! this.getValue().equals(other.getValue())) return false;
+            if ( ! this.getKey().equals(other.getKey())) return false;
+            return true;
+        }
+
+        @Override
+        public int hashCode() {
+            return getKey().hashCode() ^ getValue().hashCode(); // by Map.Entry spec
+        }
+
+    }
+
     interface Builder {
 
         /** Creates a suitable builder for the given type */
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 030fdb754de..d55fe9e1ac8 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java
@@ -100,7 +100,7 @@ public class Join extends PrimitiveTensorFunction {
     /** When both tensors have the same dimensions, at most one cell matches a cell in the other tensor */
     private Tensor singleSpaceJoin(Tensor a, Tensor b, TensorType joinedType) {
         Tensor.Builder builder = Tensor.Builder.of(joinedType);
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = a.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = a.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> aCell = i.next();
             double bCellValue = b.get(aCell.getKey());
             if (Double.isNaN(bCellValue)) continue; // no match
@@ -188,7 +188,7 @@ public class Join extends PrimitiveTensorFunction {
     private Tensor generalSubspaceJoin(Tensor subspace, Tensor superspace, TensorType joinedType, boolean reversedArgumentOrder) {
         int[] subspaceIndexes = subspaceIndexes(superspace.type(), subspace.type());
         Tensor.Builder builder = Tensor.Builder.of(joinedType);
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = superspace.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = superspace.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> supercell = i.next();
             TensorAddress subaddress = mapAddressToSubspace(supercell.getKey(), subspaceIndexes);
             double subspaceValue = subspace.get(subaddress);
@@ -220,9 +220,9 @@ public class Join extends PrimitiveTensorFunction {
         int[] aToIndexes = mapIndexes(a.type(), joinedType);
         int[] bToIndexes = mapIndexes(b.type(), joinedType);
         Tensor.Builder builder = Tensor.Builder.of(joinedType);
-        for (Iterator<Map.Entry<TensorAddress, Double>> aIterator = a.cellIterator(); aIterator.hasNext(); ) {
+        for (Iterator<Tensor.Cell> aIterator = a.cellIterator(); aIterator.hasNext(); ) {
             Map.Entry<TensorAddress, Double> aCell = aIterator.next();
-            for (Iterator<Map.Entry<TensorAddress, Double>> bIterator = b.cellIterator(); bIterator.hasNext(); ) {
+            for (Iterator<Tensor.Cell> bIterator = b.cellIterator(); bIterator.hasNext(); ) {
                 Map.Entry<TensorAddress, Double> bCell = bIterator.next();
                 TensorAddress combinedAddress = combineAddresses(aCell.getKey(), aToIndexes,
                                                                  bCell.getKey(), bToIndexes, joinedType);
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Map.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Map.java
index b5ca6d3ccfb..e91182ab06d 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Map.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Map.java
@@ -53,7 +53,7 @@ public class Map extends PrimitiveTensorFunction {
     public Tensor evaluate(EvaluationContext context) {
         Tensor argument = argument().evaluate(context);
         Tensor.Builder builder = Tensor.Builder.of(argument.type());
-        for (Iterator<java.util.Map.Entry<TensorAddress, Double>> i = argument.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = argument.cellIterator(); i.hasNext(); ) {
             java.util.Map.Entry<TensorAddress, Double> cell = i.next();
             builder.cell(cell.getKey(), mapper.applyAsDouble(cell.getValue()));
         }
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 d528b4d3dc5..afe98d4bc07 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Reduce.java
@@ -112,7 +112,7 @@ public class Reduce extends PrimitiveTensorFunction {
         
         // Reduce cells
         Map<TensorAddress, ValueAggregator> aggregatingCells = new HashMap<>();
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = argument.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = argument.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> cell = i.next();
             TensorAddress reducedAddress = reduceDimensions(cell.getKey(), argument.type(), reducedType);
             aggregatingCells.putIfAbsent(reducedAddress, ValueAggregator.ofType(aggregator));
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
index 29dcb218843..c7f1ca2546f 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Rename.java
@@ -71,7 +71,7 @@ public class Rename extends PrimitiveTensorFunction {
         }
             
         Tensor.Builder builder = Tensor.Builder.of(renamedType);
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = tensor.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = tensor.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> cell = i.next();
             TensorAddress renamedAddress = rename(cell.getKey(), toIndexes);
             builder.cell(renamedAddress, cell.getValue());
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/serialization/SparseBinaryFormat.java b/vespajlib/src/main/java/com/yahoo/tensor/serialization/SparseBinaryFormat.java
index 8e15a994ed9..27a009b5e7e 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/SparseBinaryFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/SparseBinaryFormat.java
@@ -40,7 +40,7 @@ class SparseBinaryFormat implements BinaryFormat {
 
     private static void encodeCells(GrowableByteBuffer buffer, Tensor tensor) {
         buffer.putInt1_4Bytes(tensor.size());
-        for (Iterator<Map.Entry<TensorAddress, Double>> i = tensor.cellIterator(); i.hasNext(); ) {
+        for (Iterator<Tensor.Cell> i = tensor.cellIterator(); i.hasNext(); ) {
             Map.Entry<TensorAddress, Double> cell = i.next();
             encodeAddress(buffer, cell.getKey());
             buffer.putDouble(cell.getValue());