1 files changed, 42 insertions, 27 deletions
diff --git a/jsonpath.c b/jsonpath.c
index b8e4527..9d53259 100644
--- a/jsonpath.c
+++ b/jsonpath.c
@@ -331,20 +331,6 @@ static size_t utf8_substring(
 	return sub_length;
 }
 
-static json_node *json_head(json_node *parent)
-{
-	switch (parent->type) {
-		case JSON_ARRAY:
-		case JSON_OBJECT:
-			return parent->v.children.head;
-		default:
-			return NULL;
-	}
-}
-
-#define json_foreach(child, parent) \
-	for ((child) = json_head(parent); (child) != NULL; (child) = (child)->next)
-
 static JPRef *json_index_subscript(JPRef *ref, long index)
 {
 	json_node *json;
@@ -412,7 +398,8 @@ JPRef could merged with json_node by adding JPRef's specialty types to the
 json_type enum and json_node union.  JPRef is currently not merged with
 json_node in an attempt to keep the codebase tidy and easier to extend.
 */
-static void match_recurse(List **results, ListCell *path, JPRef *ref)
+static void match_recurse(void on_match(void *ctx, JPRef *ref), void *ctx,
+	ListCell *path, JPRef *ref)
 {
 	jp_element	*elem;
 	JPRef		*child_ref;
@@ -420,7 +407,7 @@ static void match_recurse(List **results, ListCell *path, JPRef *ref)
 
 	if (path == NULL) {
 		/* The end of the JSONPath list is the "accept" state. */
-		*results = lappend(*results, ref);
+		on_match(ctx, ref);
 		return;
 	}
 
@@ -435,23 +422,25 @@ static void match_recurse(List **results, ListCell *path, JPRef *ref)
 		case JP_WILDCARD:
 			if (json) {
 				json_foreach(child, json)
-					match_recurse(results, lnext(path), mkRefNode(child));
+					match_recurse(on_match, ctx, lnext(path), mkRefNode(child));
 			}
 			break;
 
 		case JP_INDEX_SUBSCRIPT:
 			child_ref = json_index_subscript(ref, elem->data.index);
 			if (child_ref != NULL)
-				match_recurse(results, lnext(path), child_ref);
+				match_recurse(on_match, ctx, lnext(path), child_ref);
 			break;
 
 		case JP_KEY_SUBSCRIPT:
-			json_foreach(child, json) {
-				if (child->key != NULL &&
-				    child->key_length == elem->data.key.length &&
-					!memcmp(child->key, elem->data.key.ptr, child->key_length))
-				{
-					match_recurse(results, lnext(path), mkRefNode(child));
+			if (json->type == JSON_OBJECT) {
+				json_foreach(child, json) {
+					if (child->key != NULL &&
+						child->key_length == elem->data.key.length &&
+						!memcmp(child->key, elem->data.key.ptr, child->key_length))
+					{
+						match_recurse(on_match, ctx, lnext(path), mkRefNode(child));
+					}
 				}
 			}
 			break;
@@ -460,21 +449,47 @@ static void match_recurse(List **results, ListCell *path, JPRef *ref)
 	}
 
 	if (elem->recursive_descent && json) {
-		json_foreach(child, json)
-			match_recurse(results, path, mkRefNode(child));
+		json_foreach(child, json) {
+			if (!child->jp_changed)
+				match_recurse(on_match, ctx, path, mkRefNode(child));
+		}
 	}
 }
 
+static void jp_match_callback(List **results, JPRef *ref)
+{
+	*results = lappend(*results, ref);
+}
+
 List *jp_match(JSONPath *jp, json_node *json)
 {
 	ListCell	*lc			= jp_head(jp);
 	List		*results	= NIL;
 
-	match_recurse(&results, lc, mkRefNode(json));
+	match_recurse((void*)jp_match_callback, &results, lc, mkRefNode(json));
 
 	return results;
 }
 
+static void jp_set_callback(json_node *value, JPRef *ref)
+{
+	switch (ref->type) {
+		case JP_REF_NODE:
+			json_replace_value(ref->u.node, value);
+			ref->u.node->jp_changed = true;
+			break;
+
+		default:; /* Do nothing if ref is immutable. */
+	}
+}
+
+void jp_set(JSONPath *jp, json_node *json, json_node *value)
+{
+	ListCell	*lc			= jp_head(jp);
+
+	match_recurse((void*)jp_set_callback, value, lc, mkRefNode(json));
+}
+
 char *jpref_encode(JPRef *ref)
 {
 	switch (ref->type) {