From eb2380ddd8fd637cd11a4a039f18b29778f55684 Mon Sep 17 00:00:00 2001
From: Arul Ajmani <arulajmani@gmail.com>
Date: Fri, 28 Apr 2023 15:21:59 -0400
Subject: [PATCH] lockspanset: cleanup some loops

Followup from https://github.com/cockroachdb/cockroach/pull/102396.
This patch converts some loops to be more idiomatic for go that that
patch mistakenly didn't.

Epic: none

Release note: None
---
 pkg/kv/kvserver/concurrency/lock/locking.go |  7 +++++++
 pkg/kv/kvserver/lockspanset/BUILD.bazel     |  1 +
 pkg/kv/kvserver/lockspanset/lockspanset.go  | 14 ++++++--------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/pkg/kv/kvserver/concurrency/lock/locking.go b/pkg/kv/kvserver/concurrency/lock/locking.go
index 0747620ccfa8..6d477a82d06e 100644
--- a/pkg/kv/kvserver/concurrency/lock/locking.go
+++ b/pkg/kv/kvserver/concurrency/lock/locking.go
@@ -53,6 +53,9 @@ var ExclusiveLocksBlockNonLockingReads = settings.RegisterBoolSetting(
 // MaxStrength is the maximum value in the Strength enum.
 const MaxStrength = Intent
 
+// NumLockStrength is the total number of lock strengths in the Strength enum.
+const NumLockStrength = MaxStrength + 1
+
 // MaxDurability is the maximum value in the Durability enum.
 const MaxDurability = Unreplicated
 
@@ -62,6 +65,10 @@ func init() {
 			panic(fmt.Sprintf("Strength (%s) with value larger than MaxDurability", st))
 		}
 	}
+	if int(NumLockStrength) != len(Strength_name) {
+		panic(fmt.Sprintf("mismatched numer of lock strengths: NumLockStrength %d, lock strengths %d",
+			int(NumLockStrength), len(Strength_name)))
+	}
 	for v := range Durability_name {
 		if d := Durability(v); d > MaxDurability {
 			panic(fmt.Sprintf("Durability (%s) with value larger than MaxDurability", d))
diff --git a/pkg/kv/kvserver/lockspanset/BUILD.bazel b/pkg/kv/kvserver/lockspanset/BUILD.bazel
index 65aa3c825c65..b9d70f4e2f48 100644
--- a/pkg/kv/kvserver/lockspanset/BUILD.bazel
+++ b/pkg/kv/kvserver/lockspanset/BUILD.bazel
@@ -18,6 +18,7 @@ go_test(
     args = ["-test.timeout=295s"],
     embed = [":lockspanset"],
     deps = [
+        "//pkg/keys",
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/roachpb",
         "//pkg/util/leaktest",
diff --git a/pkg/kv/kvserver/lockspanset/lockspanset.go b/pkg/kv/kvserver/lockspanset/lockspanset.go
index 4fe1b2e9fa2d..5f09b4970d1a 100644
--- a/pkg/kv/kvserver/lockspanset/lockspanset.go
+++ b/pkg/kv/kvserver/lockspanset/lockspanset.go
@@ -19,10 +19,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 )
 
-const NumLockStrength = lock.MaxStrength + 1
-
 type LockSpanSet struct {
-	spans [NumLockStrength][]roachpb.Span
+	spans [lock.NumLockStrength][]roachpb.Span
 }
 
 var lockSpanSetPool = sync.Pool{
@@ -78,10 +76,10 @@ func (l *LockSpanSet) Empty() bool {
 // String prints a string representation of the LockSpanSet.
 func (l *LockSpanSet) String() string {
 	var buf strings.Builder
-	for st := lock.Strength(0); st < NumLockStrength; st++ {
-		for _, span := range l.GetSpans(st) {
+	for st, spans := range l.spans {
+		for _, span := range spans {
 			fmt.Fprintf(&buf, "%s: %s\n",
-				st, span)
+				lock.Strength(st), span)
 		}
 	}
 	return buf.String()
@@ -90,8 +88,8 @@ func (l *LockSpanSet) String() string {
 // Len returns the total number of spans tracked across all strengths.
 func (l *LockSpanSet) Len() int {
 	var count int
-	for st := lock.Strength(0); st < NumLockStrength; st++ {
-		count += len(l.GetSpans(st))
+	for _, spans := range l.spans {
+		count += len(spans)
 	}
 	return count
 }