From 1221f4cedf410d7efedd5094de457510b44afea1 Mon Sep 17 00:00:00 2001
From: wenovus <wenbli@google.com>
Date: Tue, 7 Mar 2023 17:20:14 -0800
Subject: [PATCH] Improve ygot.Diff performance

---
 ygot/diff.go                     | 32 ++++++++++++++++++++++++++++++++
 ygot/schema_tests/schema_test.go | 20 ++++++++++++++++----
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/ygot/diff.go b/ygot/diff.go
index e5add1589..b6afe7f26 100644
--- a/ygot/diff.go
+++ b/ygot/diff.go
@@ -50,9 +50,34 @@ func joingNMIPaths(parent *gnmipb.Path, child *gnmipb.Path) *gnmipb.Path {
 
 // pathSpec is a wrapper type used to store a set of gNMI paths to which
 // a value within a GoStruct corresponds to.
+//
+// NOTE: Once Equal() is called, no more changes must be made to pathSpec;
+// otherwise, future calls to Equal() will give the incorrect result.
 type pathSpec struct {
 	// gNMIPaths is the set of gNMI paths that the path represents.
 	gNMIPaths []*gnmipb.Path
+	// gNMIPathStrs is the set of gNMI paths that the path represents as strings.
+	gNMIPathStrs map[string]struct{}
+}
+
+func (p *pathSpec) populateStrs() {
+	if p == nil {
+		return
+	}
+
+	if p.gNMIPathStrs != nil {
+		return
+	}
+
+	p.gNMIPathStrs = map[string]struct{}{}
+	for _, path := range p.gNMIPaths {
+		str, err := PathToString(path)
+		if err != nil {
+			p.gNMIPathStrs = nil
+			return
+		}
+		p.gNMIPathStrs[str] = struct{}{}
+	}
 }
 
 // Equal compares two pathSpecs, returning true if all paths within the pathSpec
@@ -62,6 +87,13 @@ func (p *pathSpec) Equal(o *pathSpec) bool {
 		return p == o
 	}
 
+	p.populateStrs()
+	o.populateStrs()
+
+	if p.gNMIPathStrs != nil && o.gNMIPathStrs != nil {
+		return reflect.DeepEqual(p.gNMIPathStrs, o.gNMIPathStrs)
+	}
+
 	for _, path := range p.gNMIPaths {
 		var found bool
 		for _, otherPath := range o.gNMIPaths {
diff --git a/ygot/schema_tests/schema_test.go b/ygot/schema_tests/schema_test.go
index 5ecf98297..48dd4596d 100644
--- a/ygot/schema_tests/schema_test.go
+++ b/ygot/schema_tests/schema_test.go
@@ -17,6 +17,7 @@
 package schematest
 
 import (
+	"fmt"
 	"io/ioutil"
 	"reflect"
 	"testing"
@@ -591,8 +592,11 @@ func TestNotificationOutput(t *testing.T) {
 	}
 }
 
-func getBenchmarkDevice() *exampleoc.Device {
-	intfs := []string{"eth0", "eth1", "eth2", "eth3"}
+func getBenchmarkDevice(n int) *exampleoc.Device {
+	var intfs []string
+	for i := 0; i < n; i++ {
+		intfs = append(intfs, fmt.Sprintf("eth%d", i))
+	}
 	d := &exampleoc.Device{}
 	for _, intf := range intfs {
 		d.GetOrCreateInterface(intf)
@@ -603,7 +607,7 @@ func getBenchmarkDevice() *exampleoc.Device {
 }
 
 func BenchmarkNotificationOutput(b *testing.B) {
-	d := getBenchmarkDevice()
+	d := getBenchmarkDevice(4)
 	b.ResetTimer()
 	for i := 0; i != b.N; i++ {
 		if _, err := ygot.TogNMINotifications(d, 0, ygot.GNMINotificationsConfig{
@@ -615,7 +619,7 @@ func BenchmarkNotificationOutput(b *testing.B) {
 }
 
 func BenchmarkNotificationOutputElement(b *testing.B) {
-	d := getBenchmarkDevice()
+	d := getBenchmarkDevice(4)
 	b.ResetTimer()
 	for i := 0; i != b.N; i++ {
 		if _, err := ygot.TogNMINotifications(d, 0, ygot.GNMINotificationsConfig{
@@ -625,3 +629,11 @@ func BenchmarkNotificationOutputElement(b *testing.B) {
 		}
 	}
 }
+
+func BenchmarkDiff(b *testing.B) {
+	d1 := getBenchmarkDevice(1000)
+	d2 := getBenchmarkDevice(500)
+	for n := 0; n != b.N; n++ {
+		ygot.Diff(d1, d2)
+	}
+}