Retain information about sources when compiling

7 files changed, 89 insertions, 43 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/AllValuesQueryProfileVisitor.java b/container-search/src/main/java/com/yahoo/search/query/profile/AllValuesQueryProfileVisitor.java
index d2e0f64f968..93025a2d68c 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/AllValuesQueryProfileVisitor.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/AllValuesQueryProfileVisitor.java
@@ -2,8 +2,9 @@
 package com.yahoo.search.query.profile;
 
 import com.yahoo.processing.request.CompoundName;
-import com.yahoo.search.query.profile.compiled.ValueSource;
+import com.yahoo.search.query.profile.compiled.ValueWithSource;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -12,8 +13,7 @@ import java.util.Map;
  */
 final class AllValuesQueryProfileVisitor extends PrefixQueryProfileVisitor {
 
-    private Map<String, Object> values = new HashMap<>();
-    private Map<String, ValueSource> sources = new HashMap<>();
+    private Map<String, ValueWithSource> values = new HashMap<>();
 
     /* Lists all values starting at prefix */
     public AllValuesQueryProfileVisitor(CompoundName prefix) {
@@ -43,12 +43,19 @@ final class AllValuesQueryProfileVisitor extends PrefixQueryProfileVisitor {
         if (fullName.isEmpty()) return; // Avoid putting a non-leaf (subtree) root in the list
         if (values.containsKey(fullName.toString())) return; // The first value encountered has priority
 
-        values.put(fullName.toString(), value);
-        sources.put(fullName.toString(), new ValueSource(owner, variant));
+        values.put(fullName.toString(), new ValueWithSource(value, owner.getIdString(), variant));
     }
 
     /** Returns the values resulting from this visiting */
-    public Map<String, Object> getResult() { return values; }
+    public Map<String, Object> values() {
+        Map<String, Object> values = new HashMap<>();
+        for (var entry : this.values.entrySet())
+            values.put(entry.getKey(), entry.getValue().value());
+        return values;
+    }
+
+    /** Returns the values with source resulting from this visiting */
+    public Map<String, ValueWithSource> valuesWithSource() { return Collections.unmodifiableMap(values); }
 
     /** Returns false - we are not done until we have seen all */
     public boolean isDone() { return false; }
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfile.java b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfile.java
index 7b3f43d75e5..a311b55339c 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfile.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfile.java
@@ -234,11 +234,7 @@ public class QueryProfile extends FreezableSimpleComponent implements Cloneable
      * will return {"d" =&gt; "a.d-value","e" =&gt; "a.e-value"}
      */
     public Map<String, Object> listValues(CompoundName prefix, Map<String, String> context, Properties substitution) {
-        DimensionBinding dimensionBinding = DimensionBinding.createFrom(getDimensions(),context);
-
-        AllValuesQueryProfileVisitor visitor = new AllValuesQueryProfileVisitor(prefix);
-        accept(visitor,dimensionBinding, null);
-        Map<String, Object>  values = visitor.getResult();
+        Map<String, Object> values = visitValues(prefix, context).values();
 
         if (substitution == null) return values;
         for (Map.Entry<String, Object> entry : values.entrySet()) {
@@ -250,6 +246,14 @@ public class QueryProfile extends FreezableSimpleComponent implements Cloneable
         return values;
     }
 
+    AllValuesQueryProfileVisitor visitValues(CompoundName prefix, Map<String, String> context) {
+        DimensionBinding dimensionBinding = DimensionBinding.createFrom(getDimensions(), context);
+
+        AllValuesQueryProfileVisitor visitor = new AllValuesQueryProfileVisitor(prefix);
+        accept(visitor, dimensionBinding, null);
+        return visitor;
+    }
+
     /**
      * Lists types reachable from this, indexed by the prefix having that type.
      * If this is itself typed, this' type will be included with an empty prefix
@@ -510,7 +514,8 @@ public class QueryProfile extends FreezableSimpleComponent implements Cloneable
     protected Object lookup(CompoundName name,
                             boolean allowQueryProfileResult,
                             DimensionBinding dimensionBinding) {
-        SingleValueQueryProfileVisitor visitor = new SingleValueQueryProfileVisitor(name.asList(), allowQueryProfileResult);
+        SingleValueQueryProfileVisitor visitor = new SingleValueQueryProfileVisitor(name.asList(),
+                                                                                    allowQueryProfileResult);
         accept(visitor, dimensionBinding, null);
         return visitor.getResult();
     }
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileCompiler.java b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileCompiler.java
index 588071459f0..3cc7570576b 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileCompiler.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/QueryProfileCompiler.java
@@ -5,6 +5,7 @@ import com.yahoo.processing.request.CompoundName;
 import com.yahoo.search.query.profile.compiled.CompiledQueryProfile;
 import com.yahoo.search.query.profile.compiled.CompiledQueryProfileRegistry;
 import com.yahoo.search.query.profile.compiled.DimensionalMap;
+import com.yahoo.search.query.profile.compiled.ValueWithSource;
 import com.yahoo.search.query.profile.types.QueryProfileType;
 
 import java.util.HashSet;
@@ -31,7 +32,7 @@ public class QueryProfileCompiler {
 
     public static CompiledQueryProfile compile(QueryProfile in, CompiledQueryProfileRegistry registry) {
         try {
-            DimensionalMap.Builder<CompoundName, Object> values = new DimensionalMap.Builder<>();
+            DimensionalMap.Builder<CompoundName, ValueWithSource> values = new DimensionalMap.Builder<>();
             DimensionalMap.Builder<CompoundName, QueryProfileType> types = new DimensionalMap.Builder<>();
             DimensionalMap.Builder<CompoundName, Object> references = new DimensionalMap.Builder<>();
             DimensionalMap.Builder<CompoundName, Object> unoverridables = new DimensionalMap.Builder<>();
@@ -42,7 +43,7 @@ public class QueryProfileCompiler {
             log.fine(() -> "Compiling " + in.toString() + " having " + variants.size() + " variants");
             for (DimensionBindingForPath variant : variants) {
                 log.finer(() -> "  Compiling variant " + variant);
-                for (Map.Entry<String, Object> entry : in.listValues(variant.path(), variant.binding().getContext(), null).entrySet()) {
+                for (Map.Entry<String, ValueWithSource> entry : in.visitValues(variant.path(), variant.binding().getContext()).valuesWithSource().entrySet()) {
                     values.put(variant.path().append(entry.getKey()), variant.binding(), entry.getValue());
                 }
                 for (Map.Entry<CompoundName, QueryProfileType> entry : in.listTypes(variant.path(), variant.binding().getContext()).entrySet())
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/CompiledQueryProfile.java b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/CompiledQueryProfile.java
index ea85a2be242..06a6ac48cb6 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/CompiledQueryProfile.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/CompiledQueryProfile.java
@@ -12,7 +12,6 @@ import com.yahoo.search.query.profile.types.QueryProfileType;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -31,7 +30,7 @@ public class CompiledQueryProfile extends AbstractComponent implements Cloneable
     private final QueryProfileType type;
 
     /** The values of this */
-    private final DimensionalMap<CompoundName, Object> entries;
+    private final DimensionalMap<CompoundName, ValueWithSource> entries;
 
     /** Keys which have a type in this */
     private final DimensionalMap<CompoundName, QueryProfileType> types;
@@ -47,7 +46,7 @@ public class CompiledQueryProfile extends AbstractComponent implements Cloneable
      */
     public CompiledQueryProfile(ComponentId id,
                                 QueryProfileType type,
-                                DimensionalMap<CompoundName, Object> entries,
+                                DimensionalMap<CompoundName, ValueWithSource> entries,
                                 DimensionalMap<CompoundName, QueryProfileType> types,
                                 DimensionalMap<CompoundName, Object> references,
                                 DimensionalMap<CompoundName, Object> unoverridables,
@@ -131,11 +130,14 @@ public class CompiledQueryProfile extends AbstractComponent implements Cloneable
      */
     public Map<String, Object> listValues(CompoundName prefix, Map<String, String> context, Properties substitution) {
         Map<String, Object> values = new HashMap<>();
-        for (Map.Entry<CompoundName, DimensionalValue<Object>> entry : entries.entrySet()) {
+        for (Map.Entry<CompoundName, DimensionalValue<ValueWithSource>> entry : entries.entrySet()) {
             if ( entry.getKey().size() <= prefix.size()) continue;
             if ( ! entry.getKey().hasPrefix(prefix)) continue;
 
-            Object value = entry.getValue().get(context);
+            ValueWithSource valueWithSource = entry.getValue().get(context);
+            if (valueWithSource == null) continue;
+
+            Object value = valueWithSource.value();
             if (value == null) continue;
 
             value = substitute(value, context, substitution);
@@ -155,7 +157,9 @@ public class CompiledQueryProfile extends AbstractComponent implements Cloneable
         return get(new CompoundName(name), context, substitution);
     }
     public final Object get(CompoundName name, Map<String, String> context, Properties substitution) {
-        return substitute(entries.get(name, context), context, substitution);
+        ValueWithSource value = entries.get(name, context);
+        if (value == null) return null;
+        return substitute(value.value(), context, substitution);
     }
 
     private Object substitute(Object value, Map<String, String> context, Properties substitution) {
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/DimensionalValue.java b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/DimensionalValue.java
index 50d0a2de46f..b5481059ac0 100644
--- a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/DimensionalValue.java
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/DimensionalValue.java
@@ -168,12 +168,14 @@ public class DimensionalValue<VALUE> {
                 return value;
             }
 
+            // TODO: Move this
             @SuppressWarnings("unchecked")
             private VALUE substituteIfRelative(VALUE value,
                                                DimensionBinding variant,
                                                Map<CompoundName, DimensionalValue.Builder<VALUE>> entries) {
-                if (value instanceof SubstituteString) {
-                    SubstituteString substitute = (SubstituteString)value;
+                if (value instanceof ValueWithSource && ((ValueWithSource)value).value() instanceof SubstituteString) {
+                    ValueWithSource valueWithSource = (ValueWithSource)value;
+                    SubstituteString substitute = (SubstituteString)valueWithSource.value();
                     if (substitute.hasRelative()) {
                         List<SubstituteString.Component> resolvedComponents = new ArrayList<>(substitute.components().size());
                         for (SubstituteString.Component component : substitute.components()) {
@@ -184,14 +186,14 @@ public class DimensionalValue<VALUE> {
                                     throw new IllegalArgumentException("Could not resolve local substitution '" +
                                                                        relativeComponent.fieldName() + "' in variant " +
                                                                        variant);
-                                String resolved = substituteValues.valueFor(variant).toString();
-                                resolvedComponents.add(new SubstituteString.StringComponent(resolved));
+                                ValueWithSource resolved = (ValueWithSource)substituteValues.valueFor(variant);
+                                resolvedComponents.add(new SubstituteString.StringComponent(resolved.value().toString()));
                             }
                             else {
                                 resolvedComponents.add(component);
                             }
                         }
-                        return (VALUE)new SubstituteString(resolvedComponents, substitute.stringValue());
+                        return (VALUE)valueWithSource.withValue(new SubstituteString(resolvedComponents, substitute.stringValue()));
                     }
                 }
                 return value;
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueSource.java b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueSource.java
deleted file mode 100644
index 17326efb220..00000000000
--- a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueSource.java
+++ /dev/null
@@ -1,18 +0,0 @@
-// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.search.query.profile.compiled;
-
-import com.yahoo.search.query.profile.DimensionValues;
-import com.yahoo.search.query.profile.QueryProfile;
-
-/**
- * The source of a value in a compiled query profile
- *
- * @author bratseth
- */
-public class ValueSource {
-
-    public ValueSource(QueryProfile owner, DimensionValues variant) {
-
-    }
-
-}
diff --git a/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueWithSource.java b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueWithSource.java
new file mode 100644
index 00000000000..c2ce34c3f47
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/query/profile/compiled/ValueWithSource.java
@@ -0,0 +1,45 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.query.profile.compiled;
+
+import com.yahoo.search.query.profile.DimensionValues;
+
+import java.util.Optional;
+
+/**
+ * A value in a query profile with information about its source.
+ *
+ * @author bratseth
+ */
+public class ValueWithSource {
+
+    private final Object value;
+
+    /** The id of the query profile having a value */
+    private final String ownerId;
+
+    /** The dimension values specifying a variant in that profile, or null if it is not in a variant */
+    private final DimensionValues variant;
+
+    public ValueWithSource(Object value, String ownerId, DimensionValues variant) {
+        this.value = value;
+        this.ownerId = ownerId;
+        this.variant = variant;
+    }
+
+    public Object value() { return value; }
+
+    public String ownerId() { return ownerId; }
+
+    public ValueWithSource withValue(Object value) {
+        return new ValueWithSource(value, ownerId, variant);
+    }
+
+    /** Returns the variant having this value, or empty if it's not in a variant */
+    public Optional<DimensionValues> variant() { return Optional.ofNullable(variant); }
+
+    @Override
+    public String toString() {
+        return value + " from " + ownerId + ( variant != null ? " variant " + variant : "");
+    }
+
+}