From 4d6547d09926aef30d3350ce63b7d76d9fb560e4 Mon Sep 17 00:00:00 2001
From: Josh Humphries <2035234+jhump@users.noreply.github.com>
Date: Wed, 6 Dec 2023 07:15:40 -0500
Subject: [PATCH] new protomessage package; move As from protoresolve to
 protomessage... still debating if this is the right organization...

---
 {protoresolve => protomessage}/as.go      |  6 ++-
 {protoresolve => protomessage}/as_test.go |  2 +-
 protomessage/walk.go                      | 54 +++++++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)
 rename {protoresolve => protomessage}/as.go (92%)
 rename {protoresolve => protomessage}/as_test.go (98%)
 create mode 100644 protomessage/walk.go

diff --git a/protoresolve/as.go b/protomessage/as.go
similarity index 92%
rename from protoresolve/as.go
rename to protomessage/as.go
index 9dc6e345..5fc97ff6 100644
--- a/protoresolve/as.go
+++ b/protomessage/as.go
@@ -1,4 +1,4 @@
-package protoresolve
+package protomessage
 
 import (
 	"fmt"
@@ -6,6 +6,8 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/reflect/protoregistry"
+
+	"github.com/jhump/protoreflect/v2/protoresolve"
 )
 
 // PointerMessage is a pointer type that implements [proto.Message].
@@ -34,7 +36,7 @@ func As[M PointerMessage[T], T any](msg proto.Message) (M, error) {
 			if exts == nil {
 				exts = &protoregistry.Types{}
 			}
-			err = exts.RegisterExtension(ExtensionType(fd))
+			err = exts.RegisterExtension(protoresolve.ExtensionType(fd))
 			return err == nil
 		}
 		return true
diff --git a/protoresolve/as_test.go b/protomessage/as_test.go
similarity index 98%
rename from protoresolve/as_test.go
rename to protomessage/as_test.go
index 5e69466d..45be201e 100644
--- a/protoresolve/as_test.go
+++ b/protomessage/as_test.go
@@ -1,4 +1,4 @@
-package protoresolve
+package protomessage
 
 import (
 	"testing"
diff --git a/protomessage/walk.go b/protomessage/walk.go
new file mode 100644
index 00000000..b126053b
--- /dev/null
+++ b/protomessage/walk.go
@@ -0,0 +1,54 @@
+package protomessage
+
+import "google.golang.org/protobuf/reflect/protoreflect"
+
+// Walk traverses the given root messages, iterating through its fields and
+// through all values in maps and lists, calling the given action for all
+// message values encountered. The given action is called for root first
+// before being called for any contained message values.
+//
+// The path provided to the callback is the sequence of field numbers,
+// map keys, and list indices that identifies the location of the given
+// message. It is empty when called for the root message.
+//
+// If the callback returns false, the traversal is terminated and the
+// callback will not be invoked again.
+func Walk(root protoreflect.Message, action func(path []any, val protoreflect.Message) bool) {
+	walk(root, nil, action)
+}
+
+func walk(root protoreflect.Message, path []any, action func(path []any, val protoreflect.Message) bool) bool {
+	ok := action(path, root)
+	root.Range(func(field protoreflect.FieldDescriptor, val protoreflect.Value) bool {
+		path = append(path, field.Number())
+		switch {
+		case field.IsMap() && isMessageKind(field.MapValue().Kind()):
+			mapVal := val.Map()
+			mapVal.Range(func(key protoreflect.MapKey, val protoreflect.Value) bool {
+				path = append(path, key, protoreflect.FieldNumber(2) /* field 2 is the value in an entry */)
+				ok = walk(val.Message(), path, action)
+				path = path[:len(path)-2]
+				return ok
+			})
+		case field.IsList() && isMessageKind(field.Kind()):
+			listVal := val.List()
+			for i, length := 0, listVal.Len(); i < length; i++ {
+				path = append(path, i)
+				ok = walk(listVal.Get(i).Message(), path, action)
+				path = path[:len(path)-1]
+				if !ok {
+					break
+				}
+			}
+		case isMessageKind(field.Kind()):
+			ok = walk(val.Message(), path, action)
+		}
+		path = path[:len(path)-1] // pop field number
+		return ok
+	})
+	return ok
+}
+
+func isMessageKind(k protoreflect.Kind) bool {
+	return k == protoreflect.MessageKind || k == protoreflect.GroupKind
+}