Extend tensor mapper to handle unbound indexed dimensions in type spec.

2 files changed, 162 insertions, 21 deletions

diff --git a/vespalib/src/tests/tensor/tensor_mapper/tensor_mapper_test.cpp b/vespalib/src/tests/tensor/tensor_mapper/tensor_mapper_test.cpp
index 18d8a8dd508..f4edd8901e4 100644
--- a/vespalib/src/tests/tensor/tensor_mapper/tensor_mapper_test.cpp
+++ b/vespalib/src/tests/tensor/tensor_mapper/tensor_mapper_test.cpp
@@ -189,6 +189,52 @@ testTensorMapper(FixtureType &f)
                                   {{{"x","10"},{"y","1"}}, 7}
                               },
                               { "x", "y" }));
+    TEST_DO(f.assertDenseMap({
+                                  {{{"x",0},{"y",0}}, 1},
+                                  {{{"x",0},{"y",1}}, 5},
+                                  {{{"x",1},{"y",0}}, 3},
+                                  {{{"x",1},{"y",1}}, 0}
+                              },
+                             "tensor(x[2], y[])",
+                              {
+                                  {{{"x","0"},{"y","0"}}, 1},
+                                  {{{"x","1"},{"y","0"}}, 3},
+                                  {{{"x","0"},{"y","1"}}, 5},
+                                  {{{"x","10"},{"y","1"}}, 7}
+                              },
+                              { "x", "y" }));
+    TEST_DO(f.assertDenseMap({
+                                  {{{"x",0},{"y",0}}, 1},
+                                  {{{"x",0},{"y",1}}, 5},
+                                  {{{"x",1},{"y",0}}, 3},
+                                  {{{"x",1},{"y",1}}, 0},
+                                  {{{"x",2},{"y",0}}, 7},
+                                  {{{"x",2},{"y",1}}, 0}
+                              },
+                             "tensor(x[], y[])",
+                              {
+                                  {{{"x","0"},{"y","0"}}, 1},
+                                  {{{"x","1"},{"y","0"}}, 3},
+                                  {{{"x","0"},{"y","1"}}, 5},
+                                  {{{"x","2"},{"y","0"}}, 7}
+                              },
+                              { "x", "y" }));
+    TEST_DO(f.assertDenseMap({
+                                  {{{"x",0},{"y",0}}, 1},
+                                  {{{"x",0},{"y",1}}, 5},
+                                  {{{"x",0},{"y",2}}, 0},
+                                  {{{"x",1},{"y",0}}, 3},
+                                  {{{"x",1},{"y",1}}, 0},
+                                  {{{"x",1},{"y",2}}, 0}
+                              },
+                             "tensor(x[], y[3])",
+                              {
+                                  {{{"x","0"},{"y","0"}}, 1},
+                                  {{{"x","1"},{"y","0"}}, 3},
+                                  {{{"x","0"},{"y","1"}}, 5},
+                                  {{{"x","10"},{"y","3"}}, 7}
+                              },
+                              { "x", "y" }));
 }
 
 TEST_F("test tensor mapper for SparseTensor", SparseFixture)
diff --git a/vespalib/src/vespa/vespalib/tensor/tensor_mapper.cpp b/vespalib/src/vespa/vespalib/tensor/tensor_mapper.cpp
index f740ffbf348..c14601cd5cf 100644
--- a/vespalib/src/vespa/vespalib/tensor/tensor_mapper.cpp
+++ b/vespalib/src/vespa/vespalib/tensor/tensor_mapper.cpp
@@ -98,15 +98,124 @@ SparseTensorMapper<TensorT>::map(const Tensor &tensor,
     return mapper.build();
 }
 
