From cabbea0d9752b9ee5de2ec900b4caef01ab1e07a Mon Sep 17 00:00:00 2001
From: Mengqi Yu <mengqiy@google.com>
Date: Mon, 7 Jun 2021 11:29:31 -0700
Subject: [PATCH] handle comments in list correctly

---
 kyaml/comments/comments.go      |  5 +++-
 kyaml/comments/comments_test.go | 42 +++++++++++++++++++++++++++++++++
 kyaml/kio/byteio_writer.go      |  8 +++++--
 kyaml/yaml/serialization.go     | 39 ++++--------------------------
 4 files changed, 56 insertions(+), 38 deletions(-)

diff --git a/kyaml/comments/comments.go b/kyaml/comments/comments.go
index 9d6dc22dc..d67517e99 100644
--- a/kyaml/comments/comments.go
+++ b/kyaml/comments/comments.go
@@ -54,7 +54,10 @@ func (c *copier) VisitList(s walk.Sources, _ *openapi.ResourceSchema, _ walk.Lis
 		origin := originItems[i]
 
 		if dest.Value == origin.Value {
-			copy(yaml.NewRNode(dest), yaml.NewRNode(origin))
+			// We should do it recursively on each node in the list.
+			if err := CopyComments(yaml.NewRNode(dest), yaml.NewRNode(origin)); err != nil {
+				return nil, err
+			}
 		}
 	}
 
diff --git a/kyaml/comments/comments_test.go b/kyaml/comments/comments_test.go
index 5238e8754..e96047417 100644
--- a/kyaml/comments/comments_test.go
+++ b/kyaml/comments/comments_test.go
@@ -85,6 +85,48 @@ spec:
 `,
 		},
 
+		{
+			name: "associative_list_2",
+			from: `
+apiVersion: constraints.gatekeeper.sh/v1beta1
+kind: EnforceFoo
+metadata:
+  name: enforce-foo
+spec:
+  parameters:
+    naming_rules:
+    - kind: Bar
+      patterns:
+      # comment 1
+      - ^(dev|prod|staging|qa|shared)$
+`,
+			to: `
+apiVersion: constraints.gatekeeper.sh/v1beta1
+kind: EnforceFoo
+metadata:
+  name: enforce-foo
+spec:
+  parameters:
+    naming_rules:
+    - kind: Bar
+      patterns:
+      - ^(dev|prod|staging|qa|shared)$
+`,
+			expected: `
+apiVersion: constraints.gatekeeper.sh/v1beta1
+kind: EnforceFoo
+metadata:
+  name: enforce-foo
+spec:
+  parameters:
+    naming_rules:
+    - kind: Bar
+      patterns:
+      # comment 1
+      - ^(dev|prod|staging|qa|shared)$
+`,
+		},
+
 		{
 			name: "keep_comments",
 			from: `# A
diff --git a/kyaml/kio/byteio_writer.go b/kyaml/kio/byteio_writer.go
index cd3da26de..a0ecdb381 100644
--- a/kyaml/kio/byteio_writer.go
+++ b/kyaml/kio/byteio_writer.go
@@ -50,7 +50,12 @@ type ByteWriter struct {
 
 var _ Writer = ByteWriter{}
 
-func (w ByteWriter) Write(nodes []*yaml.RNode) error {
+func (w ByteWriter) Write(inputNodes []*yaml.RNode) error {
+	// Copy the nodes to prevent writer from mutating the original nodes.
+	var nodes []*yaml.RNode
+	for i := range inputNodes {
+		nodes = append(nodes, inputNodes[i].Copy())
+	}
 	yaml.DoSerializationHacksOnNodes(nodes)
 	if w.Sort {
 		if err := kioutil.SortNodes(nodes); err != nil {
@@ -131,7 +136,6 @@ func (w ByteWriter) Write(nodes []*yaml.RNode) error {
 		items.Content = append(items.Content, nodes[i].YNode())
 	}
 	err := encoder.Encode(doc)
-	yaml.UndoSerializationHacksOnNodes(nodes)
 	return err
 }
 
diff --git a/kyaml/yaml/serialization.go b/kyaml/yaml/serialization.go
index 92510c518..ff47157df 100644
--- a/kyaml/yaml/serialization.go
+++ b/kyaml/yaml/serialization.go
@@ -34,43 +34,12 @@ func DoSerializationHacks(node *yaml.Node) {
 			// https://github.com/go-yaml/yaml/issues/587 in go-yaml.v3
 			// Remove this hack when the issue has been resolved
 			if len(node.Content) > 0 && node.Content[0].Kind == ScalarNode {
-				node.HeadComment = node.Content[0].HeadComment
+				// Don't clobber the head comment if it's not empty.
+				if node.HeadComment == "" && node.Content[0].HeadComment != "" {
+					node.HeadComment = node.Content[0].HeadComment
+				}
 				node.Content[0].HeadComment = ""
 			}
 		}
 	}
 }
-
-func UndoSerializationHacksOnNodes(nodes []*RNode) {
-	for _, node := range nodes {
-		UndoSerializationHacks(node.YNode())
-	}
-}
-
-// UndoSerializationHacks reverts the changes made by DoSerializationHacks
-// Refer to https://github.com/go-yaml/yaml/issues/587 for more details
-func UndoSerializationHacks(node *yaml.Node) {
-	switch node.Kind {
-	case DocumentNode:
-		for _, node := range node.Content {
-			DoSerializationHacks(node)
-		}
-
-	case MappingNode:
-		for _, node := range node.Content {
-			DoSerializationHacks(node)
-		}
-
-	case SequenceNode:
-		for _, node := range node.Content {
-			// revert the changes made in DoSerializationHacks
-			// This is necessary to address serialization issue
-			// https://github.com/go-yaml/yaml/issues/587 in go-yaml.v3
-			// Remove this hack when the issue has been resolved
-			if len(node.Content) > 0 && node.Content[0].Kind == ScalarNode {
-				node.Content[0].HeadComment = node.HeadComment
-				node.HeadComment = ""
-			}
-		}
-	}
-}