From ae4c0d878e559ba741c7e684a78ffe7271c6f10b Mon Sep 17 00:00:00 2001
From: Sven Rebhan <36194019+srebhan@users.noreply.github.com>
Date: Thu, 3 Aug 2023 15:20:24 +0200
Subject: [PATCH] fix: Handle maps correctly (#4)

---
 node.go | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/node.go b/node.go
index c781f46..cbcacdd 100644
--- a/node.go
+++ b/node.go
@@ -138,9 +138,9 @@ func visit(parent *Node, msg protoreflect.Message, level int) {
 func traverse(parent *Node, field protoreflect.FieldDescriptor, value protoreflect.Value, level int) {
 	node := &Node{Type: ElementNode, Name: string(field.Name()), level: level}
 	nodeChildren := 0
-	if field.IsList() {
+	switch {
+	case field.IsList():
 		l := value.List()
-
 		for i := 0; i < l.Len(); i++ {
 			subNode := handleValue(field.Kind(), l.Get(i), level+1)
 			if subNode.Type == ElementNode {
@@ -156,7 +156,25 @@ func traverse(parent *Node, field protoreflect.FieldDescriptor, value protorefle
 				nodeChildren++
 			}
 		}
-	} else {
+	case field.IsMap():
+		key := field.MapKey()
+		value.Map().Range(func(k protoreflect.MapKey, v protoreflect.Value) bool {
+			subNode := handleValue(key.Kind(), v, level+1)
+			if subNode.Type == ElementNode {
+				// Add element nodes directly to the parent
+				subNode.Name = k.String()
+				addNode(parent, subNode)
+			} else {
+				// Add basic nodes to the local collection node
+				elementNode := &Node{Type: ElementNode, Name: k.String(), level: level + 1}
+				subNode.level += 2
+				addNode(elementNode, subNode)
+				addNode(node, elementNode)
+				nodeChildren++
+			}
+			return true
+		})
+	default:
 		subNode := handleValue(field.Kind(), value, level+1)
 		if subNode.Type == ElementNode {
 			// Add element nodes directly to the parent