From 4b3b47ed9c7099500eba1f8d80c358559f1c3aef Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Tue, 6 Feb 2024 15:49:36 +0100
Subject: [PATCH] Ensure every variable reference is passed to lookup function

References to keys that themselves are also variable references were
shortcircuited in the previous approach. This meant that certain fields were
resolved even if the lookup function would have instructed to skip resolution.

To fix this we separate memoization of resolved variable references from
memoization of lookups. Now, every variable reference is passed through the
lookup function.
---
 libs/dyn/dynvar/resolve.go      | 72 +++++++++++++++++++++++----------
 libs/dyn/dynvar/resolve_test.go | 25 ++++++++++++
 2 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/libs/dyn/dynvar/resolve.go b/libs/dyn/dynvar/resolve.go
index b4e119b6d7..b5417cac20 100644
--- a/libs/dyn/dynvar/resolve.go
+++ b/libs/dyn/dynvar/resolve.go
@@ -38,12 +38,20 @@ func Resolve(in dyn.Value, fn Lookup) (out dyn.Value, err error) {
 	return resolver{in: in, fn: fn}.run()
 }
 
+type lookupResult struct {
+	v   dyn.Value
+	err error
+}
+
 type resolver struct {
 	in dyn.Value
 	fn Lookup
 
 	refs     map[string]ref
 	resolved map[string]dyn.Value
+
+	// Memoization for lookups.
+	lookups map[string]lookupResult
 }
 
 func (r resolver) run() (out dyn.Value, err error) {
@@ -84,8 +92,10 @@ func (r *resolver) collectVariableReferences() (err error) {
 }
 
 func (r *resolver) resolveVariableReferences() (err error) {
-	// Initialize map for resolved variables.
-	// We use this for memoization.
+	// Initialize cache for lookups.
+	r.lookups = make(map[string]lookupResult)
+
+	// Initialize cache for resolved variable references.
 	r.resolved = make(map[string]dyn.Value)
 
 	// Resolve each variable reference (in order).
@@ -95,7 +105,7 @@ func (r *resolver) resolveVariableReferences() (err error) {
 	keys := maps.Keys(r.refs)
 	sort.Strings(keys)
 	for _, key := range keys {
-		_, err := r.resolve(key, []string{key})
+		_, err := r.resolveRef(key, r.refs[key], []string{key})
 		if err != nil {
 			return err
 		}
@@ -104,29 +114,12 @@ func (r *resolver) resolveVariableReferences() (err error) {
 	return nil
 }
 
-func (r *resolver) resolve(key string, seen []string) (dyn.Value, error) {
+func (r *resolver) resolveRef(key string, ref ref, seen []string) (dyn.Value, error) {
 	// Check if we have already resolved this variable reference.
 	if v, ok := r.resolved[key]; ok {
 		return v, nil
 	}
 
-	ref, ok := r.refs[key]
-	if !ok {
-		// Perform lookup in the input.
-		p, err := dyn.NewPathFromString(key)
-		if err != nil {
-			return dyn.InvalidValue, err
-		}
-		v, err := r.fn(p)
-		if err != nil && dyn.IsNoSuchKeyError(err) {
-			return dyn.InvalidValue, fmt.Errorf(
-				"reference does not exist: ${%s}",
-				key,
-			)
-		}
-		return v, err
-	}
-
 	// This is an unresolved variable reference.
 	deps := ref.references()
 
@@ -143,7 +136,7 @@ func (r *resolver) resolve(key string, seen []string) (dyn.Value, error) {
 			)
 		}
 
-		v, err := r.resolve(dep, append(seen, dep))
+		v, err := r.resolveKey(dep, append(seen, dep))
 
 		// If we should skip resolution of this key, index j will hold an invalid [dyn.Value].
 		if errors.Is(err, ErrSkipResolution) {
@@ -191,6 +184,41 @@ func (r *resolver) resolve(key string, seen []string) (dyn.Value, error) {
 	return v, nil
 }
 
+func (r *resolver) resolveKey(key string, seen []string) (dyn.Value, error) {
+	// Check if we have already looked up this key.
+	if v, ok := r.lookups[key]; ok {
+		return v.v, v.err
+	}
+
+	// Parse the key into a path.
+	p, err := dyn.NewPathFromString(key)
+	if err != nil {
+		return dyn.InvalidValue, err
+	}
+
+	// Look up the value for the given key.
+	v, err := r.fn(p)
+	if err != nil {
+		if dyn.IsNoSuchKeyError(err) {
+			err = fmt.Errorf("reference does not exist: ${%s}", key)
+		}
+
+		// Cache the return value and return to the caller.
+		r.lookups[key] = lookupResult{v: dyn.InvalidValue, err: err}
+		return dyn.InvalidValue, err
+	}
+
+	// If the returned value is a valid variable reference, resolve it.
+	ref, ok := newRef(v)
+	if ok {
+		v, err = r.resolveRef(key, ref, seen)
+	}
+
+	// Cache the return value and return to the caller.
+	r.lookups[key] = lookupResult{v: v, err: err}
+	return v, err
+}
+
 func (r *resolver) replaceVariableReferences() (dyn.Value, error) {
 	// Walk the input and replace all variable references.
 	return dyn.Walk(r.in, func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
diff --git a/libs/dyn/dynvar/resolve_test.go b/libs/dyn/dynvar/resolve_test.go
index ba700503ef..1234b7cbfc 100644
--- a/libs/dyn/dynvar/resolve_test.go
+++ b/libs/dyn/dynvar/resolve_test.go
@@ -182,3 +182,28 @@ func TestResolveWithSkip(t *testing.T) {
 	assert.Equal(t, "a ${b}", getByPath(t, out, "e").MustString())
 	assert.Equal(t, "${b} a a ${b}", getByPath(t, out, "f").MustString())
 }
+
+func TestResolveWithSkipEverything(t *testing.T) {
+	in := dyn.V(map[string]dyn.Value{
+		"a": dyn.V("a"),
+		"b": dyn.V("b"),
+		"c": dyn.V("${a}"),
+		"d": dyn.V("${b}"),
+		"e": dyn.V("${a} ${b}"),
+		"f": dyn.V("${b} ${a} ${a} ${b}"),
+		"g": dyn.V("${d} ${c} ${c} ${d}"),
+	})
+
+	// The call must not replace anything if the lookup function returns ErrSkipResolution.
+	out, err := dynvar.Resolve(in, func(path dyn.Path) (dyn.Value, error) {
+		return dyn.InvalidValue, dynvar.ErrSkipResolution
+	})
+	require.NoError(t, err)
+	assert.Equal(t, "a", getByPath(t, out, "a").MustString())
+	assert.Equal(t, "b", getByPath(t, out, "b").MustString())
+	assert.Equal(t, "${a}", getByPath(t, out, "c").MustString())
+	assert.Equal(t, "${b}", getByPath(t, out, "d").MustString())
+	assert.Equal(t, "${a} ${b}", getByPath(t, out, "e").MustString())
+	assert.Equal(t, "${b} ${a} ${a} ${b}", getByPath(t, out, "f").MustString())
+	assert.Equal(t, "${d} ${c} ${c} ${d}", getByPath(t, out, "g").MustString())
+}
