Simplify and test type check

1 files changed, 5 insertions, 27 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java b/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
index 3ff82ea774b..1c6d8170885 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
@@ -51,7 +51,11 @@ public class DenseBinaryFormat implements BinaryFormat {
         DimensionSizes sizes;
         if (optionalType.isPresent()) {
             type = optionalType.get();
-            sizes = decodeAndValidateDimensionSizes(type, buffer);
+            TensorType serializedType = decodeType(buffer);
+            if ( ! type.isAssignableTo(serializedType))
+                throw new IllegalArgumentException("Type/instance mismatch: A tensor of type " + serializedType + 
+                                                   " cannot be assigned to type " + type);
+            sizes = sizesFromType(serializedType);
         }
         else {
             type = decodeType(buffer);
@@ -62,32 +66,6 @@ public class DenseBinaryFormat implements BinaryFormat {
         return builder.build();
     }
 
-    private DimensionSizes decodeAndValidateDimensionSizes(TensorType type, GrowableByteBuffer buffer) {
-        int dimensionCount = buffer.getInt1_4Bytes();
-        if (type.dimensions().size() != dimensionCount)
-            throw new IllegalArgumentException("Type/instance mismatch: Instance has " + dimensionCount + 
-                                               " dimensions but type is " + type);
-        
-        DimensionSizes.Builder builder = new DimensionSizes.Builder(dimensionCount);
-        for (int i = 0; i < dimensionCount; i++) {
-            TensorType.Dimension expectedDimension = type.dimensions().get(i);
-
-            String encodedName = buffer.getUtf8String();
-            int encodedSize = buffer.getInt1_4Bytes();
-
-            if ( ! expectedDimension.name().equals(encodedName))
-                throw new IllegalArgumentException("Type/instance mismatch: Instance has '" + encodedName +
-                                                   "' as dimension " + i + " but type is " + type);
-
-            if (expectedDimension.size().isPresent() && expectedDimension.size().get() < encodedSize)
-                throw new IllegalArgumentException("Type/instance mismatch: Instance has size " + encodedSize + 
-                                                   " in " + expectedDimension  + " in type " + type);
-
-            builder.set(i, encodedSize);
-        }
-        return builder.build();
-    }
-
     private TensorType decodeType(GrowableByteBuffer buffer) {
         int dimensionCount = buffer.getInt1_4Bytes();
         TensorType.Builder builder = new TensorType.Builder();