Support values cached during execution of a script

2 files changed, 76 insertions, 0 deletions

diff --git a/indexinglanguage/src/main/java/com/yahoo/vespa/indexinglanguage/expressions/ExecutionContext.java b/indexinglanguage/src/main/java/com/yahoo/vespa/indexinglanguage/expressions/ExecutionContext.java
index 39e8d10facd..1935664cddc 100644
--- a/indexinglanguage/src/main/java/com/yahoo/vespa/indexinglanguage/expressions/ExecutionContext.java
+++ b/indexinglanguage/src/main/java/com/yahoo/vespa/indexinglanguage/expressions/ExecutionContext.java
@@ -21,6 +21,7 @@ public class ExecutionContext implements FieldTypeAdapter, FieldValueAdapter {
     private final FieldValueAdapter adapter;
     private FieldValue value;
     private Language language;
+    private Map<String, Object> cache = null;
 
     public ExecutionContext() {
         this(null);
@@ -118,6 +119,19 @@ public class ExecutionContext implements FieldTypeAdapter, FieldValueAdapter {
         return this;
     }
 
+    public void putCachedValue(String key, Object value) {
+        if (cache == null)
+            cache = new HashMap<>();
+        cache.put(key, value);
+    }
+
+    /** Returns a cached value, or null if not present. */
+    public Object getCachedValue(String key) {
+        if (cache == null) return null;
+        return cache.get(key);
+    }
+
+    /** Clears all state in this except the cache. */
     public ExecutionContext clear() {
         variables.clear();
         value = null;
diff --git a/indexinglanguage/src/test/java/com/yahoo/vespa/indexinglanguage/expressions/ScriptTestCase.java b/indexinglanguage/src/test/java/com/yahoo/vespa/indexinglanguage/expressions/ScriptTestCase.java
index 75f852f0331..df7f99d22d2 100644
--- a/indexinglanguage/src/test/java/com/yahoo/vespa/indexinglanguage/expressions/ScriptTestCase.java
+++ b/indexinglanguage/src/test/java/com/yahoo/vespa/indexinglanguage/expressions/ScriptTestCase.java
@@ -75,6 +75,20 @@ public class ScriptTestCase {
     }
 
     @Test
+    public void testCache() {
+        SimpleTestAdapter adapter = new SimpleTestAdapter(new Field("field1", DataType.STRING));
+        var script = newScript(newStatement(new InputExpression("field1"),
+                                            new PutCacheExpression("myCacheKey", "myCacheValue")),
+                               newStatement(new ClearStateExpression()), // inserted by config model
+                               newStatement(new InputExpression("field1"),
+                                            new AssertCacheExpression("myCacheKey", "myCacheValue")));
+        adapter.setValue("field1", new StringFieldValue("foo1"));
+        ExecutionContext context = new ExecutionContext(adapter);
+        script.execute(context);
+        assertEquals("myCacheValue", context.getCachedValue("myCacheKey"));
+    }
+
+    @Test
     public void requireThatStatementsProcessingMissingInputsAreSkipped() {
         SimpleTestAdapter adapter = new SimpleTestAdapter(new Field("foo", DataType.STRING),
                                                           new Field("bar", DataType.STRING));
@@ -155,4 +169,52 @@ public class ScriptTestCase {
 
     }
 
+    private static class PutCacheExpression extends Expression {
+
+        private final String keyToSet;
+        private final String valueToSet;
+
+        public PutCacheExpression(String keyToSet, String valueToSet) {
+            super(null);
+            this.keyToSet = keyToSet;
+            this.valueToSet = valueToSet;
+        }
+
+        @Override
+        protected void doExecute(ExecutionContext context) {
+            context.putCachedValue(keyToSet, valueToSet);
+        }
+
+        @Override
+        protected void doVerify(VerificationContext context) {}
+
+        @Override
+        public DataType createdOutputType() { return null; }
+
+    }
+
+    private static class AssertCacheExpression extends Expression {
+
+        private final String expectedKey;
+        private final String expectedValue;
+
+        public AssertCacheExpression(String expectedKey, String expectedValue) {
+            super(null);
+            this.expectedKey = expectedKey;
+            this.expectedValue = expectedValue;
+        }
+
+        @Override
+        protected void doExecute(ExecutionContext context) {
+            assertEquals(expectedValue, context.getCachedValue(expectedKey));
+        }
+
+        @Override
+        protected void doVerify(VerificationContext context) {}
+
+        @Override
+        public DataType createdOutputType() { return null; }
+
+    }
+
 }