types/logid: add Add method (#12478)

The Add method derives a new ID by adding a signed integer
to the ID, treating it as an unsigned 256-bit big-endian integer.

We also add Less and Compare methods to PrivateID to provide
feature parity with existing methods on PublicID.

Updates tailscale/corp#11038

Signed-off-by: Joe Tsai <joetsai@digital-static.net>

types/logid/id.go

@@ -11,6 +11,7 @@
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
+	"math/bits"
 	"slices"
 	"unicode/utf8"
 )
@@ -38,6 +39,12 @@ func ParsePrivateID(in string) (out PrivateID, err error) {
 	return out, err
 }
 
+// Add adds i to the id, treating it as an unsigned 256-bit big-endian integer,
+// and returns the resulting ID.
+func (id PrivateID) Add(i int64) PrivateID {
+	return add(id, i)
+}
+
 func (id PrivateID) AppendText(b []byte) ([]byte, error) {
 	return hex.AppendEncode(b, id[:]), nil
 }
@@ -54,6 +61,14 @@ func (id PrivateID) String() string {
 	return string(hex.AppendEncode(nil, id[:]))
 }
 
+func (id1 PrivateID) Less(id2 PrivateID) bool {
+	return id1.Compare(id2) < 0
+}
+
+func (id1 PrivateID) Compare(id2 PrivateID) int {
+	return slices.Compare(id1[:], id2[:])
+}
+
 func (id PrivateID) IsZero() bool {
 	return id == PrivateID{}
 }
@@ -74,6 +89,12 @@ func ParsePublicID(in string) (out PublicID, err error) {
 	return out, err
 }
 
+// Add adds i to the id, treating it as an unsigned 256-bit big-endian integer,
+// and returns the resulting ID.
+func (id PublicID) Add(i int64) PublicID {
+	return add(id, i)
+}
+
 func (id PublicID) AppendText(b []byte) ([]byte, error) {
 	return hex.AppendEncode(b, id[:]), nil
 }
@@ -118,3 +139,22 @@ func parseID[Bytes []byte | string](funcName string, out *[32]byte, in Bytes) (e
 	}
 	return nil
 }
+
+func add(id [32]byte, i int64) [32]byte {
+	var out uint64
+	switch {
+	case i < 0:
+		borrow := ^uint64(i) + 1 // twos-complement inversion
+		for i := 0; i < 4 && borrow > 0; i++ {
+			out, borrow = bits.Sub64(binary.BigEndian.Uint64(id[8*(3-i):]), borrow, 0)
+			binary.BigEndian.PutUint64(id[8*(3-i):], out)
+		}
+	case i > 0:
+		carry := uint64(i)
+		for i := 0; i < 4 && carry > 0; i++ {
+			out, carry = bits.Add64(binary.BigEndian.Uint64(id[8*(3-i):]), carry, 0)
+			binary.BigEndian.PutUint64(id[8*(3-i):], out)
+		}
+	}
+	return id
+}

types/logid/id_test.go

@@ -4,9 +4,11 @@
 package logid
 
 import (
+	"math"
 	"testing"
 
 	"tailscale.com/tstest"
+	"tailscale.com/util/must"
 )
 
 func TestIDs(t *testing.T) {
@@ -77,3 +79,89 @@ func TestIDs(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestAdd(t *testing.T) {
+	tests := []struct {
+		in   string
+		add  int64
+		want string
+	}{{
+		in:   "0000000000000000000000000000000000000000000000000000000000000000",
+		add:  0,
+		want: "0000000000000000000000000000000000000000000000000000000000000000",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000000000000000000000",
+		add:  1,
+		want: "0000000000000000000000000000000000000000000000000000000000000001",
+	}, {
+		in:   "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
+		add:  1,
+		want: "0000000000000000000000000000000000000000000000000000000000000000",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000000000000000000000",
+		add:  -1,
+		want: "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000000000000000000000",
+		add:  math.MinInt64,
+		want: "ffffffffffffffffffffffffffffffffffffffffffffffff8000000000000000",
+	}, {
+		in:   "000000000000000000000000000000000000000000000000ffffffffffffffff",
+		add:  math.MinInt64,
+		want: "0000000000000000000000000000000000000000000000007fffffffffffffff",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000000000000000000000",
+		add:  math.MaxInt64,
+		want: "0000000000000000000000000000000000000000000000007fffffffffffffff",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000007fffffffffffffff",
+		add:  math.MaxInt64,
+		want: "000000000000000000000000000000000000000000000000fffffffffffffffe",
+	}, {
+		in:   "000000000000000000000000000000000000000000000000ffffffffffffffff",
+		add:  1,
+		want: "0000000000000000000000000000000000000000000000010000000000000000",
+	}, {
+		in:   "00000000000000000000000000000000fffffffffffffffffffffffffffffffe",
+		add:  3,
+		want: "0000000000000000000000000000000100000000000000000000000000000001",
+	}, {
+		in:   "0000000000000000fffffffffffffffffffffffffffffffffffffffffffffffd",
+		add:  5,
+		want: "0000000000000001000000000000000000000000000000000000000000000002",
+	}, {
+		in:   "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffc",
+		add:  7,
+		want: "0000000000000000000000000000000000000000000000000000000000000003",
+	}, {
+		in:   "ffffffffffffffffffffffffffffffffffffffffffffffff0000000000000000",
+		add:  -1,
+		want: "fffffffffffffffffffffffffffffffffffffffffffffffeffffffffffffffff",
+	}, {
+		in:   "ffffffffffffffffffffffffffffffff00000000000000000000000000000001",
+		add:  -3,
+		want: "fffffffffffffffffffffffffffffffefffffffffffffffffffffffffffffffe",
+	}, {
+		in:   "ffffffffffffffff000000000000000000000000000000000000000000000002",
+		add:  -5,
+		want: "fffffffffffffffefffffffffffffffffffffffffffffffffffffffffffffffd",
+	}, {
+		in:   "0000000000000000000000000000000000000000000000000000000000000003",
+		add:  -7,
+		want: "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffc",
+	}}
+	for _, tt := range tests {
+		in := must.Get(ParsePublicID(tt.in))
+		want := must.Get(ParsePublicID(tt.want))
+		got := in.Add(tt.add)
+		if got != want {
+			t.Errorf("%s.Add(%d):\n\tgot  %s\n\twant %s", in, tt.add, got, want)
+		}
+		if tt.add != math.MinInt64 {
+			got = got.Add(-tt.add)
+			if got != in {
+				t.Errorf("%s.Add(%d):\n\tgot  %s\n\twant %s", want, -tt.add, got, in)
+			}
+		}
+	}
+}