sql/schemachanger: remove cycle DepEdge rules, add SameStage kind

pkg/sql/schemachanger/scgraph/BUILD.bazel

@@ -1,10 +1,14 @@
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+load("//build:STRINGER.bzl", "stringer")
 
 go_library(
     name = "scgraph",
     srcs = [
+        "dep_edge_tree.go",
+        "edge.go",
         "graph.go",
         "iteration.go",
+        ":gen-depedgekind-stringer",  # keep
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scgraph",
     visibility = ["//visibility:public"],
@@ -21,12 +25,23 @@ go_library(
 
 go_test(
     name = "scgraph_test",
-    srcs = ["graph_test.go"],
+    srcs = [
+        "dep_edge_tree_test.go",
+        "graph_test.go",
+    ],
+    embed = [":scgraph"],
     deps = [
-        ":scgraph",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/schemachanger/scop",
         "//pkg/sql/schemachanger/scpb",
+        "//pkg/util/iterutil",
+        "//pkg/util/leaktest",
         "@com_github_stretchr_testify//require",
     ],
 )
+
+stringer(
+    name = "gen-depedgekind-stringer",
+    src = "edge.go",
+    typ = "DepEdgeKind",
+)

pkg/sql/schemachanger/scgraph/dep_edge_tree.go

@@ -0,0 +1,102 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package scgraph
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
+	"github.com/cockroachdb/cockroach/pkg/util/iterutil"
+	"github.com/google/btree"
+)
+
+type depEdgeTree struct {
+	t     *btree.BTree
+	order edgeTreeOrder
+	cmp   nodeCmpFn
+}
+
+type nodeCmpFn func(a, b *scpb.Node) (less, eq bool)
+
+func newDepEdgeTree(order edgeTreeOrder, cmp nodeCmpFn) *depEdgeTree {
+	const degree = 8 // arbitrary
+	return &depEdgeTree{
+		t:     btree.New(degree),
+		order: order,
+		cmp:   cmp,
+	}
+}
+
+// edgeTreeOrder order in which the edge tree is sorted,
+// either based on from/to node indexes.
+type edgeTreeOrder bool
+
+func (o edgeTreeOrder) first(e Edge) *scpb.Node {
+	if o == fromTo {
+		return e.From()
+	}
+	return e.To()
+}
+
+func (o edgeTreeOrder) second(e Edge) *scpb.Node {
+	if o == toFrom {
+		return e.From()
+	}
+	return e.To()
+}
+
+const (
+	fromTo edgeTreeOrder = true
+	toFrom edgeTreeOrder = false
+)
+
+// edgeTreeEntry BTree items for tracking edges
+// in an ordered manner.
+type edgeTreeEntry struct {
+	t    *depEdgeTree
+	edge *DepEdge
+}
+
+func (et *depEdgeTree) insert(e *DepEdge) {
+	et.t.ReplaceOrInsert(&edgeTreeEntry{
+		t:    et,
+		edge: e,
+	})
+}
+
+func (et *depEdgeTree) iterateSourceNode(n *scpb.Node, it DepEdgeIterator) (err error) {
+	e := &edgeTreeEntry{t: et, edge: &DepEdge{}}
+	if et.order == fromTo {
+		e.edge.from = n
+	} else {
+		e.edge.to = n
+	}
+	et.t.AscendGreaterOrEqual(e, func(i btree.Item) (wantMore bool) {
+		e := i.(*edgeTreeEntry)
+		if et.order.first(e.edge) != n {
+			return false
+		}
+		err = it(e.edge)
+		return err == nil
+	})
+	if iterutil.Done(err) {
+		err = nil
+	}
+	return err
+}
+
+// Less implements btree.Item.
+func (e *edgeTreeEntry) Less(otherItem btree.Item) bool {
+	o := otherItem.(*edgeTreeEntry)
+	if less, eq := e.t.cmp(e.t.order.first(e.edge), e.t.order.first(o.edge)); !eq {
+		return less
+	}
+	less, _ := e.t.cmp(e.t.order.second(e.edge), e.t.order.second(o.edge))
+	return less
+}

pkg/sql/schemachanger/scgraph/dep_edge_tree_test.go

@@ -0,0 +1,181 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package scgraph
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
+	"github.com/cockroachdb/cockroach/pkg/util/iterutil"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/stretchr/testify/require"
+)
+
+// TestDepEdgeTree exercises the depEdgeTree data structure to ensure it works
+// as expected.
+func TestDepEdgeTree(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	type nodeID int
+	type edge [2]nodeID
+	// queryCase runs a query to iterate all edges sources at the node with id q.
+	type queryCase struct {
+		q    nodeID
+		take int    // if > 1, indicates a desire to stop early
+		res  []edge // expected results
+	}
+	// testCase describes the edges to be added and the queries to run.
+	type testCase struct {
+		order   edgeTreeOrder
+		edges   []edge
+		queries []queryCase
+	}
+	testCases := []testCase{
+		{
+			order: fromTo,
+			edges: []edge{
+				{2, 4}, {2, 3}, {4, 5}, {1, 2},
+			},
+			queries: []queryCase{
+				{q: 1, res: []edge{{1, 2}}},
+				{q: 2, res: []edge{{2, 3}, {2, 4}}},
+				{q: 2, take: 1, res: []edge{{2, 3}}},
+			},
+		},
+		{
+			order: toFrom,
+			edges: []edge{
+				{2, 4}, {2, 3}, {4, 5}, {1, 2}, {2, 5}, {1, 5},
+			},
+			queries: []queryCase{
+				{q: 1, res: nil},
+				{q: 2, res: []edge{{1, 2}}},
+				{q: 5, res: []edge{{1, 5}, {2, 5}, {4, 5}}},
+				{q: 5, take: 1, res: []edge{{1, 5}}},
+			},
+		},
+	}
+
+	// testCaseState is used for each queryCase in a testCase.
+	type testCaseState struct {
+		tree      *depEdgeTree
+		nodes     []*scpb.Node // nodes with lower indexes sort lower
+		nodesToID map[*scpb.Node]nodeID
+	}
+	makeTestCaseState := func(tc testCase) testCaseState {
+		tcs := testCaseState{
+			nodesToID: make(map[*scpb.Node]nodeID),
+		}
+		target := scpb.Target{}
+		getNode := func(i nodeID) *scpb.Node {
+			if i > nodeID(len(tcs.nodes)-1) {
+				for j := nodeID(len(tcs.nodes)); j <= i; j++ {
+					tcs.nodes = append(tcs.nodes, &scpb.Node{
+						Target: &target,
+						Status: scpb.Status(j),
+					})
+					tcs.nodesToID[tcs.nodes[j]] = j
+				}
+			}
+			return tcs.nodes[i]
+		}
+		tcs.tree = newDepEdgeTree(tc.order, func(a, b *scpb.Node) (less, eq bool) {
+			ai, bi := tcs.nodesToID[a], tcs.nodesToID[b]
+			return ai < bi, ai == bi
+		})
+		for _, e := range tc.edges {
+			tcs.tree.insert(&DepEdge{
+				from: getNode(e[0]),
+				to:   getNode(e[1]),
+			})
+		}
+		return tcs
+	}
+	runQueryCase := func(t *testing.T, tcs testCaseState, qc queryCase) {
+		i := 0
+		var res []edge
+		require.NoError(t, tcs.tree.iterateSourceNode(tcs.nodes[qc.q], func(de *DepEdge) error {
+			if i++; qc.take > 0 && i > qc.take {
+				return iterutil.StopIteration()
+			}
+			res = append(res, edge{
+				tcs.nodesToID[de.From()],
+				tcs.nodesToID[de.To()],
+			})
+			return nil
+		}))
+		require.Equal(t, qc.res, res)
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("%v,%v", tc.order, tc.edges), func(t *testing.T) {
+			tcs := makeTestCaseState(tc)
+			for _, qc := range tc.queries {
+				t.Run(fmt.Sprintf("%d,%d", qc.q, qc.take), func(t *testing.T) {
+					runQueryCase(t, tcs, qc)
+				})
+			}
+		})
+	}
+}
+
+// TestGraphCompareNodes ensures the semantics of (*Graph).compareNodes is sane.
+func TestGraphCompareNodes(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	t1 := scpb.NewTarget(scpb.Target_ADD, &scpb.Table{TableID: 1}, nil)
+	t2 := scpb.NewTarget(scpb.Target_DROP, &scpb.Table{TableID: 2}, nil)
+	mkNode := func(t *scpb.Target, s scpb.Status) *scpb.Node {
+		return &scpb.Node{Target: t, Status: s}
+	}
+	t1ABSENT := mkNode(t1, scpb.Status_ABSENT)
+	t2PUBLIC := mkNode(t2, scpb.Status_PUBLIC)
+	g, err := New(scpb.State{
+		Nodes: []*scpb.Node{t1ABSENT, t2PUBLIC},
+	})
+	targetStr := func(target *scpb.Target) string {
+		switch target {
+		case t1:
+			return "t1"
+		case t2:
+			return "t2"
+		default:
+			panic("unexpected target")
+		}
+	}
+	nodeStr := func(n *scpb.Node) string {
+		if n == nil {
+			return "nil"
+		}
+		return fmt.Sprintf("%s:%s", targetStr(n.Target), n.Status.String())
+	}
+
+	require.NoError(t, err)
+	for _, tc := range []struct {
+		a, b     *scpb.Node
+		less, eq bool
+	}{
+		{a: nil, b: nil, less: false, eq: true},
+		{a: t1ABSENT, b: nil, less: false, eq: false},
+		{a: nil, b: t1ABSENT, less: true, eq: false},
+		{a: t1ABSENT, b: t1ABSENT, less: false, eq: true},
+		{a: t2PUBLIC, b: t1ABSENT, less: false, eq: false},
+		{a: t1ABSENT, b: t2PUBLIC, less: true, eq: false},
+		{a: t1ABSENT, b: mkNode(t1, scpb.Status_PUBLIC), less: true, eq: false},
+		{a: mkNode(t1, scpb.Status_PUBLIC), b: t1ABSENT, less: false, eq: false},
+	} {
+		t.Run(fmt.Sprintf("cmp(%s,%s)", nodeStr(tc.a), nodeStr(tc.b)), func(t *testing.T) {
+			less, eq := g.compareNodes(tc.a, tc.b)
+			require.Equal(t, tc.less, less, "less")
+			require.Equal(t, tc.eq, eq, "eq")
+		})
+	}
+}