use generics to make graph reusable

Signed-off-by: Nicolas De Loof <[email protected]>

Author: ndeloof

graph/graph.go

@@ -20,74 +20,39 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/compose-spec/compose-go/v2/types"
 	"github.com/compose-spec/compose-go/v2/utils"
 	"golang.org/x/exp/slices"
 )
 
 // graph represents project as service dependencies
-type graph struct {
-	vertices map[string]*vertex
+type graph[T any] struct {
+	vertices map[string]*vertex[T]
 }
 
 // vertex represents a service in the dependencies structure
-type vertex struct {
+type vertex[T any] struct {
 	key      string
-	service  *types.ServiceConfig
-	children map[string]*vertex
-	parents  map[string]*vertex
+	service  *T
+	children map[string]*vertex[T]
+	parents  map[string]*vertex[T]
 }
 
-// newGraph creates a service graph from project
-func newGraph(project *types.Project) (*graph, error) {
-	g := &graph{
-		vertices: map[string]*vertex{},
-	}
-
-	for name, s := range project.Services {
-		g.addVertex(name, s)
-	}
-
-	for name, s := range project.Services {
-		src := g.vertices[name]
-		for dep, condition := range s.DependsOn {
-			dest, ok := g.vertices[dep]
-			if !ok {
-				if condition.Required {
-					if ds, exists := project.DisabledServices[dep]; exists {
-						return nil, fmt.Errorf("service %q is required by %q but is disabled. Can be enabled by profiles %s", dep, name, ds.Profiles)
-					}
-					return nil, fmt.Errorf("service %q depends on unknown service %q", name, dep)
-				}
-				delete(s.DependsOn, name)
-				project.Services[name] = s
-				continue
-			}
-			src.children[dep] = dest
-			dest.parents[name] = src
-		}
-	}
-
-	err := g.checkCycle()
-	return g, err
-}
-
-func (g *graph) addVertex(name string, service types.ServiceConfig) {
-	g.vertices[name] = &vertex{
+func (g *graph[T]) addVertex(name string, service T) {
+	g.vertices[name] = &vertex[T]{
 		key:      name,
 		service:  &service,
-		parents:  map[string]*vertex{},
-		children: map[string]*vertex{},
+		parents:  map[string]*vertex[T]{},
+		children: map[string]*vertex[T]{},
 	}
 }
 
-func (g *graph) addEdge(src, dest string) {
+func (g *graph[T]) addEdge(src, dest string) {
 	g.vertices[src].children[dest] = g.vertices[dest]
 	g.vertices[dest].parents[src] = g.vertices[src]
 }
 
-func (g *graph) roots() []*vertex {
-	var res []*vertex
+func (g *graph[T]) roots() []*vertex[T] {
+	var res []*vertex[T]
 	for _, v := range g.vertices {
 		if len(v.parents) == 0 {
 			res = append(res, v)
@@ -96,8 +61,8 @@ func (g *graph) roots() []*vertex {
 	return res
 }
 
-func (g *graph) leaves() []*vertex {
-	var res []*vertex
+func (g *graph[T]) leaves() []*vertex[T] {
+	var res []*vertex[T]
 	for _, v := range g.vertices {
 		if len(v.children) == 0 {
 			res = append(res, v)
@@ -107,8 +72,8 @@ func (g *graph) leaves() []*vertex {
 	return res
 }
 
-func (g *graph) checkCycle() error {
-	// iterate on verticles in a name-order to render a predicable error message
+func (g *graph[T]) checkCycle() error {
+	// iterate on vertices in a name-order to render a predicable error message
 	// this is required by tests and enforce command reproducibility by user, which otherwise could be confusing
 	names := utils.MapKeys(g.vertices)
 	for _, name := range names {
@@ -120,7 +85,7 @@ func (g *graph) checkCycle() error {
 	return nil
 }
 
-func searchCycle(path []string, v *vertex) error {
+func searchCycle[T any](path []string, v *vertex[T]) error {
 	names := utils.MapKeys(v.children)
 	for _, name := range names {
 		if i := slices.Index(path, name); i > 0 {
@@ -136,7 +101,7 @@ func searchCycle(path []string, v *vertex) error {
 }
 
 // descendents return all descendents for a vertex, might contain duplicates
-func (v *vertex) descendents() []string {
+func (v *vertex[T]) descendents() []string {
 	var vx []string
 	for _, n := range v.children {
 		vx = append(vx, n.key)

graph/graph_test.go

@@ -112,7 +112,7 @@ func TestBuildGraph(t *testing.T) {
 		desc             string
 		services         types.Services
 		disabled         types.Services
-		expectedVertices map[string]*vertex
+		expectedVertices map[string]*vertex[types.ServiceConfig]
 		expectedError    string
 	}{
 		{
@@ -123,12 +123,12 @@ func TestBuildGraph(t *testing.T) {
 					DependsOn: types.DependsOnConfig{},
 				},
 			},
-			expectedVertices: map[string]*vertex{
+			expectedVertices: map[string]*vertex[types.ServiceConfig]{
 				"test": {
 					key:      "test",
 					service:  &types.ServiceConfig{Name: "test"},
-					children: map[string]*vertex{},
-					parents:  map[string]*vertex{},
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents:  map[string]*vertex[types.ServiceConfig]{},
 				},
 			},
 		},
@@ -144,18 +144,18 @@ func TestBuildGraph(t *testing.T) {
 					DependsOn: types.DependsOnConfig{},
 				},
 			},
-			expectedVertices: map[string]*vertex{
+			expectedVertices: map[string]*vertex[types.ServiceConfig]{
 				"test": {
 					key:      "test",
 					service:  &types.ServiceConfig{Name: "test"},
-					children: map[string]*vertex{},
-					parents:  map[string]*vertex{},
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents:  map[string]*vertex[types.ServiceConfig]{},
 				},
 				"another": {
 					key:      "another",
 					service:  &types.ServiceConfig{Name: "another"},
-					children: map[string]*vertex{},
-					parents:  map[string]*vertex{},
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents:  map[string]*vertex[types.ServiceConfig]{},
 				},
 			},
 		},
@@ -173,20 +173,20 @@ func TestBuildGraph(t *testing.T) {
 					DependsOn: types.DependsOnConfig{},
 				},
 			},
-			expectedVertices: map[string]*vertex{
+			expectedVertices: map[string]*vertex[types.ServiceConfig]{
 				"test": {
 					key:     "test",
 					service: &types.ServiceConfig{Name: "test"},
-					children: map[string]*vertex{
+					children: map[string]*vertex[types.ServiceConfig]{
 						"another": {},
 					},
-					parents: map[string]*vertex{},
+					parents: map[string]*vertex[types.ServiceConfig]{},
 				},
 				"another": {
 					key:      "another",
 					service:  &types.ServiceConfig{Name: "another"},
-					children: map[string]*vertex{},
-					parents: map[string]*vertex{
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents: map[string]*vertex[types.ServiceConfig]{
 						"test": {},
 					},
 				},
@@ -204,12 +204,12 @@ func TestBuildGraph(t *testing.T) {
 					},
 				},
 			},
-			expectedVertices: map[string]*vertex{
+			expectedVertices: map[string]*vertex[types.ServiceConfig]{
 				"test": {
 					key:      "test",
 					service:  &types.ServiceConfig{Name: "test"},
-					children: map[string]*vertex{},
-					parents:  map[string]*vertex{},
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents:  map[string]*vertex[types.ServiceConfig]{},
 				},
 			},
 		},
@@ -268,30 +268,30 @@ func TestBuildGraph(t *testing.T) {
 					DependsOn: types.DependsOnConfig{},
 				},
 			},
-			expectedVertices: map[string]*vertex{
+			expectedVertices: map[string]*vertex[types.ServiceConfig]{
 				"test": {
 					key:     "test",
 					service: &types.ServiceConfig{Name: "test"},
-					children: map[string]*vertex{
+					children: map[string]*vertex[types.ServiceConfig]{
 						"another": {},
 					},
-					parents: map[string]*vertex{},
+					parents: map[string]*vertex[types.ServiceConfig]{},
 				},
 				"another": {
 					key:     "another",
 					service: &types.ServiceConfig{Name: "another"},
-					children: map[string]*vertex{
+					children: map[string]*vertex[types.ServiceConfig]{
 						"another_dep": {},
 					},
-					parents: map[string]*vertex{
+					parents: map[string]*vertex[types.ServiceConfig]{
 						"test": {},
 					},
 				},
 				"another_dep": {
 					key:      "another_dep",
 					service:  &types.ServiceConfig{Name: "another_dep"},
-					children: map[string]*vertex{},
-					parents: map[string]*vertex{
+					children: map[string]*vertex[types.ServiceConfig]{},
+					parents: map[string]*vertex[types.ServiceConfig]{
 						"another": {},
 					},
 				},
@@ -384,7 +384,7 @@ func TestWith_RootNodesAndUp(t *testing.T) {
 	}
 }
 
-func assertVertexEqual(t *testing.T, a, b vertex) {
+func assertVertexEqual(t *testing.T, a, b vertex[types.ServiceConfig]) {
 	assert.Equal(t, a.key, b.key)
 	assert.Equal(t, a.service.Name, b.service.Name)
 	for c := range a.children {
@@ -397,9 +397,9 @@ func assertVertexEqual(t *testing.T, a, b vertex) {
 	}
 }
 
-func exampleGraph() *graph {
-	graph := &graph{
-		vertices: map[string]*vertex{},
+func exampleGraph() *graph[types.ServiceConfig] {
+	graph := &graph[types.ServiceConfig]{
+		vertices: map[string]*vertex[types.ServiceConfig]{},
 	}
 
 	/** graph topology:

graph/services.go

@@ -0,0 +1,80 @@
+/*
+   Copyright 2020 The Compose Specification Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package graph
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/compose-spec/compose-go/v2/types"
+)
+
+// InDependencyOrder walk the service graph an invoke VisitorFn in respect to dependency order
+func InDependencyOrder(ctx context.Context, project *types.Project, fn VisitorFn[types.ServiceConfig], options ...func(*Options)) error {
+	_, err := CollectInDependencyOrder[any](ctx, project, func(ctx context.Context, s string, config types.ServiceConfig) (any, error) {
+		return nil, fn(ctx, s, config)
+	}, options...)
+	return err
+}
+
+// CollectInDependencyOrder walk the service graph an invoke CollectorFn in respect to dependency order, then return result for each call
+func CollectInDependencyOrder[T any](ctx context.Context, project *types.Project, fn CollectorFn[types.ServiceConfig, T], options ...func(*Options)) (map[string]T, error) {
+	graph, err := newGraph(project)
+	if err != nil {
+		return nil, err
+	}
+	t := newTraversal(fn)
+	for _, option := range options {
+		option(t.Options)
+	}
+	err = walk(ctx, graph, t)
+	return t.results, err
+}
+
+// newGraph creates a service graph from project
+func newGraph(project *types.Project) (*graph[types.ServiceConfig], error) {
+	g := &graph[types.ServiceConfig]{
+		vertices: map[string]*vertex[types.ServiceConfig]{},
+	}
+
+	for name, s := range project.Services {
+		g.addVertex(name, s)
+	}
+
+	for name, s := range project.Services {
+		src := g.vertices[name]
+		for dep, condition := range s.DependsOn {
+			dest, ok := g.vertices[dep]
+			if !ok {
+				if condition.Required {
+					if ds, exists := project.DisabledServices[dep]; exists {
+						return nil, fmt.Errorf("service %q is required by %q but is disabled. Can be enabled by profiles %s", dep, name, ds.Profiles)
+					}
+					return nil, fmt.Errorf("service %q depends on unknown service %q", name, dep)
+				}
+				delete(s.DependsOn, name)
+				project.Services[name] = s
+				continue
+			}
+			src.children[dep] = dest
+			dest.parents[name] = src
+		}
+	}
+
+	err := g.checkCycle()
+	return g, err
+}