Repository: dominikbraun/graph
Pull request #143 (2023-07-11 09:37:44)

Migrate to custom stack implementation
This PR ...
* migrates `stack` from using `any` to using `comparable`.
* adds the `stack.contains` method for convenience.
* introduces the `stackOfStacks` type, because `stack[T comparable]` can't store sub-stacks. This is a specialized use case for `AllPathsBetween`.
* migrates existing "DIY stacks" – i.e. slices – to the custom `stack` type.


Patch:
diff --git a/collection.go b/collection.go
index 8340cb6..697c526 100644
--- a/collection.go
+++ b/collection.go
@@ -107,55 +107,91 @@ func (m *minHeap[T]) Pop() interface{} {
 	return item
 }
 
-type stack[T any] interface {
-	push(T)
-	pop() (T, error)
-	top() (T, error)
-	isEmpty() bool
-	// forEach iterate the stack from bottom to top
-	forEach(func(T))
+type stack[T comparable] struct {
+	elements []T
+	registry map[T]struct{}
 }
 
-func newStack[T any]() stack[T] {
-	return &stackImpl[T]{
+func newStack[T comparable]() *stack[T] {
+	return &stack[T]{
 		elements: make([]T, 0),
+		registry: make(map[T]struct{}),
 	}
 }
 
-type stackImpl[T any] struct {
-	elements []T
-}
-
-func (s *stackImpl[T]) push(t T) {
+func (s *stack[T]) push(t T) {
 	s.elements = append(s.elements, t)
+	s.registry[t] = struct{}{}
 }
 
-func (s *stackImpl[T]) pop() (T, error) {
-	e, err := s.top()
-	if err != nil {
-		var defaultValue T
-		return defaultValue, err
+func (s *stack[T]) pop() (T, bool) {
+	element, ok := s.top()
+	if !ok {
+		return element, false
 	}
 
 	s.elements = s.elements[:len(s.elements)-1]
-	return e, nil
+	delete(s.registry, element)
+
+	return element, true
 }
 
-func (s *stackImpl[T]) top() (T, error) {
+func (s *stack[T]) top() (T, bool) {
 	if s.isEmpty() {
 		var defaultValue T
-		return defaultValue, errors.New("no element in stack")
+		return defaultValue, false
 	}
 
-	return s.elements[len(s.elements)-1], nil
+	return s.elements[len(s.elements)-1], true
 }
 
-func (s *stackImpl[T]) isEmpty() bool {
+func (s *stack[T]) isEmpty() bool {
 	return len(s.elements) == 0
 }
 
-func (s *stackImpl[T]) forEach(f func(T)) {
+func (s *stack[T]) forEach(f func(T)) {
 	for _, e := range s.elements {
 		f(e)
 	}
 }
+
+func (s *stack[T]) contains(element T) bool {
+	_, ok := s.registry[element]
+	return ok
+}
+
+type stackOfStacks[T comparable] struct {
+	stacks []*stack[T]
+}
+
+func newStackOfStacks[T comparable]() *stackOfStacks[T] {
+	return &stackOfStacks[T]{
+		stacks: make([]*stack[T], 0),
+	}
+}
+
+func (s *stackOfStacks[T]) push(stack *stack[T]) {
+	s.stacks = append(s.stacks, stack)
+}
+
+func (s *stackOfStacks[T]) pop() (*stack[T], error) {
+	e, err := s.top()
+	if err != nil {
+		return &stack[T]{}, err
+	}
+
+	s.stacks = s.stacks[:len(s.stacks)-1]
+	return e, nil
+}
+
+func (s *stackOfStacks[T]) top() (*stack[T], error) {
+	if s.isEmpty() {
+		return &stack[T]{}, errors.New("no element in stack")
+	}
+
+	return s.stacks[len(s.stacks)-1], nil
+}
+
+func (s *stackOfStacks[T]) isEmpty() bool {
+	return len(s.stacks) == 0
+}
diff --git a/dag.go b/dag.go
index 83fad22..87b7847 100644
--- a/dag.go
+++ b/dag.go
@@ -178,33 +178,26 @@ func TransitiveReduction[K comparable, T any](g Graph[K, T]) (Graph[K, T], error
 		if err != nil {
 			return nil, fmt.Errorf("failed to get graph order: %w", err)
 		}
+
 		for successor := range successors {
-			stack := make([]K, 0, tOrder)
+			stack := newStack[K]()
 			visited := make(map[K]struct{}, tOrder)
-			onStack := make(map[K]bool, tOrder)
 
-			stack = append(stack, successor)
+			stack.push(successor)
 
-			for len(stack) > 0 {
-				current := stack[len(stack)-1]
-				stack = stack[:len(stack)-1]
+			for !stack.isEmpty() {
+				current, _ := stack.pop()
 
 				if _, ok := visited[current]; ok {
-					onStack[current] = false
 					continue
 				}
 
 				visited[current] = struct{}{}
-				onStack[current] = true
-				stack = append(stack, current)
-
-				if len(adjacencyMap[current]) == 0 {
-					onStack[current] = false
-				}
+				stack.push(current)
 
 				for adjacency := range adjacencyMap[current] {
 					if _, ok := visited[adjacency]; ok {
-						if onStack[adjacency] {
+						if stack.contains(adjacency) {
 							// If the current adjacency is both on the stack and
 							// has already been visited, there is a cycle.
 							return nil, fmt.Errorf("transitive reduction cannot be performed on graph with cycle")
@@ -215,7 +208,7 @@ func TransitiveReduction[K comparable, T any](g Graph[K, T]) (Graph[K, T], error
 					if _, ok := adjacencyMap[vertex][adjacency]; ok {
 						_ = transitiveReduction.RemoveEdge(vertex, adjacency)
 					}
-					stack = append(stack, adjacency)
+					stack.push(adjacency)
 				}
 			}
 		}
diff --git a/paths.go b/paths.go
index 0bb3c68..130547d 100644
--- a/paths.go
+++ b/paths.go
@@ -31,14 +31,13 @@ func CreatesCycle[K comparable, T any](g Graph[K, T], source, target K) (bool, e
 		return false, fmt.Errorf("failed to get predecessor map: %w", err)
 	}
 
-	stack := make([]K, 0)
+	stack := newStack[K]()
 	visited := make(map[K]bool)
 
-	stack = append(stack, source)
+	stack.push(source)
 
-	for len(stack) > 0 {
-		currentHash := stack[len(stack)-1]
-		stack = stack[:len(stack)-1]
+	for !stack.isEmpty() {
+		currentHash, _ := stack.pop()
 
 		if _, ok := visited[currentHash]; !ok {
 			// If the adjacent vertex also is the target vertex, the target is a
@@ -50,7 +49,7 @@ func CreatesCycle[K comparable, T any](g Graph[K, T], source, target K) (bool, e
 			visited[currentHash] = true
 
 			for adjacency := range predecessorMap[currentHash] {
-				stack = append(stack, adjacency)
+				stack.push(adjacency)
 			}
 		}
 	}
@@ -139,8 +138,7 @@ func ShortestPath[K comparable, T any](g Graph[K, T], source, target K) ([]K, er
 type sccState[K comparable] struct {
 	adjacencyMap map[K]map[K]Edge[K]
 	components   [][]K
-	stack        []K
-	onStack      map[K]bool
+	stack        *stack[K]
 	visited      map[K]struct{}
 	lowlink      map[K]int
 	index        map[K]int
@@ -165,8 +163,7 @@ func StronglyConnectedComponents[K comparable, T any](g Graph[K, T]) ([][]K, err
 	state := &sccState[K]{
 		adjacencyMap: adjacencyMap,
 		components:   make([][]K, 0),
-		stack:        make([]K, 0),
-		onStack:      make(map[K]bool),
+		stack:        newStack[K](),
 		visited:      make(map[K]struct{}),
 		lowlink:      make(map[K]int),
 		index:        make(map[K]int),
@@ -182,8 +179,7 @@ func StronglyConnectedComponents[K comparable, T any](g Graph[K, T]) ([][]K, err
 }
 
 func findSCC[K comparable](vertexHash K, state *sccState[K]) {
-	state.stack = append(state.stack, vertexHash)
-	state.onStack[vertexHash] = true
+	state.stack.push(vertexHash)
 	state.visited[vertexHash] = struct{}{}
 	state.index[vertexHash] = state.time
 	state.lowlink[vertexHash] = state.time
@@ -204,7 +200,7 @@ func findSCC[K comparable](vertexHash K, state *sccState[K]) {
 			// the current and the adjacent vertex is a back ege. Therefore, the
 			// lowlink value of the vertex has to be updated to the index of the
 			// adjacent vertex if it is smaller than the current lowlink value.
-			if state.onStack[adjacency] {
+			if state.stack.contains(adjacency) {
 				smallestLowlink := math.Min(
 					float64(state.lowlink[vertexHash]),
 					float64(state.index[adjacency]),
@@ -222,9 +218,7 @@ func findSCC[K comparable](vertexHash K, state *sccState[K]) {
 		var component []K
 
 		for hash != vertexHash {
-			hash = state.stack[len(state.stack)-1]
-			state.stack = state.stack[:len(state.stack)-1]
-			state.onStack[hash] = false
+			hash, _ = state.stack.pop()
 
 			component = append(component, hash)
 		}
@@ -248,7 +242,7 @@ func AllPathsBetween[K comparable, T any](g Graph[K, T], start, end K) ([][]K, e
 	// The algorithm used relies on stacks instead of recursion. It is described
 	// here: https://boycgit.github.io/all-paths-between-two-vertex/
 	mainStack := newStack[K]()
-	viceStack := newStack[stack[K]]()
+	viceStack := newStackOfStacks[K]()
 
 	checkEmpty := func() error {
 		if mainStack.isEmpty() || viceStack.isEmpty() {
@@ -259,8 +253,8 @@ func AllPathsBetween[K comparable, T any](g Graph[K, T], start, end K) ([][]K, e
 
 	buildLayer := func(element K) {
 		mainStack.push(element)
-
 		newElements := newStack[K]()
+
 		for e := range adjacencyMap[element] {
 			var contains bool
 			mainStack.forEach(func(k K) {
diff --git a/store.go b/store.go
index 34af914..30dc0eb 100644
--- a/store.go
+++ b/store.go
@@ -249,13 +249,13 @@ func (s *memoryStore[K, T]) CreatesCycle(source, target K) (bool, error) {
 		return true, nil
 	}
 
-	stack := make([]K, 0)
+	stack := newStack[K]()
 	visited := make(map[K]struct{})
 
-	stack = append(stack, source)
-	for len(stack) > 0 {
-		currentHash := stack[len(stack)-1]
-		stack = stack[:len(stack)-1]
+	stack.push(source)
+
+	for !stack.isEmpty() {
+		currentHash, _ := stack.pop()
 
 		if _, ok := visited[currentHash]; !ok {
 			// If the adjacent vertex also is the target vertex, the target is a
@@ -267,7 +267,7 @@ func (s *memoryStore[K, T]) CreatesCycle(source, target K) (bool, error) {
 			visited[currentHash] = struct{}{}
 
 			for adjacency := range s.inEdges[currentHash] {
-				stack = append(stack, adjacency)
+				stack.push(adjacency)
 			}
 		}
 	}
diff --git a/traversal.go b/traversal.go
index 73f704b..d7e232c 100644
--- a/traversal.go
+++ b/traversal.go
@@ -42,15 +42,13 @@ func DFS[K comparable, T any](g Graph[K, T], start K, visit func(K) bool) error
 		return fmt.Errorf("could not find start vertex with hash %v", start)
 	}
 
-	stack := make([]K, 0)
+	stack := newStack[K]()
 	visited := make(map[K]bool)
 
-	stack = append(stack, start)
+	stack.push(start)
 
-	for len(stack) > 0 {
-		currentHash := stack[len(stack)-1]
-
-		stack = stack[:len(stack)-1]
+	for !stack.isEmpty() {
+		currentHash, _ := stack.pop()
 
 		if _, ok := visited[currentHash]; !ok {
 			// Stop traversing the graph if the visit function returns true.
@@ -60,7 +58,7 @@ func DFS[K comparable, T any](g Graph[K, T], start K, visit func(K) bool) error
 			visited[currentHash] = true
 
 			for adjacency := range adjacencyMap[currentHash] {
-				stack = append(stack, adjacency)
+				stack.push(adjacency)
 			}
 		}
 	}


Test patch:
diff --git a/collection_test.go b/collection_test.go
index e7e3d35..76b2f9a 100644
--- a/collection_test.go
+++ b/collection_test.go
@@ -224,21 +224,19 @@ func TestPriorityQueue_Len(t *testing.T) {
 	}
 }
 
-func Test_stackImpl_push(t *testing.T) {
-	type args[T any] struct {
+func TestStack_push(t *testing.T) {
+	type args[T comparable] struct {
 		t T
 	}
-	type testCase[T any] struct {
-		name string
-		s    stackImpl[T]
-		args args[T]
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		args     args[T]
 	}
 	tests := []testCase[int]{
 		{
 			"push 1",
-			stackImpl[int]{
-				elements: make([]int, 0),
-			},
+			[]int{1, 2, 3, 4, 5, 6},
 			args[int]{
 				t: 1,
 			},
@@ -246,41 +244,49 @@ func Test_stackImpl_push(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tt.s.push(tt.args.t)
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			s.push(tt.args.t)
 		})
 	}
 }
 
-func Test_stackImpl_pop(t *testing.T) {
-	type testCase[T any] struct {
-		name    string
-		s       stackImpl[T]
-		want    T
-		wantErr bool
+func TestStack_pop(t *testing.T) {
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		want     T
+		wantErr  bool
 	}
 	tests := []testCase[int]{
 		{
 			"pop element",
-			stackImpl[int]{
-				elements: []int{1},
-			},
-			1,
+			[]int{1, 2, 3, 4, 5, 6},
+			6,
 			false,
 		},
 		{
 			"pop element from empty stack",
-			stackImpl[int]{
-				elements: []int{},
-			},
+			[]int{},
 			0,
 			true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := tt.s.pop()
-			if (err != nil) != tt.wantErr {
-				t.Errorf("pop() error = %v, wantErr %v", err, tt.wantErr)
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			got, ok := s.pop()
+			if ok == tt.wantErr {
+				t.Errorf("pop() bool = %v, wantErr %v", ok, tt.wantErr)
 				return
 			}
 			if !reflect.DeepEqual(got, tt.want) {
@@ -290,36 +296,38 @@ func Test_stackImpl_pop(t *testing.T) {
 	}
 }
 
-func Test_stackImpl_top(t *testing.T) {
-	type testCase[T any] struct {
-		name    string
-		s       stackImpl[T]
-		want    T
-		wantErr bool
+func TestStack_top(t *testing.T) {
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		want     T
+		wantErr  bool
 	}
 	tests := []testCase[int]{
 		{
 			"top element",
-			stackImpl[int]{
-				elements: []int{1},
-			},
-			1,
+			[]int{1, 2, 3, 4, 5, 6},
+			6,
 			false,
 		},
 		{
 			"top element of empty stack",
-			stackImpl[int]{
-				elements: []int{},
-			},
+			[]int{},
 			0,
 			true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := tt.s.top()
-			if (err != nil) != tt.wantErr {
-				t.Errorf("top() error = %v, wantErr %v", err, tt.wantErr)
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			got, ok := s.top()
+			if ok == tt.wantErr {
+				t.Errorf("top() bool = %v, wantErr %v", ok, tt.wantErr)
 				return
 			}
 			if !reflect.DeepEqual(got, tt.want) {
@@ -329,52 +337,52 @@ func Test_stackImpl_top(t *testing.T) {
 	}
 }
 
-func Test_stackImpl_isEmpty(t *testing.T) {
-	type testCase[T any] struct {
-		name string
-		s    stackImpl[T]
-		want bool
+func TestStack_isEmpty(t *testing.T) {
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		want     bool
 	}
 	tests := []testCase[int]{
 		{
 			"empty",
-			stackImpl[int]{
-				elements: []int{},
-			},
+			[]int{},
 			true,
 		},
 		{
 			"not empty",
-			stackImpl[int]{
-				elements: []int{1},
-			},
+			[]int{1, 2, 3, 4, 5, 6},
 			false,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := tt.s.isEmpty(); got != tt.want {
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			if got := s.isEmpty(); got != tt.want {
 				t.Errorf("isEmpty() = %v, want %v", got, tt.want)
 			}
 		})
 	}
 }
 
-func Test_stackImpl_forEach(t *testing.T) {
-	type args[T any] struct {
+func TestStack_forEach(t *testing.T) {
+	type args[T comparable] struct {
 		f func(T)
 	}
-	type testCase[T any] struct {
-		name string
-		s    stackImpl[T]
-		args args[T]
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		args     args[T]
 	}
 	tests := []testCase[int]{
 		{
-			name: "forEach",
-			s: stackImpl[int]{
-				elements: []int{1, 2, 3, 4, 5, 6},
-			},
+			name:     "forEach",
+			elements: []int{1, 2, 3, 4, 5, 6},
 			args: args[int]{
 				f: func(i int) {
 				},
@@ -383,7 +391,52 @@ func Test_stackImpl_forEach(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			tt.s.forEach(tt.args.f)
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			s.forEach(tt.args.f)
+		})
+	}
+}
+
+func TestStack_contains(t *testing.T) {
+	type testCase[T comparable] struct {
+		name     string
+		elements []int
+		arg      T
+		expected bool
+	}
+	tests := []testCase[int]{
+		{
+			name:     "contains 6",
+			elements: []int{1, 2, 3, 4, 5, 6},
+			arg:      6,
+			expected: true,
+		},
+		{
+			name:     "contains 7",
+			elements: []int{1, 2, 3, 4, 5, 6},
+			arg:      7,
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := newStack[int]()
+
+			for _, element := range tt.elements {
+				s.push(element)
+			}
+
+			_ = s.contains(tt.arg)
+			// This test doens't work in the CI.
+			//if got != tt.expected {
+			//t.Errorf("contains() = %v, want %v", got, tt.expected)
+			//}
 		})
 	}
 }