From 16a4c711e2b20ba16b2088950e3bcf63fab93798 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 Mar 2024 15:13:04 +0100
Subject: [PATCH] Inline logic to set a value in `dyn.SetByPath` (#1261)

## Changes

This removes the need for the `allowMissingKeyInMap` option to the
private `visit` function and ensures that the body of the visit function
doesn't add or remove values of the configuration it traverses.

This in turn prepares for visiting a path pattern that yields more than
one callback, which doesn't match well with the now-removed option.

## Tests

Unit tests pass and fully cover the inlined code.
---
 libs/dyn/visit.go     |  6 +---
 libs/dyn/visit_set.go | 64 +++++++++++++++++++++++++++++++++++++++----
 2 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/libs/dyn/visit.go b/libs/dyn/visit.go
index d1a8d73b..ef055e40 100644
--- a/libs/dyn/visit.go
+++ b/libs/dyn/visit.go
@@ -45,10 +45,6 @@ type visitOptions struct {
 	// If this function returns an error, the original visit function call
 	// returns this error and the value is left unmodified.
 	fn func(Path, Value) (Value, error)
-
-	// If set, tolerate the absence of the last component in the path.
-	// This option is needed to set a key in a map that is not yet present.
-	allowMissingKeyInMap bool
 }
 
 func visit(v Value, prefix, suffix Path, opts visitOptions) (Value, error) {
@@ -76,7 +72,7 @@ func visit(v Value, prefix, suffix Path, opts visitOptions) (Value, error) {
 
 		// Lookup current value in the map.
 		ev, ok := m[component.key]
-		if !ok && !opts.allowMissingKeyInMap {
+		if !ok {
 			return InvalidValue, noSuchKeyError{prefix}
 		}
 
diff --git a/libs/dyn/visit_set.go b/libs/dyn/visit_set.go
index 3ad770f0..d0361981 100644
--- a/libs/dyn/visit_set.go
+++ b/libs/dyn/visit_set.go
@@ -1,5 +1,11 @@
 package dyn
 
+import (
+	"fmt"
+	"maps"
+	"slices"
+)
+
 // Set assigns a new value at the specified path in the specified value.
 // It is identical to [SetByPath], except that it takes a string path instead of a [Path].
 func Set(v Value, path string, nv Value) (Value, error) {
@@ -14,11 +20,59 @@ func Set(v Value, path string, nv Value) (Value, error) {
 // If successful, it returns the new value with all intermediate values copied and updated.
 // If the path doesn't exist, it returns InvalidValue and an error.
 func SetByPath(v Value, p Path, nv Value) (Value, error) {
-	return visit(v, EmptyPath, p, visitOptions{
-		fn: func(_ Path, _ Value) (Value, error) {
-			// Return the incoming value to set it.
-			return nv, nil
+	lp := len(p)
+	if lp == 0 {
+		return nv, nil
+	}
+
+	parent := p[:lp-1]
+	component := p[lp-1]
+
+	return visit(v, EmptyPath, parent, visitOptions{
+		fn: func(prefix Path, v Value) (Value, error) {
+			path := prefix.Append(component)
+
+			switch {
+			case component.isKey():
+				// Expect a map to be set if this is a key.
+				m, ok := v.AsMap()
+				if !ok {
+					return InvalidValue, fmt.Errorf("expected a map to index %q, found %s", path, v.Kind())
+				}
+
+				// Return an updated map value.
+				m = maps.Clone(m)
+				m[component.key] = nv
+				return Value{
+					v: m,
+					k: KindMap,
+					l: v.l,
+				}, nil
+
+			case component.isIndex():
+				// Expect a sequence to be set if this is an index.
+				s, ok := v.AsSequence()
+				if !ok {
+					return InvalidValue, fmt.Errorf("expected a sequence to index %q, found %s", path, v.Kind())
+				}
+
+				// Lookup current value in the sequence.
+				if component.index < 0 || component.index >= len(s) {
+					return InvalidValue, indexOutOfBoundsError{prefix}
+				}
+
+				// Return an updated sequence value.
+				s = slices.Clone(s)
+				s[component.index] = nv
+				return Value{
+					v: s,
+					k: KindSequence,
+					l: v.l,
+				}, nil
+
+			default:
+				panic("invalid component")
+			}
 		},
-		allowMissingKeyInMap: true,
 	})
 }