Merge pull request #15533 from vespa-engine/arnej/renumber-peek-children-2

renumber Peek children, the sequel

5 files changed, 41 insertions, 41 deletions

diff --git a/eval/src/tests/eval/reference_operations/reference_operations_test.cpp b/eval/src/tests/eval/reference_operations/reference_operations_test.cpp
index b137f5e4a23..0495923018e 100644
--- a/eval/src/tests/eval/reference_operations/reference_operations_test.cpp
+++ b/eval/src/tests/eval/reference_operations/reference_operations_test.cpp
@@ -231,26 +231,26 @@ TEST(ReferencePeekTest, verbatim_labels) {
     ReferenceOperations::PeekSpec spec;
     spec.emplace("c", "qux");
     // peek 1 mapped dimension, verbatim label
-    auto output = ReferenceOperations::peek(input, spec, {});
+    auto output = ReferenceOperations::peek(spec, {input});
     auto expect = TensorSpec("tensor(e{})")
         .add({{"e","foo"}}, 16.0)
         .add({{"e","qux"}}, 25.0);
     EXPECT_EQ(output, expect);
     spec.emplace("e", "foo");
     // peek all mapped dimensions, verbatim labels
-    output = ReferenceOperations::peek(input, spec, {});    
+    output = ReferenceOperations::peek(spec, {input});    
     expect = TensorSpec("double").add({}, 16.0);
     EXPECT_EQ(output, expect);    
 
     spec.clear();
     spec.emplace("c", "nomatch");
     // peek 1 mapped dimension, non-matching verbatim label
-    output = ReferenceOperations::peek(input, spec, {});
+    output = ReferenceOperations::peek(spec, {input});
     expect = TensorSpec("tensor(e{})");
     EXPECT_EQ(output, expect);    
     spec.emplace("e", "nomatch");
     // peek all mapped dimensions, non-matching verbatim labels
-    output = ReferenceOperations::peek(input, spec, {});
+    output = ReferenceOperations::peek(spec, {input});
     expect = TensorSpec("double");
     EXPECT_EQ(output, expect);    
 
@@ -258,14 +258,14 @@ TEST(ReferencePeekTest, verbatim_labels) {
     spec.clear();
     spec.emplace("a", TensorSpec::Label(1));
     // peek 1 indexed dimension, verbatim label
-    output = ReferenceOperations::peek(input, spec, {});
+    output = ReferenceOperations::peek(spec, {input});
     expect = TensorSpec("tensor(d[5])")
         .add({{"d", 2}}, 3.0)
         .add({{"d", 0}}, 5.0);
     EXPECT_EQ(output, expect);
     spec.emplace("d", TensorSpec::Label(2));
     // peek all indexed dimensions, verbatim labels
-    output = ReferenceOperations::peek(input, spec, {});    
+    output = ReferenceOperations::peek(spec, {input});    
     expect = TensorSpec("double").add({}, 3.0);
     EXPECT_EQ(output, expect);    
 }    
@@ -275,32 +275,33 @@ TEST(ReferencePeekTest, labels_from_children) {
     auto zero_ch = TensorSpec("double").add({}, 0.0);
     auto neg_ch = TensorSpec("double").add({}, -2.0);
     auto too_big_ch = TensorSpec("double").add({}, 42.0);
-    std::vector<TensorSpec> children = {too_big_ch, too_big_ch, zero_ch, pos_ch, neg_ch, too_big_ch};
+    std::vector<TensorSpec> children = {TensorSpec(""), too_big_ch, too_big_ch, zero_ch, pos_ch, neg_ch, too_big_ch};
+    auto &input = children[0];
+    input = dense_2d_some_cells(false);
 
-    auto input = dense_2d_some_cells(false);
     ReferenceOperations::PeekSpec spec;
-    spec.emplace("a", size_t(3));
+    spec.emplace("a", size_t(4));
     // peek 1 indexed dimension, child (evaluating to 1.0)
-    auto output = ReferenceOperations::peek(input, spec, children);
+    auto output = ReferenceOperations::peek(spec, children);
     auto expect = TensorSpec("tensor(d[5])")
         .add({{"d", 2}}, 3.0)
         .add({{"d", 0}}, 5.0);
     EXPECT_EQ(output, expect);
-    spec.emplace("d", size_t(2));
+    spec.emplace("d", size_t(3));
     // peek 2 indexed dimensions (both children)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("double").add({}, 5.0);
     EXPECT_EQ(output, expect);
     spec.clear();
-    spec.emplace("a", size_t(0));
+    spec.emplace("a", size_t(1));
     // peek 1 indexed dimension, child (evaluating to 42.0)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("tensor(d[5])");
     EXPECT_EQ(output, expect);
     spec.clear();
-    spec.emplace("a", size_t(4));
+    spec.emplace("a", size_t(5));
     // peek 1 indexed dimension, child (evaluating to -2.0)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("tensor(d[5])");
     EXPECT_EQ(output, expect);
 
@@ -311,32 +312,32 @@ TEST(ReferencePeekTest, labels_from_children) {
         .add({{"c", "-2"}, {"e", "1"}}, 5.0)
         .add({{"c", "-2"}, {"e", "-2"}}, 6.0);
     spec.clear();
-    spec.emplace("c", size_t(3));
+    spec.emplace("c", size_t(4));
     // peek 1 mapped dimension, child (evaluating to 1.0)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("tensor(e{})")
         .add({{"e", "1"}}, 3.0)
         .add({{"e", "0"}}, 4.0);
     EXPECT_EQ(output, expect);
-    spec.emplace("e", size_t(2));
+    spec.emplace("e", size_t(3));
     // peek 2 mapped dimensions (both children)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("double").add({}, 4.0);
     EXPECT_EQ(output, expect);
 
     spec.clear();
-    spec.emplace("c", size_t(4));
+    spec.emplace("c", size_t(5));
     // peek 1 mapped dimension, child (evaluating to -2.0)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("tensor(e{})")
         .add({{"e", "1"}}, 5.0)
         .add({{"e", "-2"}}, 6.0);
     EXPECT_EQ(output, expect);
 
     spec.clear();
-    spec.emplace("c", size_t(0));
+    spec.emplace("c", size_t(1));
     // peek 1 indexed dimension, child (evaluating to 42.0)
-    output = ReferenceOperations::peek(input, spec, children);
+    output = ReferenceOperations::peek(spec, children);
     expect = TensorSpec("tensor(e{})");
     EXPECT_EQ(output, expect);
 }
@@ -346,7 +347,6 @@ TEST(ReferencePeekTest, peek_mixed) {
     auto zero_ch = TensorSpec("double").add({}, 0.0);
     auto neg_ch = TensorSpec("double").add({}, -2.0);
     auto too_big_ch = TensorSpec("double").add({}, 42.0);
-    std::vector<TensorSpec> children = {too_big_ch, too_big_ch, zero_ch, pos_ch, neg_ch, too_big_ch};
     auto input = TensorSpec("tensor(a[3],b[1],c{},d[5],e{})")
         .add({{"a", 0}, {"b", 0}, {"c", "-2"}, {"d", 1}, {"e", "foo"}},  1.0)
         .add({{"a", 0}, {"b", 0}, {"c", "1"}, {"d", 4}, {"e", "foo"}},   2.0) 
@@ -368,12 +368,13 @@ TEST(ReferencePeekTest, peek_mixed) {
         .add({{"a", 2}, {"b", 0}, {"c", "-2"}, {"d", 2}, {"e", "foo"}}, 18.0)
         .add({{"a", 2}, {"b", 0}, {"c", "0"}, {"d", 3}, {"e", "bar"}},  19.0) 
         .add({{"a", 2}, {"b", 0}, {"c", "1"}, {"d", 1}, {"e", "foo"}},  20.0);
+    std::vector<TensorSpec> children = {input, too_big_ch, too_big_ch, zero_ch, pos_ch, neg_ch, too_big_ch};
     ReferenceOperations::PeekSpec spec;
-    spec.emplace("a", size_t(3));
-    spec.emplace("b", size_t(2));
-    spec.emplace("c", size_t(4));
+    spec.emplace("a", size_t(4));
+    spec.emplace("b", size_t(3));
+    spec.emplace("c", size_t(5));
     spec.emplace("e", "foo");
-    auto output = ReferenceOperations::peek(input, spec, children);
+    auto output = ReferenceOperations::peek(spec, children);
     auto expect = TensorSpec("tensor(d[5])")
         .add({{"d", 1}}, 6.0)
         .add({{"d", 2}}, 8.0)
diff --git a/eval/src/tests/instruction/generic_peek/generic_peek_test.cpp b/eval/src/tests/instruction/generic_peek/generic_peek_test.cpp
index 08e445ca701..6841215038a 100644
--- a/eval/src/tests/instruction/generic_peek/generic_peek_test.cpp
+++ b/eval/src/tests/instruction/generic_peek/generic_peek_test.cpp
@@ -39,6 +39,7 @@ using PeekSpec = GenericPeek::SpecMap;
 
 TensorSpec reference_peek(const TensorSpec &param, const PeekSpec &spec) {
     std::vector<TensorSpec> children;
+    children.push_back(param);
     PeekSpec with_indexes;
     for (const auto & [dim_name, label_or_child] : spec) {
         const vespalib::string &dim = dim_name;
@@ -59,7 +60,7 @@ TensorSpec reference_peek(const TensorSpec &param, const PeekSpec &spec) {
                        }
                    }, label_or_child);
     }
-    return ReferenceOperations::peek(param, with_indexes, children);
+    return ReferenceOperations::peek(with_indexes, children);
 }
 
 TensorSpec perform_generic_peek(const TensorSpec &a, const ValueType &result_type,
diff --git a/eval/src/vespa/eval/eval/test/reference_evaluation.cpp b/eval/src/vespa/eval/eval/test/reference_evaluation.cpp
index 840e9a13f25..c20d8af32ec 100644
--- a/eval/src/vespa/eval/eval/test/reference_evaluation.cpp
+++ b/eval/src/vespa/eval/eval/test/reference_evaluation.cpp
@@ -107,12 +107,6 @@ struct EvalNode : public NodeVisitor {
     }
 
     void eval_peek(const TensorPeek &node) {
-        // TODO: fix Peek API so that the 'child index' sent in the
-        // spec is actually 'child index' (as defined by the function
-        // AST and Peek TensorFunction subclass) and not 'child index'
-        // - 1. This also means that the param (the object being
-        // peeked) should be sent as the first child and not as a
-        // separate parameter.
         TensorSpec param = eval_node(node.param(), params);
         ValueType param_type = ValueType::from_spec(param.type());
         auto is_indexed = [&](const vespalib::string &dim_name) {
@@ -121,6 +115,7 @@ struct EvalNode : public NodeVisitor {
                     (param_type.dimensions()[dim_idx].is_indexed()));
         };
         std::vector<TensorSpec> children;
+        children.push_back(param);
         std::map<vespalib::string, std::variant<TensorSpec::Label, size_t>> spec;
         for (const auto &[name, label]: node.dim_list()) {
             if (label.is_expr()) {
@@ -134,7 +129,7 @@ struct EvalNode : public NodeVisitor {
                 }
             }
         }
-        result = ReferenceOperations::peek(param, spec, children);
+        result = ReferenceOperations::peek(spec, children);
     }
 
     //-------------------------------------------------------------------------
diff --git a/eval/src/vespa/eval/eval/test/reference_operations.cpp b/eval/src/vespa/eval/eval/test/reference_operations.cpp
index bb053ce4dc0..27ecb0fd224 100644
--- a/eval/src/vespa/eval/eval/test/reference_operations.cpp
+++ b/eval/src/vespa/eval/eval/test/reference_operations.cpp
@@ -175,14 +175,15 @@ TensorSpec ReferenceOperations::merge(const TensorSpec &a, const TensorSpec &b,
 }
 
 
-TensorSpec ReferenceOperations::peek(const TensorSpec &param, const PeekSpec &peek_spec, const std::vector<TensorSpec> &children) {
-    if (peek_spec.empty()) {
+TensorSpec ReferenceOperations::peek(const PeekSpec &peek_spec, const std::vector<TensorSpec> &children) {
+    if (peek_spec.empty() || children.empty()) {
         return TensorSpec(ValueType::error_type().to_spec());
     }
     std::vector<vespalib::string> peek_dims;
     for (const auto & [dim_name, label_or_child] : peek_spec) {
         peek_dims.push_back(dim_name);
     }
+    const TensorSpec &param = children[0];
     ValueType param_type = ValueType::from_spec(param.type());
     ValueType result_type = param_type.reduce(peek_dims);
     TensorSpec result(result_type.to_spec());
diff --git a/eval/src/vespa/eval/eval/test/reference_operations.h b/eval/src/vespa/eval/eval/test/reference_operations.h
index 4663d935383..588215853f9 100644
--- a/eval/src/vespa/eval/eval/test/reference_operations.h
+++ b/eval/src/vespa/eval/eval/test/reference_operations.h
@@ -23,7 +23,9 @@ struct ReferenceOperations {
     // mapping from cell address to index of child that computes the cell value
     using CreateSpec = tensor_function::Create::Spec;
 
-    // mapping from dimension name to verbatim label or child
+    // mapping from dimension name to verbatim label or child index.
+    // Note: child 0 is the input param, so indexes in the spec must
+    // start at 1.
     using PeekSpec = tensor_function::Peek::Spec;
 
     static TensorSpec concat(const TensorSpec &a, const TensorSpec &b, const std::string &concat_dim);
@@ -31,7 +33,7 @@ struct ReferenceOperations {
     static TensorSpec join(const TensorSpec &a, const TensorSpec &b, join_fun_t function);
     static TensorSpec map(const TensorSpec &a, map_fun_t func);
     static TensorSpec merge(const TensorSpec &a, const TensorSpec &b, join_fun_t fun);
-    static TensorSpec peek(const TensorSpec &param, const PeekSpec &spec, const std::vector<TensorSpec> &children);
+    static TensorSpec peek(const PeekSpec &spec, const std::vector<TensorSpec> &children);
     static TensorSpec reduce(const TensorSpec &a, Aggr aggr, const std::vector<vespalib::string> &dims);
     static TensorSpec rename(const TensorSpec &a, const std::vector<vespalib::string> &from, const std::vector<vespalib::string> &to);
     static TensorSpec lambda(const vespalib::string &type, lambda_fun_t fun);