+static constexpr uint32_t BAD_LABEL = std::numeric_limits<uint32_t>::max();
+static constexpr uint32_t BAD_ADDRESS = std::numeric_limits<uint32_t>::max();
+
+uint32_t mapLabelToNumber(vespalib::stringref label) {
+    uint32_t result = 0;
+    for (char c : label) {
+        if (c < '0' || c > '9') {
+            return BAD_LABEL; // bad char
+        }
+        result = result * 10 + (c - '0');
+        if (result > 100000000) {
+            return BAD_LABEL; // overflow
+        }
+    }
+    return result;
+}
+
+class DenseTensorBoundsMapper : public TensorVisitor
+{
+    ValueType _type;
+    std::vector<ValueType::Dimension> _dimensions;
+
+    bool addressOK(const TensorAddress &address);
+    void expandUnboundDimensions(const TensorAddress &address);
+
+    virtual void visit(const TensorAddress &address, double value) override;
+
+    DenseTensorBoundsMapper(const ValueType &type);
+    ~DenseTensorBoundsMapper();
+
+    ValueType build();
+public:
+    static ValueType map(const Tensor &tensor, const ValueType &type);
+};
+
+bool
+DenseTensorBoundsMapper::addressOK(const TensorAddress &address)
+{
+    TensorAddressElementIterator<TensorAddress> addressIterator(address);
+    auto dimIterator = _dimensions.begin();
+    for (const auto &dimension : _type.dimensions()) {
+        if (addressIterator.skipToDimension(dimension.name)) {
+            uint32_t label = mapLabelToNumber(addressIterator.label());
+            if (label == BAD_LABEL ||
+                (dimension.is_bound() && label >= dimIterator->size)) {
+                return false;
+            }
+            addressIterator.next();
+        }
+        ++dimIterator;
+    }
+    assert(dimIterator == _dimensions.end());
+    return true;
+}
+
+
+void
+DenseTensorBoundsMapper::expandUnboundDimensions(const TensorAddress &address)
+{
+    TensorAddressElementIterator<TensorAddress> addressIterator(address);
+    auto dimIterator = _dimensions.begin();
+    for (const auto &dimension : _type.dimensions()) {
+        if (addressIterator.skipToDimension(dimension.name)) {
+            uint32_t label = mapLabelToNumber(addressIterator.label());
+            if (label != BAD_LABEL &&
+                !dimension.is_bound() &&
+                label >= dimIterator->size) {
+                dimIterator->size = label + 1;
+            }
+            addressIterator.next();
+        }
+        ++dimIterator;
+    }
+    assert(dimIterator == _dimensions.end());
+}
+
+void
+DenseTensorBoundsMapper::visit(const TensorAddress &address, double value)
+{
+    (void) value;
+    if (addressOK(address)) {
+        expandUnboundDimensions(address);
+    }
+}
+
+DenseTensorBoundsMapper::DenseTensorBoundsMapper(const ValueType &type)
+    : _type(type),
+      _dimensions(type.dimensions())
+{
+    for (auto &dimension : _dimensions) {
+        if (!dimension.is_bound())
+            dimension.size = 1;
+    }
+}
+
+DenseTensorBoundsMapper::~DenseTensorBoundsMapper()
+{
+}
+
+ValueType
+DenseTensorBoundsMapper::build()
+{
+    return ValueType::tensor_type(std::move(_dimensions));
+}
+
+ValueType
+DenseTensorBoundsMapper::map(const Tensor &tensor, const ValueType &type)
+{
+    DenseTensorBoundsMapper mapper(type);
+    tensor.accept(mapper);
+    return mapper.build();
+}
+
 class DenseTensorMapper : public TensorVisitor
 {
     eval::ValueType _type;
     DenseTensor::Cells _cells;
-    static constexpr uint32_t BAD_LABEL = std::numeric_limits<uint32_t>::max();
-    static constexpr uint32_t BAD_ADDRESS =
-        std::numeric_limits<uint32_t>::max();
 
-    static uint32_t mapLabelToNumber(vespalib::stringref label);
     uint32_t mapAddressToIndex(const TensorAddress &address);
     virtual void visit(const TensorAddress &address, double value) override;
 
@@ -142,22 +251,6 @@ DenseTensorMapper::build()
 }
 
 uint32_t
-DenseTensorMapper::mapLabelToNumber(vespalib::stringref label)
-{
-    uint32_t result = 0;
-    for (char c : label) {
-        if (c < '0' || c > '9') {
-            return BAD_LABEL; // bad char
-        }
-        result = result * 10 + (c - '0');
-        if (result > 100000000) {
-            return BAD_LABEL; // overflow
-        }
-    }
-    return result;
-}
-
-uint32_t
 DenseTensorMapper::mapAddressToIndex(const TensorAddress &address)
 {
     uint32_t idx = 0;
@@ -191,7 +284,9 @@ DenseTensorMapper::visit(const TensorAddress &address, double value)
 std::unique_ptr<Tensor>
 DenseTensorMapper::map(const Tensor &tensor, const ValueType &type)
 {
-    DenseTensorMapper mapper(type);
+    DenseTensorMapper mapper(type.is_abstract() ?
+                             DenseTensorBoundsMapper::map(tensor, type) :
+                             type);
     tensor.accept(mapper);
     return mapper.build();
 }