Merge pull request #30 from Nearoo/master

implement `OrderedWalk`: Walk in Topological Order
heimdalr · Dec 7, 2023 · 9f62a94 · 9f62a94
2 parents ee341df + 003273b
commit 9f62a94
Show file tree

Hide file tree

Showing 2 changed files with 156 additions and 29 deletions.
diff --git a/visitor.go b/visitor.go
@@ -105,3 +105,53 @@ func reversedVertexIDs(vertices map[string]interface{}) []string {
 	}
 	return ids
 }
+
+// OrderedWalk implements the Topological Sort algorithm to traverse the entire DAG.
+// This means that for any edge a -> b, node a will be visited before node b.
+func (d *DAG) OrderedWalk(visitor Visitor) {
+
+	d.muDAG.RLock()
+	defer d.muDAG.RUnlock()
+
+	queue := llq.New()
+	vertices := d.getRoots()
+	for _, id := range vertexIDs(vertices) {
+		v := vertices[id]
+		sv := storableVertex{WrappedID: id, Value: v}
+		queue.Enqueue(sv)
+	}
+
+	visited := make(map[string]bool, d.getOrder())
+
+Main:
+	for !queue.Empty() {
+		v, _ := queue.Dequeue()
+		sv := v.(storableVertex)
+
+		if visited[sv.WrappedID] {
+			continue
+		}
+
+		// if the current vertex has any parent that hasn't been visited yet,
+		// put it back into the queue, and work on the next element
+		parents, _ := d.GetParents(sv.WrappedID)
+		for parent := range parents {
+			if !visited[parent] {
+				queue.Enqueue(sv)
+				continue Main
+			}
+		}
+
+		if !visited[sv.WrappedID] {
+			visited[sv.WrappedID] = true
+			visitor.Visit(sv)
+		}
+
+		vertices, _ := d.getChildren(sv.WrappedID)
+		for _, id := range vertexIDs(vertices) {
+			v := vertices[id]
+			sv := storableVertex{WrappedID: id, Value: v}
+			queue.Enqueue(sv)
+		}
+	}
+}
diff --git a/visitor_test.go b/visitor_test.go
@@ -16,16 +16,17 @@ func (pv *testVisitor) Visit(v Vertexer) {
 }
 
 // schematic diagram:
-//     v5
-//     ^
-//     |
-//     v4
-//     ^
-//     |
-//     v2 --> v3
-//     ^
-//     |
-//     v1
+//
+//	v5
+//	^
+//	|
+//	v4
+//	^
+//	|
+//	v2 --> v3
+//	^
+//	|
+//	v1
 func getTestWalkDAG() *DAG {
 	dag := NewDAG()
 
@@ -44,13 +45,14 @@ func getTestWalkDAG() *DAG {
 }
 
 // schematic diagram:
-//    v4 --> v5
-//           ^
-//           |
-//    v1 --> v3
-//           ^
-//           |
-//          v2
+//
+//	v4 --> v5
+//	       ^
+//	       |
+//	v1 --> v3
+//	       ^
+//	       |
+//	      v2
 func getTestWalkDAG2() *DAG {
 	dag := NewDAG()
 
@@ -69,13 +71,14 @@ func getTestWalkDAG2() *DAG {
 }
 
 // schematic diagram:
-//    v4 --> v5
 //
+//	v4 --> v5
 //
-//    v1 --> v3
-//           ^
-//           |
-//          v2
+//
+//	v1 --> v3
+//	       ^
+//	       |
+//	      v2
 func getTestWalkDAG3() *DAG {
 	dag := NewDAG()
 
@@ -93,13 +96,14 @@ func getTestWalkDAG3() *DAG {
 }
 
 // schematic diagram:
-//     v4     v5
-//     ^      ^
-//     |      |
-//     v2 --> v3
-//     ^
-//     |
-//     v1
+//
+//	v4     v5
+//	^      ^
+//	|      |
+//	v2 --> v3
+//	^
+//	|
+//	v1
 func getTestWalkDAG4() *DAG {
 	dag := NewDAG()
 
@@ -117,6 +121,32 @@ func getTestWalkDAG4() *DAG {
 	return dag
 }
 
+// schematic diagram:
+//
+//	v5
+//	^
+//	|
+//	v3 <-- v4
+//	^      ^
+//	|      |
+//	v1     v2
+func getTestWalkDAG5() *DAG {
+	dag := NewDAG()
+
+	v1, v2, v3, v4, v5 := "1", "2", "3", "4", "5"
+	_ = dag.AddVertexByID(v1, "v1")
+	_ = dag.AddVertexByID(v2, "v2")
+	_ = dag.AddVertexByID(v3, "v3")
+	_ = dag.AddVertexByID(v4, "v4")
+	_ = dag.AddVertexByID(v5, "v5")
+	_ = dag.AddEdge(v1, v3)
+	_ = dag.AddEdge(v2, v4)
+	_ = dag.AddEdge(v4, v3)
+	_ = dag.AddEdge(v3, v5)
+
+	return dag
+}
+
 func TestDFSWalk(t *testing.T) {
 	cases := []struct {
 		dag      *DAG
@@ -138,6 +168,10 @@ func TestDFSWalk(t *testing.T) {
 			dag:      getTestWalkDAG4(),
 			expected: []string{"v1", "v2", "v3", "v5", "v4"},
 		},
+		{
+			dag:      getTestWalkDAG5(),
+			expected: []string{"v1", "v3", "v5", "v2", "v4"},
+		},
 	}
 
 	for _, c := range cases {
@@ -173,6 +207,10 @@ func TestBFSWalk(t *testing.T) {
 			dag:      getTestWalkDAG4(),
 			expected: []string{"v1", "v2", "v3", "v4", "v5"},
 		},
+		{
+			dag:      getTestWalkDAG5(),
+			expected: []string{"v1", "v2", "v3", "v4", "v5"},
+		},
 	}
 
 	for _, c := range cases {
@@ -186,3 +224,42 @@ func TestBFSWalk(t *testing.T) {
 		}
 	}
 }
+
+func TestOrderedWalk(t *testing.T) {
+	cases := []struct {
+		dag      *DAG
+		expected []string
+	}{
+		{
+			dag:      getTestWalkDAG(),
+			expected: []string{"v1", "v2", "v3", "v4", "v5"},
+		},
+		{
+			dag:      getTestWalkDAG2(),
+			expected: []string{"v1", "v2", "v4", "v3", "v5"},
+		},
+		{
+			dag:      getTestWalkDAG3(),
+			expected: []string{"v1", "v2", "v4", "v3", "v5"},
+		},
+		{
+			dag:      getTestWalkDAG4(),
+			expected: []string{"v1", "v2", "v3", "v4", "v5"},
+		},
+		{
+			dag:      getTestWalkDAG5(),
+			expected: []string{"v1", "v2", "v4", "v3", "v5"},
+		},
+	}
+
+	for _, c := range cases {
+		pv := &testVisitor{}
+		c.dag.OrderedWalk(pv)
+
+		expected := c.expected
+		actual := pv.Values
+		if deep.Equal(expected, actual) != nil {
+			t.Errorf("OrderedWalk() = %v, want %v", actual, expected)
+		}
+	}
+}