Augment path search

pkg/schema/expand.go

@@ -41,7 +41,7 @@ func (sc *Schema) ExpandPath(p *sdcpb.Path, dt sdcpb.DataType) ([]*sdcpb.Path, e
 	for _, k := range strings.Fields(e.Key) {
 		keys[k] = struct{}{}
 	}
-	for _, c := range e.Dir {
+	for _, c := range getChildren(e) {
 		// skip keys
 		if _, ok := keys[c.Name]; ok {
 			continue
@@ -109,7 +109,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 			kmap[k] = struct{}{}
 		}
 
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			if _, ok := kmap[c.Name]; ok {
 				continue
 			}
@@ -126,7 +126,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 	case e.IsContainer():
 		log.Debugf("got container: %s", e.Name)
 		containerPE := &sdcpb.PathElem{Name: e.Name, Key: make(map[string]string)}
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			log.Debugf("container parent adding child: %s", c.Name)
 			childrenPE := sc.getPathElems(c, dt)
 

pkg/schema/object.go

@@ -120,7 +120,7 @@ func (sc *Schema) FindPossibleModulesForPathElement(e *yang.Entry, pathElement s
 	case e.Node == nil:
 		entries := make([]*yang.Entry, 0)
 		for _, entry := range sc.root.Dir {
-			if ee, ok := entry.Dir[path]; ok && (!foundPrefix || ee.Prefix.Name == prefix) {
+			if ee, ok := entry.Dir[path]; ok && (!foundPrefix || ee.Prefix.Name == prefix || entry.Name == prefix) {
 				entries = append(entries, entry)
 			}
 		}
@@ -190,8 +190,7 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 
 			// prefix will be in [0] if exists, so path will always be in last index
 			// compare name without prefix
-			pathElements := strings.SplitN(pe[0], ":", 2)
-			if ee.Name != pathElements[len(pathElements)-1] {
+			if ee.Name != pathElemLocalName(pe[0]) {
 				continue
 			}
 			return getEntry(ee, pe[1:])
@@ -201,6 +200,22 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 	}
 }
 
+// pathElemLocalName returns the YANG identifier after an optional single "prefix:".
+// It matches the name comparison rule used in getEntry for prefixed path elements.
+func pathElemLocalName(s string) string {
+	parts := strings.SplitN(s, ":", 2)
+	return parts[len(parts)-1]
+}
+
+// entryChildByName returns a direct child of parent by name from Dir.
+// Goyang merges augment children into Dir; Augmented holds shallow augment-block copies only.
+func entryChildByName(parent *yang.Entry, name string) *yang.Entry {
+	if parent == nil {
+		return nil
+	}
+	return parent.Dir[name]
+}
+
 func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 	if len(pe) == 0 {
 		return nil
@@ -221,7 +236,7 @@ func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 		if e == nil {
 			return fmt.Errorf("module %q not found", first)
 		}
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pathElemLocalName(pe[0])); ee != nil {
 			err := sc.buildPath(pe, p, ee)
 			if err != nil {
 				return err
@@ -234,7 +249,7 @@ func (sc *Schema) BuildPath(pe []string, p *sdcpb.Path) error {
 	}
 	// try children
 	for _, e := range sc.root.Dir {
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pathElemLocalName(pe[0])); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 	}
@@ -281,7 +296,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 			return nil
 		}
 		nxt := pe[count]
-		if ee, ok := e.Dir[nxt]; ok {
+		if ee := entryChildByName(e, nxt); ee != nil {
 			return sc.buildPath(pe[count:], p, ee)
 		}
 		// find choices/cases
@@ -291,45 +306,56 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		}
 		return sc.buildPath(pe[count:], p, ee)
 	case e.IsChoice():
-		p.Elem = append(p.Elem, cpe)
-		for _, entry := range e.Dir {
-			if entry.IsCase() {
-				if ee, ok := entry.Dir[pe[0]]; ok {
-					return sc.buildPath(pe[1:], p, ee)
+		{
+			tryChoiceChild := func(choice *yang.Entry, entry *yang.Entry) (handled bool, err error) {
+				if entry == nil {
+					return false, nil
 				}
-			} else {
-				if ee, ok := e.Dir[pe[0]]; ok {
-					return sc.buildPath(pe[1:], p, ee)
+				if entry.IsCase() {
+					if ee := entryChildByName(entry, pe[0]); ee != nil {
+						return true, sc.buildPath(pe[1:], p, ee)
+					}
+				} else {
+					if ee := entryChildByName(choice, pe[0]); ee != nil {
+						return true, sc.buildPath(pe[1:], p, ee)
+					}
 				}
+				return false, nil
 			}
+			p.Elem = append(p.Elem, cpe)
+			for _, entry := range e.Dir {
+				if handled, err := tryChoiceChild(e, entry); handled {
+					return err
+				}
+			}
+			return fmt.Errorf("choice %s - unknown element %s", e.Name, pe[0])
 		}
-		return fmt.Errorf("choice %s - unknown element %s", e.Name, pe[0])
 	case e.IsCase():
 		// RFC7950 7.9.2: A case node does not exist in the data tree.
 		// p.Elem = append(p.Elem, cpe)
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
-		if ee, ok := e.Dir[e.Name]; ok {
+		if ee := entryChildByName(e, e.Name); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		return fmt.Errorf("case %s - unknown element %s", e.Name, pe[0])
 	case e.IsContainer():
 		// implicit case: child with same name which is a choice
-		if ee, ok := e.Dir[pe[0]]; ee != nil && ok {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			if ee.IsChoice() {
 				return sc.buildPath(pe[1:], p, ee)
 			}
 		}
 
 		p.Elem = append(p.Elem, cpe)
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := entryChildByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		if lpe == 1 {
 			return nil
 		}
-		if ee, ok := e.Dir[pe[1]]; ok {
+		if ee := entryChildByName(e, pe[1]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
 		// find choice/case
@@ -562,21 +588,36 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 	if len(pe) == 0 {
 		return e, nil
 	}
-	for _, ee := range e.Dir {
-		if !ee.IsChoice() {
-			continue
-		}
-		if eee, ok := ee.Dir[pe[1]]; ok && !eee.IsCase() {
-			return eee, nil
+	// pe is expected to contain at least the current element name at index 0 and
+	// the sought child element name at index 1.
+	//
+	// This is used from container/list resolution paths where the next element may
+	// live under a choice/case. Case nodes do not exist in the data tree, so we
+	// must search through cases to find the actual schema node.
+	if len(pe) < 2 {
+		return nil, fmt.Errorf("unknown element %s", pe[0])
+	}
+
+	// scan choice nodes under e (goyang merges augmented choices into Dir)
+	choices := make([]*yang.Entry, 0)
+	for _, child := range e.Dir {
+		if child != nil && child.IsChoice() {
+			choices = append(choices, child)
 		}
-		// assume there was a case obj,
-		// search one step deeper
-		for _, eee := range ee.Dir {
-			if !eee.IsCase() {
+	}
+
+	for _, choice := range choices {
+		// implicit case: choice directly contains the data node
+		if direct := entryChildByName(choice, pe[1]); direct != nil && !direct.IsCase() {
+			return direct, nil
+		}
+		// explicit cases: the data node is under a case
+		for _, cc := range choice.Dir {
+			if cc == nil || !cc.IsCase() {
 				continue
 			}
-			if eeee, ok := eee.Dir[pe[1]]; ok {
-				return eeee, nil
+			if target := entryChildByName(cc, pe[1]); target != nil {
+				return target, nil
 			}
 		}
 	}

pkg/schema/object_test.go

@@ -15,8 +15,10 @@
 package schema
 
 import (
+	"path/filepath"
 	"reflect"
 	"sort"
+	"sync"
 	"testing"
 
 	"github.com/openconfig/goyang/pkg/yang"
@@ -476,6 +478,84 @@ func TestSchema_BuildPath(t *testing.T) {
 	}
 }
 
+func TestSchema_BuildPath_AugmentedUnderCase(t *testing.T) {
+	td := filepath.Join("testdata", "buildpath-augment")
+	sc, err := NewSchema(&config.SchemaConfig{
+		Name:    "bpcc",
+		Vendor:  "test",
+		Version: "0",
+		Files: []string{
+			filepath.Join(td, "bpcc-base.yang"),
+			filepath.Join(td, "bpcc-aug.yang"),
+		},
+	})
+	if err != nil {
+		t.Fatalf("NewSchema: %v", err)
+	}
+	p := &sdcpb.Path{}
+	err = sc.BuildPath([]string{"bpcc-base:top", "augment-leaf"}, p)
+	if err != nil {
+		t.Fatalf("BuildPath: %v", err)
+	}
+	want := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "bpcc-base:top"},
+		{Name: "augment-leaf"},
+	}}
+	if !comparePaths(p, want) {
+		t.Fatalf("got %v want %v", p, want)
+	}
+
+	p2 := &sdcpb.Path{}
+	if err := sc.BuildPath([]string{"bpcc-base:top", "leaf-a"}, p2); err != nil {
+		t.Fatalf("BuildPath native leaf under case: %v", err)
+	}
+	want2 := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "bpcc-base:top"},
+		{Name: "leaf-a"},
+	}}
+	if !comparePaths(p2, want2) {
+		t.Fatalf("native choice path got %v want %v", p2, want2)
+	}
+}
+
+// TestSchema_BuildPath_bootstrapModuleChild verifies BuildPath resolves the first
+// in-module segment using Dir lookup on the module entry (goyang merges augments into Dir).
+func TestSchema_BuildPath_bootstrapModuleChild(t *testing.T) {
+	augLeaf := &yang.Entry{
+		Name:   "aug-only",
+		Kind:   yang.LeafEntry,
+		Type:   &yang.YangType{Kind: yang.Ystring},
+		Prefix: &yang.Value{Name: "pfx"},
+	}
+	mod := &yang.Entry{
+		Name:   "testmod",
+		Kind:   yang.DirectoryEntry,
+		Dir:    map[string]*yang.Entry{"aug-only": augLeaf},
+		Prefix: &yang.Value{Name: "pfx"},
+	}
+	augLeaf.Parent = mod
+	root := &yang.Entry{
+		Name: RootName,
+		Kind: yang.DirectoryEntry,
+		Dir:  map[string]*yang.Entry{"testmod": mod},
+	}
+	sc := &Schema{
+		root:   root,
+		config: &config.SchemaConfig{Name: "t", Vendor: "t", Version: "0"},
+		m:      new(sync.RWMutex),
+	}
+	p := &sdcpb.Path{}
+	if err := sc.BuildPath([]string{"testmod:aug-only"}, p); err != nil {
+		t.Fatalf("BuildPath: %v", err)
+	}
+	want := &sdcpb.Path{Elem: []*sdcpb.PathElem{
+		{Name: "testmod:aug-only"},
+	}}
+	if !comparePaths(p, want) {
+		t.Fatalf("got %v want %v", p, want)
+	}
+}
+
 func comparePathElem(pe1, pe2 *sdcpb.PathElem) bool {
 	if pe1 == nil {
 		return pe2 == nil

pkg/schema/schema.go

@@ -147,8 +147,8 @@ func (s *Schema) Walk(e *yang.Entry, fn func(ec *yang.Entry) error) error {
 	if err != nil {
 		return err
 	}
-	for _, e := range e.Dir {
-		err = s.Walk(e, fn)
+	for _, ce := range e.Dir {
+		err = s.Walk(ce, fn)
 		if err != nil {
 			return err
 		}

pkg/schema/testdata/buildpath-augment/bpcc-aug.yang

@@ -0,0 +1,18 @@
+module bpcc-aug {
+  yang-version 1.1;
+  namespace "urn:bpcc:aug";
+  prefix bcx;
+
+  import bpcc-base {
+    prefix bcb;
+  }
+
+  organization "schema-server test";
+  description "Augment-only leaf under case-a";
+
+  augment "/bcb:top/bcb:variant/bcb:case-a" {
+    leaf augment-leaf {
+      type string;
+    }
+  }
+}

pkg/schema/testdata/buildpath-augment/bpcc-base.yang

@@ -0,0 +1,23 @@
+module bpcc-base {
+  yang-version 1.1;
+  namespace "urn:bpcc:base";
+  prefix bcb;
+
+  organization "schema-server test";
+  description "Base module: choice/case for BuildPath augment tests";
+
+  container top {
+    choice variant {
+      case case-a {
+        leaf leaf-a {
+          type string;
+        }
+      }
+      case case-b {
+        leaf leaf-b {
+          type string;
+        }
+      }
+    }
+  }
+}