hcl/ast: ObjectList.Prefix

hcl/ast/ast.go

@@ -43,6 +43,34 @@ func (o *ObjectList) Add(item *ObjectItem) {
 	o.Items = append(o.Items, item)
 }
 
+func (o *ObjectList) Prefix(keys ...string) *ObjectList {
+	var result ObjectList
+	for _, item := range o.Items {
+		// If there aren't enough keys, then ignore this
+		if len(item.Keys) < len(keys)+1 {
+			continue
+		}
+
+		match := true
+		for i, key := range item.Keys[:len(keys)] {
+			if key.Token.Text != keys[i] {
+				match = false
+				break
+			}
+		}
+		if !match {
+			continue
+		}
+
+		// Strip off the prefix from the children
+		newItem := *item
+		newItem.Keys = newItem.Keys[len(keys):]
+		result.Add(&newItem)
+	}
+
+	return &result
+}
+
 func (o *ObjectList) Pos() token.Pos {
 	// always returns the uninitiliazed position
 	return o.Items[0].Pos()

hcl/ast/ast_test.go

@@ -0,0 +1,62 @@
+package ast
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/hashicorp/hcl/hcl/token"
+)
+
+func TestObjectListPrefix(t *testing.T) {
+	var cases = []struct {
+		Prefix []string
+		Input  []*ObjectItem
+		Output []*ObjectItem
+	}{
+		{
+			[]string{"foo"},
+			[]*ObjectItem{
+				&ObjectItem{
+					Keys: []*ObjectKey{
+						&ObjectKey{
+							Token: token.Token{Type: token.STRING, Text: `foo`},
+						},
+					},
+				},
+			},
+			nil,
+		},
+
+		{
+			[]string{"foo"},
+			[]*ObjectItem{
+				&ObjectItem{
+					Keys: []*ObjectKey{
+						&ObjectKey{Token: token.Token{Type: token.STRING, Text: `foo`}},
+						&ObjectKey{Token: token.Token{Type: token.STRING, Text: `bar`}},
+					},
+				},
+				&ObjectItem{
+					Keys: []*ObjectKey{
+						&ObjectKey{Token: token.Token{Type: token.STRING, Text: `baz`}},
+					},
+				},
+			},
+			[]*ObjectItem{
+				&ObjectItem{
+					Keys: []*ObjectKey{
+						&ObjectKey{Token: token.Token{Type: token.STRING, Text: `bar`}},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		input := &ObjectList{Items: tc.Input}
+		expected := &ObjectList{Items: tc.Output}
+		if actual := input.Prefix(tc.Prefix...); !reflect.DeepEqual(actual, expected) {
+			t.Fatalf("in order: input, expected, actual\n\n%#v\n\n%#v\n\n%#v", input, expected, actual)
+		}
+	}
+}