hclwrite: Make block labels a node in their own right

All of the other subdivisions of a block were already nodes, but we'd
represented the labels as an undifferentiated set of nodes belonging
directly to the block's child node list.

Now that we support replacing the labels in the public API, that's a good
excuse to refactor this slightly to make the labels their own node. As
well as being consistent with everything else in Block, this also makes
it easier to implement the Block.SetLabels operation because we can
just  change the children of the labels node, rather than having to
carefully identify and extract the individual child nodes of the block
that happen to represent labels.

Internally this models the labels in a similar sort of way as the content
of a body, although we've kept the public API directly on the Block type
here because that's a more straightforward model for the use-cases we
currently know and matches better with the API of hcl.Block. This is just
an internal change for consistency.

I also added a few tests for having comments interspersed with labels
while I was here, because that helped to better exercise the new
parseBlockLabels function.

5 files changed, 236 insertions, 52 deletions

diff --git a/hclwrite/ast_block.go b/hclwrite/ast_block.go
index 99fdb1d..2bcda5b 100644
--- a/hclwrite/ast_block.go
+++ b/hclwrite/ast_block.go
@@ -10,7 +10,7 @@ type Block struct {
 
 	leadComments *node
 	typeName     *node
-	labels       nodeSet
+	labels       *node
 	open         *node
 	body         *node
 	close        *node
@@ -19,7 +19,6 @@ type Block struct {
 func newBlock() *Block {
 	return &Block{
 		inTree: newInTree(),
-		labels: newNodeSet(),
 	}
 }
 
@@ -35,12 +34,8 @@ func (b *Block) init(typeName string, labels []string) {
 	nameObj := newIdentifier(nameTok)
 	b.leadComments = b.children.Append(newComments(nil))
 	b.typeName = b.children.Append(nameObj)
-	for _, label := range labels {
-		labelToks := TokensForValue(cty.StringVal(label))
-		labelObj := newQuoted(labelToks)
-		labelNode := b.children.Append(labelObj)
-		b.labels.Add(labelNode)
-	}
+	labelsObj := newBlockLabels(labels)
+	b.labels = b.children.Append(labelsObj)
 	b.open = b.children.AppendUnstructuredTokens(Tokens{
 		{
 			Type:  hclsyntax.TokenOBrace,
@@ -88,8 +83,59 @@ func (b *Block) SetType(typeName string) {
 
 // Labels returns the labels of the block.
 func (b *Block) Labels() []string {
-	labelNames := make([]string, 0, len(b.labels))
-	list := b.labels.List()
+	return b.labelsObj().Current()
+}
+
+// SetLabels updates the labels of the block to given labels.
+// Since we cannot assume that old and new labels are equal in length,
+// remove old labels and insert new ones before TokenOBrace.
+func (b *Block) SetLabels(labels []string) {
+	b.labelsObj().Replace(labels)
+}
+
+// labelsObj returns the internal node content representation of the block
+// labels. This is not part of the public API because we're intentionally
+// exposing only a limited API to get/set labels on the block itself in a
+// manner similar to the main hcl.Block type, but our block accessors all
+// use this to get the underlying node content to work with.
+func (b *Block) labelsObj() *blockLabels {
+	return b.labels.content.(*blockLabels)
+}
+
+type blockLabels struct {
+	inTree
+
+	items nodeSet
+}
+
+func newBlockLabels(labels []string) *blockLabels {
+	ret := &blockLabels{
+		inTree: newInTree(),
+		items:  newNodeSet(),
+	}
+
+	ret.Replace(labels)
+	return ret
+}
+
+func (bl *blockLabels) Replace(newLabels []string) {
+	bl.inTree.children.Clear()
+	bl.items.Clear()
+
+	for _, label := range newLabels {
+		labelToks := TokensForValue(cty.StringVal(label))
+		// Force a new label to use the quoted form, which is the idiomatic
+		// form. The unquoted form is supported in HCL 2 only for compatibility
+		// with historical use in HCL 1.
+		labelObj := newQuoted(labelToks)
+		labelNode := bl.children.Append(labelObj)
+		bl.items.Add(labelNode)
+	}
+}
+
+func (bl *blockLabels) Current() []string {
+	labelNames := make([]string, 0, len(bl.items))
+	list := bl.items.List()
 
 	for _, label := range list {
 		switch labelObj := label.content.(type) {
@@ -123,25 +169,3 @@ func (b *Block) Labels() []string {
 
 	return labelNames
 }
-
-// SetLabels updates the labels of the block to given labels.
-// Since we cannot assume that old and new labels are equal in length,
-// remove old labels and insert new ones before TokenOBrace.
-func (b *Block) SetLabels(labels []string) {
-	// Remove old labels
-	for oldLabel := range b.labels {
-		oldLabel.Detach()
-		b.labels.Remove(oldLabel)
-	}
-
-	// Insert new labels before TokenOBrace.
-	for _, label := range labels {
-		labelToks := TokensForValue(cty.StringVal(label))
-		// Force a new label to use the quoted form even if the old one is unquoted.
-		// The unquoted form is supported in HCL 2 only for compatibility with some
-		// historical use in HCL 1.
-		labelObj := newQuoted(labelToks)
-		labelNode := b.children.Insert(b.open, labelObj)
-		b.labels.Add(labelNode)
-	}
-}
diff --git a/hclwrite/ast_block_test.go b/hclwrite/ast_block_test.go
index 5b66fd6..4ca5cb2 100644
--- a/hclwrite/ast_block_test.go
+++ b/hclwrite/ast_block_test.go
@@ -74,6 +74,13 @@ quoted "label1" "label2" {
 		},
 		{
 			`
+quoted "label1" /* foo */ "label2" {
+}
+`,
+			[]string{"label1", "label2"},
+		},
+		{
+			`
 unquoted label1 {
 }
 `,
@@ -81,6 +88,20 @@ unquoted label1 {
 		},
 		{
 			`
+unquoted label1 /* foo */ label2 {
+}
+`,
+			[]string{"label1", "label2"},
+		},
+		{
+			`
+mixed label1 "label2" {
+}
+`,
+			[]string{"label1", "label2"},
+		},
+		{
+			`
 escape "\u0041" {
 }
 `,
@@ -348,7 +369,50 @@ func TestBlockSetLabels(t *testing.T) {
 			},
 		},
 		{
-			`foo hoge {}`,
+			`foo "hoge" /* fuga */ "piyo" {}`,
+			"foo",
+			[]string{"hoge", "piyo"},
+			[]string{"fuga"}, // force quoted form even if the old one is unquoted.
+			Tokens{
+				{
+					Type:         hclsyntax.TokenIdent,
+					Bytes:        []byte(`foo`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenOQuote,
+					Bytes:        []byte(`"`),
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenQuotedLit,
+					Bytes:        []byte(`fuga`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenCQuote,
+					Bytes:        []byte(`"`),
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenOBrace,
+					Bytes:        []byte{'{'},
+					SpacesBefore: 1,
+				},
+				{
+					Type:         hclsyntax.TokenCBrace,
+					Bytes:        []byte{'}'},
+					SpacesBefore: 0,
+				},
+				{
+					Type:         hclsyntax.TokenEOF,
+					Bytes:        []byte{},
+					SpacesBefore: 0,
+				},
+			},
+		},
+		{
+			`foo "hoge" /* foo */  "" {}`,
 			"foo",
 			[]string{"hoge"},
 			[]string{"fuga"}, // force quoted form even if the old one is unquoted.
diff --git a/hclwrite/node.go b/hclwrite/node.go
index 377490a..d3a5b72 100644
--- a/hclwrite/node.go
+++ b/hclwrite/node.go
@@ -207,6 +207,12 @@ func (ns nodeSet) Remove(n *node) {
 	delete(ns, n)
 }
 
+func (ns nodeSet) Clear() {
+	for n := range ns {
+		delete(ns, n)
+	}
+}
+
 func (ns nodeSet) List() []*node {
 	if len(ns) == 0 {
 		return nil
diff --git a/hclwrite/parser.go b/hclwrite/parser.go
index ad5fcda..3df5144 100644
--- a/hclwrite/parser.go
+++ b/hclwrite/parser.go
@@ -289,7 +289,6 @@ func parseAttribute(nativeAttr *hclsyntax.Attribute, from, leadComments, lineCom
 func parseBlock(nativeBlock *hclsyntax.Block, from, leadComments, lineComments, newline inputTokens) *node {
 	block := &Block{
 		inTree: newInTree(),
-		labels: newNodeSet(),
 	}
 	children := block.inTree.children
 
@@ -312,20 +311,9 @@ func parseBlock(nativeBlock *hclsyntax.Block, from, leadComments, lineComments,
 		children.AppendNode(in)
 	}
 
-	for _, rng := range nativeBlock.LabelRanges {
-		var labelTokens inputTokens
-		before, labelTokens, from = from.Partition(rng)
-		children.AppendUnstructuredTokens(before.Tokens())
-		tokens := labelTokens.Tokens()
-		var ln *node
-		if len(tokens) == 1 && tokens[0].Type == hclsyntax.TokenIdent {
-			ln = newNode(newIdentifier(tokens[0]))
-		} else {
-			ln = newNode(newQuoted(tokens))
-		}
-		block.labels.Add(ln)
-		children.AppendNode(ln)
-	}
+	before, labelsNode, from := parseBlockLabels(nativeBlock, from)
+	block.labels = labelsNode
+	children.AppendNode(labelsNode)
 
 	before, oBrace, from := from.Partition(nativeBlock.OpenBraceRange)
 	children.AppendUnstructuredTokens(before.Tokens())
@@ -356,6 +344,34 @@ func parseBlock(nativeBlock *hclsyntax.Block, from, leadComments, lineComments,
 	return newNode(block)
 }
 
+func parseBlockLabels(nativeBlock *hclsyntax.Block, from inputTokens) (inputTokens, *node, inputTokens) {
+	labelsObj := newBlockLabels(nil)
+	children := labelsObj.children
+
+	var beforeAll inputTokens
+	for i, rng := range nativeBlock.LabelRanges {
+		var before, labelTokens inputTokens
+		before, labelTokens, from = from.Partition(rng)
+		if i == 0 {
+			beforeAll = before
+		} else {
+			children.AppendUnstructuredTokens(before.Tokens())
+		}
+		tokens := labelTokens.Tokens()
+		var ln *node
+		if len(tokens) == 1 && tokens[0].Type == hclsyntax.TokenIdent {
+			ln = newNode(newIdentifier(tokens[0]))
+		} else {
+			ln = newNode(newQuoted(tokens))
+		}
+		labelsObj.items.Add(ln)
+		children.AppendNode(ln)
+	}
+
+	after := from
+	return beforeAll, newNode(labelsObj), after
+}
+
 func parseExpression(nativeExpr hclsyntax.Expression, from inputTokens) *node {
 	expr := newExpression()
 	children := expr.inTree.children
diff --git a/hclwrite/parser_test.go b/hclwrite/parser_test.go
index 9a4339f..ed17303 100644
--- a/hclwrite/parser_test.go
+++ b/hclwrite/parser_test.go
@@ -205,6 +205,9 @@ func TestParse(t *testing.T) {
 								Val:  "b",
 							},
 							{
+								Type: "blockLabels",
+							},
+							{
 								Type: "Tokens",
 								Val:  " {",
 							},
@@ -240,8 +243,13 @@ func TestParse(t *testing.T) {
 								Val:  "b",
 							},
 							{
-								Type: "identifier",
-								Val:  ` label`,
+								Type: "blockLabels",
+								Children: []TestTreeNode{
+									{
+										Type: "identifier",
+										Val:  " label",
+									},
+								},
 							},
 							{
 								Type: "Tokens",
@@ -279,8 +287,71 @@ func TestParse(t *testing.T) {
 								Val:  "b",
 							},
 							{
-								Type: "quoted",
-								Val:  ` "label"`,
+								Type: "blockLabels",
+								Children: []TestTreeNode{
+									{
+										Type: "quoted",
+										Val:  ` "label"`,
+									},
+								},
+							},
+							{
+								Type: "Tokens",
+								Val:  " {",
+							},
+							{
+								Type: "Body",
+							},
+							{
+								Type: "Tokens",
+								Val:  "}",
+							},
+							{
+								Type: "Tokens",
+								Val:  "\n",
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			"b \"label1\" /* foo */ \"label2\" {}\n",
+			TestTreeNode{
+				Type: "Body",
+				Children: []TestTreeNode{
+					{
+						Type: "Block",
+						Children: []TestTreeNode{
+							{
+								Type: "comments",
+							},
+							{
+								Type: "identifier",
+								Val:  "b",
+							},
+							{
+								Type: "blockLabels",
+								Children: []TestTreeNode{
+									{
+										Type: "quoted",
+										Val:  ` "label1"`,
+									},
+									{
+										// The comment between the labels just
+										// becomes an "unstructured tokens"
+										// node, because this isn't a place
+										// where we expect comments to attach
+										// to a particular object as
+										// documentation.
+										Type: "Tokens",
+										Val:  ` /* foo */`,
+									},
+									{
+										Type: "quoted",
+										Val:  ` "label2"`,
+									},
+								},
 							},
 							{
 								Type: "Tokens",
@@ -318,6 +389,9 @@ func TestParse(t *testing.T) {
 								Val:  "b",
 							},
 							{
+								Type: "blockLabels",
+							},
+							{
 								Type: "Tokens",
 								Val:  " {",
 							},