hclwrite: Expression.RenameVariablePrefix

This function serves the use-case of surgically renaming references to
variables in an existing file without disturbing other tokens.

hclwrite/ast_body.go

@@ -41,6 +41,31 @@ func (b *Body) AppendUnstructuredTokens(ts Tokens) {
 	b.inTree.children.Append(ts)
 }
 
+// Attributes returns a new map of all of the attributes in the body, with
+// the attribute names as the keys.
+func (b *Body) Attributes() map[string]*Attribute {
+	ret := make(map[string]*Attribute)
+	for n := range b.items {
+		if attr, isAttr := n.content.(*Attribute); isAttr {
+			nameObj := attr.name.content.(*identifier)
+			name := string(nameObj.token.Bytes)
+			ret[name] = attr
+		}
+	}
+	return ret
+}
+
+// Blocks returns a new slice of all the blocks in the body.
+func (b *Body) Blocks() []*Block {
+	ret := make([]*Block, 0, len(b.items))
+	for n := range b.items {
+		if block, isBlock := n.content.(*Block); isBlock {
+			ret = append(ret, block)
+		}
+	}
+	return ret
+}
+
 // GetAttribute returns the attribute from the body that has the given name,
 // or returns nil if there is currently no matching attribute.
 func (b *Body) GetAttribute(name string) *Attribute {

hclwrite/ast_expression.go

@@ -1,6 +1,8 @@
 package hclwrite
 
 import (
+	"fmt"
+
 	"github.com/hashicorp/hcl2/hcl"
 	"github.com/hashicorp/hcl2/hcl/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
@@ -99,6 +101,68 @@ func NewExpressionAbsTraversal(traversal hcl.Traversal) *Expression {
 	return expr
 }
 
+// Variables returns the absolute traversals that exist within the receiving
+// expression.
+func (e *Expression) Variables() []*Traversal {
+	nodes := e.absTraversals.List()
+	ret := make([]*Traversal, len(nodes))
+	for i, node := range nodes {
+		ret[i] = node.content.(*Traversal)
+	}
+	return ret
+}
+
+// RenameVariablePrefix examines each of the absolute traversals in the
+// receiving expression to see if they have the given sequence of names as
+// a prefix prefix. If so, they are updated in place to have the given
+// replacement names instead of that prefix.
+//
+// This can be used to implement symbol renaming. The calling application can
+// visit all relevant expressions in its input and apply the same renaming
+// to implement a global symbol rename.
+//
+// The search and replacement traversals must be the same length, or this
+// method will panic. Only attribute access operations can be matched and
+// replaced. Index steps never match the prefix.
+func (e *Expression) RenameVariablePrefix(search, replacement []string) {
+	if len(search) != len(replacement) {
+		panic(fmt.Sprintf("search and replacement length mismatch (%d and %d)", len(search), len(replacement)))
+	}
+Traversals:
+	for node := range e.absTraversals {
+		traversal := node.content.(*Traversal)
+		if len(traversal.steps) < len(search) {
+			// If it's shorter then it can't have our prefix
+			continue
+		}
+
+		stepNodes := traversal.steps.List()
+		for i, name := range search {
+			step, isName := stepNodes[i].content.(*TraverseName)
+			if !isName {
+				continue Traversals // only name nodes can match
+			}
+			foundNameBytes := step.name.content.(*identifier).token.Bytes
+			if len(foundNameBytes) != len(name) {
+				continue Traversals
+			}
+			if string(foundNameBytes) != name {
+				continue Traversals
+			}
+		}
+
+		// If we get here then the prefix matched, so now we'll swap in
+		// the replacement strings.
+		for i, name := range replacement {
+			step := stepNodes[i].content.(*TraverseName)
+			token := step.name.content.(*identifier).token
+			token.Bytes = []byte(name)
+		}
+	}
+}
+
+// Traversal represents a sequence of variable, attribute, and/or index
+// operations.
 type Traversal struct {
 	inTree
 

hclwrite/examples_test.go

@@ -66,3 +66,28 @@ func Example_generateFromScratch() {
 	//   }
 	// }
 }
+
+func ExampleExpression_RenameVariablePrefix() {
+	src := []byte(
+		"foo = a.x + a.y * b.c\n" +
+			"bar = max(a.z, b.c)\n",
+	)
+	f, diags := hclwrite.ParseConfig(src, "", hcl.Pos{Line: 1, Column: 1})
+	if diags.HasErrors() {
+		fmt.Printf("errors: %s", diags)
+		return
+	}
+
+	// Rename references of variable "a" to "z"
+	for _, attr := range f.Body().Attributes() {
+		attr.Expr().RenameVariablePrefix(
+			[]string{"a"},
+			[]string{"z"},
+		)
+	}
+
+	fmt.Printf("%s", f.Bytes())
+	// Output:
+	// foo = z.x + z.y * b.c
+	// bar = max(z.z, b.c)
+}

hclwrite/parser.go

@@ -372,6 +372,7 @@ func parseTraversal(nativeTraversal hcl.Traversal, from inputTokens) (before inp
 		before, step, after := parseTraversalStep(nativeStep, stepAfter)
 		children.AppendUnstructuredTokens(before.Tokens())
 		children.AppendNode(step)
+		traversal.steps.Add(step)
 		stepAfter = after
 	}