add topological prune

Author: peiit

README.md

@@ -10,6 +10,49 @@ Sorting and pruning of DAG graphs.
 
 Ideas borrowed from [python graphlib](https://github.com/python/cpython/blob/3.14/Lib/graphlib.py)
 
+# How to install
+```bash
+go get -u github.com/aio-arch/graphlib
+```
+
+# How to use
+```golang
+    import "github.com/aio-arch/graphlib"
+
+    // New A graph
+
+	// for string type
+	g1 := graphlib.NewGraph[string]()
+
+	// for int type
+	g2 := graphlib.NewGraph[int]()
+
+	// for add node and add edge
+	g1.AddNode("A1")
+	g1.AddNode("B2")
+	g1.AddEdge("A1", "B2") // edge: A1 -> B2
+
+	// for add mulit edge,add node inline
+	g2.Add(10, 1, 9)    // edge: 1 -> 10 and 9 -> 10
+	g2.Add(100, 10, 90) // edge: 10 -> 100 and 90 -> 100
+
+	// topological order
+	topo, err := graphlib.TopologicalOrder(g1)
+	if err != nil {
+		fmt.Println(err.Error())
+	}
+	fmt.Printf("Topological Order:%v\n", topo)
+
+	// topological prune
+	g3, err := graphlib.TopologicalPrune(g2, []int{10, 90})
+	if err != nil {
+		fmt.Println(err.Error())
+	}
+	_ = g3
+
+```
+
+
 # For Go Veriosn < 1.21
 ```bash
 go mod edit -replace golang.org/x/exp=golang.org/x/exp@v0.0.0-20240904232852-e7e105dedf7e

example/main.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/aio-arch/graphlib"
+)
+
+func main() {
+
+	/*
+	                                ┌────► A-1-1
+	                                │
+	   ┌────────────────────────► A-1
+	   │                           ▲│
+	   A       ┌─────► A-2-1       │└────► A-1-2
+	   │       │                   │
+	   └────► A-2                  │
+	           │                   │
+	           └─────► A-2-2       │
+	                      │        │
+	                      ├─────► C-1
+	                      │
+	   B ───► B-1 ───► B-1-1
+	*/
+	g := graphlib.NewGraph[string]()
+
+	// add node
+	g.AddNode("A-1-1")
+	g.AddNode("A-1")
+
+	//add edge
+	g.AddEdge("A-1", "A-1-1") // edge: A-1 -> A-1-1
+
+	g.AddNode("A-1-2")
+	g.AddEdge("A-1", "A-1-2") // edge: A-1 -> A-1-2
+
+	// add multiple edges and add node inline
+	g.Add("A-1", "A", "C-1") // edge: A -> A-1 and C-1 -> A-1
+	g.Add("A-2-1", "A-2")    // edge: A-2 -> A-2-1
+	g.Add("A-2-2", "A-2")
+	g.Add("A-2", "A")
+	g.Add("B-1-1", "B-1")
+	g.Add("B-1", "B")
+	g.Add("A") // add node A
+	g.Add("B") // add node B
+	g.Add("C-1", "A-2-2", "B-1-1")
+
+	topo, err := graphlib.TopologicalOrder(g)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Topological Order:%v\n", topo)
+	//Topological Order:[A B A-2 B-1 A-2-1 A-2-2 B-1-1 C-1 A-1 A-1-1 A-1-2]
+
+	g2, err := graphlib.TopologicalPrune(g, []string{"A-1-1"})
+	if err != nil {
+		panic(err)
+	}
+	topo2, err := graphlib.TopologicalOrder(g2)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Pruned Topological Order:%v\n", topo2)
+	//Pruned Topological Order:[A B A-2 B-1 A-2-2 B-1-1 C-1 A-1 A-1-1]
+}

graph.go

@@ -1,18 +1,18 @@
 package graphlib
 
 type nodeInfo[V comparable] struct {
-	SortIdx         uint       //Sort idx
-	PredecessorNums uint       //Number of parent nodes
-	SuccessorNums   uint       //Record next Successor append index
-	Successor       map[V]uint //Child Nodes
+	SortIdx         uint //Sort idx
+	PredecessorNums uint //Number of parent nodes
+	Successors      []V  //Child Nodes
 }
 
-func (n *nodeInfo[V]) SuccessorSortSet() []V {
-	set := make([]V, n.SuccessorNums)
-	for node, idx := range n.Successor {
-		set[idx] = node
+func (ni *nodeInfo[V]) hasSuccessor(node V) bool {
+	for _, successor := range ni.Successors {
+		if successor == node {
+			return true
+		}
 	}
-	return set
+	return false
 }
 
 type Graph[V comparable] struct {
@@ -31,8 +31,8 @@ func (g *Graph[V]) AddNode(node V) *nodeInfo[V] {
 		return result
 	}
 	result := nodeInfo[V]{
-		SortIdx:   g.nodeNums,
-		Successor: make(map[V]uint, 2),
+		SortIdx:    g.nodeNums,
+		Successors: make([]V, 0, 2),
 	}
 	g.node2info[node] = &result
 	g.nodeNums++
@@ -48,9 +48,8 @@ func (g *Graph[V]) AddEdge(from, to V) {
 	if t, has = g.node2info[to]; !has {
 		panic(ErrUnknownNode[V]{node: to})
 	}
-	if _, has := f.Successor[to]; !has {
-		f.Successor[to] = f.SuccessorNums
-		f.SuccessorNums++
+	if !f.hasSuccessor(to) {
+		f.Successors = append(f.Successors, to)
 		t.PredecessorNums++
 	}
 }
@@ -98,7 +97,7 @@ func (g *Graph[V]) IsAcyclic() ([]V, bool) {
 			} else {
 				seen[node] = struct{}{}
 				itStack = append(itStack, iterItem[V]{isEnd: true})
-				for _, successor := range g.node2info[node].SuccessorSortSet() {
+				for _, successor := range g.node2info[node].Successors {
 					itStack = append(itStack, iterItem[V]{val: successor})
 				}
 				node2stacki[node] = len(stack)

order.go

@@ -0,0 +1,60 @@
+package graphlib
+
+type topologicalOrder[V comparable] struct {
+	graph         *Graph[V]
+	traversedNums uint
+	traversed     map[V]uint
+	notReadyNodes map[V]uint
+}
+
+func (ts *topologicalOrder[V]) done(nodes *[]V) []V {
+	readyNodes := make([]V, 0, 8)
+	for _, node := range *nodes {
+		ts.traversed[node] = ts.traversedNums
+		ts.traversedNums++
+
+		nodeInfo := ts.graph.node2info[node]
+		for _, successor := range nodeInfo.Successors {
+			successorInfo := ts.graph.node2info[successor]
+			if successorInfo.PredecessorNums == 1 || ts.notReadyNodes[successor] == 1 {
+				readyNodes = append(readyNodes, successor)
+				delete(ts.notReadyNodes, successor)
+			} else {
+				ts.notReadyNodes[successor] = successorInfo.PredecessorNums - 1
+			}
+		}
+	}
+	return readyNodes
+}
+
+func (ts *topologicalOrder[V]) traverse() {
+	//find roots
+	rootNodes := make([]V, 0, len(ts.graph.node2info)/2)
+	for _, v := range ts.graph.NodeSortSet() {
+		nodeInfo := ts.graph.node2info[v]
+		if nodeInfo.PredecessorNums == 0 {
+			rootNodes = append(rootNodes, v)
+		}
+	}
+	reayNodes := ts.done(&rootNodes)
+	for len(reayNodes) > 0 {
+		reayNodes = ts.done(&reayNodes)
+	}
+}
+
+func graphOrder[V comparable](g *Graph[V]) []V {
+	nodeLen := len(g.node2info)
+
+	ts := topologicalOrder[V]{
+		graph:         g,
+		traversed:     make(map[V]uint, nodeLen),
+		notReadyNodes: make(map[V]uint, nodeLen),
+	}
+	ts.traverse()
+
+	set := make([]V, ts.traversedNums)
+	for node, idx := range ts.traversed {
+		set[idx] = node
+	}
+	return set
+}

prune.go

@@ -0,0 +1,87 @@
+package graphlib
+
+type topologicalPrune[V comparable] struct {
+	graph  *Graph[V]
+	target *Graph[V]
+}
+
+// func (tp *topologicalPrune[V]) sortWaitingNodes(nodes []V) {
+// 	sort.Slice(nodes, func(i, j int) bool {
+// 		return tp.graph.node2info[nodes[i]].SortIdx < tp.graph.node2info[nodes[j]].SortIdx
+// 	})
+// }
+
+func (tp *topologicalPrune[V]) prune(nodes []V) *Graph[V] {
+
+	traversed := make(map[V]int, len(nodes))
+	readyNodes := make(map[V]uint, len(nodes))
+	childNodes := make(map[V]uint, len(nodes))
+	for _, node := range nodes {
+		nodeInfo := tp.graph.node2info[node]
+		readyNodes[node] = nodeInfo.PredecessorNums
+	}
+
+	for len(readyNodes) > 0 {
+		for node, predecessorNums := range readyNodes {
+
+			delete(readyNodes, node)
+			traversed[node] = 0
+			if predecessorNums == 0 {
+				tp.target.AddNode(node)
+			} else {
+				childNodes[node] = predecessorNums
+			}
+		}
+
+		for node, info := range tp.graph.node2info {
+			traversedSuccessorNums, ok := traversed[node]
+			if ok && len(info.Successors) == traversedSuccessorNums {
+				continue
+			}
+
+			var ready bool
+			for _, successor := range info.Successors {
+				if successorPredecessorNums, has := childNodes[successor]; has {
+					ready = true
+					tp.target.Add(successor, node)
+					traversed[node] = traversedSuccessorNums + 1
+					if successorPredecessorNums == 1 {
+						delete(childNodes, successor)
+					} else {
+						childNodes[successor] = successorPredecessorNums - 1
+					}
+				}
+
+			}
+			if ready && info.PredecessorNums > 0 {
+				readyNodes[node] = info.PredecessorNums
+			}
+			if len(childNodes) == 0 {
+				break
+			}
+		}
+	}
+
+	// sort by graph order
+	var idx uint
+	for _, node := range tp.graph.NodeSortSet() {
+		if info, has := tp.target.node2info[node]; has {
+			info.SortIdx = idx
+			idx++
+		}
+		if idx == tp.target.nodeNums {
+			break
+		}
+	}
+
+	return tp.target
+}
+
+func graphPrune[V comparable](g *Graph[V], nodes []V) *Graph[V] {
+	tp := topologicalPrune[V]{
+		graph:  g,
+		target: NewGraph[V](),
+	}
+	tp.prune(nodes)
+	return tp.target
+}