table: make AdjRib API take slice of RouteFamily

Make the code simpler

Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 39 insertions, 35 deletions

diff --git a/table/table_manager.go b/table/table_manager.go
index acd23676..f52f1ae9 100644
--- a/table/table_manager.go
+++ b/table/table_manager.go
@@ -540,55 +540,59 @@ func (adj *AdjRib) UpdateOut(pathList []*Path) {
 	adj.update(adj.adjRibOut, pathList)
 }
 
-func (adj *AdjRib) getPathList(rib map[string]*Path) []*Path {
-	pathList := make([]*Path, 0, len(rib))
-	for _, rr := range rib {
-		pathList = append(pathList, rr)
+func (adj *AdjRib) GetInPathList(rfList []bgp.RouteFamily) []*Path {
+	pathList := make([]*Path, 0, adj.GetInCount(rfList))
+	for _, rf := range rfList {
+		for _, rr := range adj.adjRibIn[rf] {
+			pathList = append(pathList, rr)
+		}
 	}
 	return pathList
 }
 
-func (adj *AdjRib) GetInPathList(rf bgp.RouteFamily) []*Path {
-	if _, ok := adj.adjRibIn[rf]; !ok {
-		return []*Path{}
+func (adj *AdjRib) GetOutPathList(rfList []bgp.RouteFamily) []*Path {
+	pathList := make([]*Path, 0, adj.GetOutCount(rfList))
+	for _, rf := range rfList {
+		for _, rr := range adj.adjRibOut[rf] {
+			pathList = append(pathList, rr)
+		}
 	}
-	return adj.getPathList(adj.adjRibIn[rf])
+	return pathList
 }
 
-func (adj *AdjRib) GetOutPathList(rf bgp.RouteFamily) []*Path {
-	if _, ok := adj.adjRibOut[rf]; !ok {
-		return []*Path{}
-	}
-	return adj.getPathList(adj.adjRibOut[rf])
-}
+func (adj *AdjRib) GetInCount(rfList []bgp.RouteFamily) int {
+	count := 0
+	for _, rf := range rfList {
+		if _, ok := adj.adjRibIn[rf]; ok {
+			count += len(adj.adjRibIn[rf])
 
-func (adj *AdjRib) GetInCount(rf bgp.RouteFamily) int {
-	if _, ok := adj.adjRibIn[rf]; !ok {
-		return 0
+		}
 	}
-	return len(adj.adjRibIn[rf])
+	return count
 }
 
-func (adj *AdjRib) GetOutCount(rf bgp.RouteFamily) int {
-	if _, ok := adj.adjRibOut[rf]; !ok {
-		return 0
+func (adj *AdjRib) GetOutCount(rfList []bgp.RouteFamily) int {
+	count := 0
+	for _, rf := range rfList {
+		if _, ok := adj.adjRibOut[rf]; ok {
+			count += len(adj.adjRibOut[rf])
+		}
 	}
-	return len(adj.adjRibOut[rf])
+	return count
 }
 
-func (adj *AdjRib) DropIn(rf bgp.RouteFamily) {
-	if _, ok := adj.adjRibIn[rf]; ok {
-		adj.adjRibIn[rf] = make(map[string]*Path)
+func (adj *AdjRib) DropIn(rfList []bgp.RouteFamily) {
+	for _, rf := range rfList {
+		if _, ok := adj.adjRibIn[rf]; ok {
+			adj.adjRibIn[rf] = make(map[string]*Path)
+		}
 	}
 }
 
-func (adj *AdjRib) DropOut(rf bgp.RouteFamily) {
-	if _, ok := adj.adjRibIn[rf]; ok {
-		adj.adjRibOut[rf] = make(map[string]*Path)
+func (adj *AdjRib) DropOut(rfList []bgp.RouteFamily) {
+	for _, rf := range rfList {
+		if _, ok := adj.adjRibIn[rf]; ok {
+			adj.adjRibOut[rf] = make(map[string]*Path)
+		}
 	}
 }
-
-func (adj *AdjRib) DropAll(rf bgp.RouteFamily) {
-	adj.DropIn(rf)
-	adj.DropOut(rf)
-}
diff --git a/table/table_manager_test.go b/table/table_manager_test.go
index 122d60a1..a836347a 100644
--- a/table/table_manager_test.go
+++ b/table/table_manager_test.go
@@ -2154,7 +2154,7 @@ func TestProcessBGPUpdate_Timestamp(t *testing.T) {
 	//t2 = path2.timestamp
 	adjRib.UpdateIn(pList2)
 
-	inList := adjRib.GetInPathList(bgp.RF_IPv4_UC)
+	inList := adjRib.GetInPathList([]bgp.RouteFamily{bgp.RF_IPv4_UC})
 	assert.Equal(t, len(inList), 1)
 	assert.Equal(t, inList[0].GetTimestamp(), t1)
 
@@ -2171,7 +2171,7 @@ func TestProcessBGPUpdate_Timestamp(t *testing.T) {
 	t3 := pList3[0].GetTimestamp()
 	adjRib.UpdateIn(pList3)
 
-	inList = adjRib.GetInPathList(bgp.RF_IPv4_UC)
+	inList = adjRib.GetInPathList([]bgp.RouteFamily{bgp.RF_IPv4_UC})
 	assert.Equal(t, len(inList), 1)
 	assert.Equal(t, inList[0].GetTimestamp(), t3)
 }