diff --git a/crrand/perm.go b/crrand/perm.go new file mode 100644 index 0000000..3f2a46e --- /dev/null +++ b/crrand/perm.go @@ -0,0 +1,81 @@ +// Copyright 2025 The Cockroach 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 crrand implements functionality related to pseudorandom number +// generation. +package crrand + +import "math/bits" + +// MakePerm64 constructs a new Perm64 from a 64-bit seed, providing a +// deterministic, pseudorandom, bijective mapping of 64-bit values X to 64-bit +// values Y. +func MakePerm64(seed uint64) Perm64 { + // derive 4 x 32-bit round keys from the 64-bit seed using only ARX ops. + const c0 = 0x9E3779B97F4A7C15 // golden ratio (used here as XOR salt) + const c1 = 0xC2B2AE3D27D4EB4F // a constant + + var m Perm64 + s0 := seed + s1 := bits.RotateLeft64(seed^c0, 13) + s2 := bits.RotateLeft64(seed^c1, 37) + s3 := bits.RotateLeft64(seed^c0^c1, 53) + + m.seed[0] = uint32(s0) + m.seed[1] = uint32(s1 >> 32) + m.seed[2] = uint32(s2) + m.seed[3] = uint32(s3 >> 32) + return m +} + +// A Perm64 provides a deterministic, pseudorandom permutation of 64-bit values. +type Perm64 struct { + seed [4]uint32 +} + +// At returns the nth value in the permutation of the 64-bit values. The return +// value may be passed to Index to recover n. The permutation is pseudorandom. +func (p Perm64) At(n uint64) uint64 { + // Use a simple Feistel network with 4 rounds to shuffle data. + L := uint32(n >> 32) + R := uint32(n) + for r := range p.seed { + t := arx(R^p.seed[r], p.seed[(r+1)&3]) + L, R = R, L^t + } + return (uint64(L) << 32) | uint64(R) +} + +// IndexOf inverts the permutation, returning the index of the provided value in +// the permutation. If y was produced by At(x), then IndexOf(y) returns x. +func (p Perm64) IndexOf(y uint64) uint64 { + L := uint32(y >> 32) + R := uint32(y) + for r := 3; r >= 0; r-- { + // reverse of: L, R = R, L ^ arx(R^k[r], k[(r+1)&3]) + prevR := L + prevL := R ^ arx(prevR^p.seed[r], p.seed[(r+1)&3]) + L, R = prevL, prevR + } + return (uint64(L) << 32) | uint64(R) +} + +// ARX-only round function. +func arx(x, k uint32) uint32 { + x ^= k + x += bits.RotateLeft32(x, 5) + x ^= bits.RotateLeft32(x, 7) + x += bits.RotateLeft32(x, 16) + return x +} diff --git a/crrand/perm_test.go b/crrand/perm_test.go new file mode 100644 index 0000000..c9672c5 --- /dev/null +++ b/crrand/perm_test.go @@ -0,0 +1,60 @@ +// Copyright 2025 The Cockroach 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 crrand + +import ( + "math" + "math/rand/v2" + "testing" + "time" +) + +var interestingUint64s = []uint64{ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 32, 33, 63, 64, 65, 129, 3050, 29356, + 297532935, 2973539791203, 0x9E3779B97F4A7C15, math.MaxUint64 - 1, + math.MaxUint64, +} + +func TestPerm64(t *testing.T) { + for _, seed := range interestingUint64s { + p := MakePerm64(seed) + for _, x := range interestingUint64s { + y := p.At(x) + x2 := p.IndexOf(y) + if x != x2 { + t.Errorf(".At(%d) = %d, .IndexOf(%d) = %d, want %d", x, y, y, x2, x) + } + } + } +} + +func TestPerm64Random(t *testing.T) { + seed := uint64(time.Now().UnixNano()) + defer func() { + if t.Failed() { + t.Logf("seed: %d", seed) + } + }() + rng := rand.New(rand.NewPCG(seed, seed)) + p := MakePerm64(rng.Uint64()) + for i := 0; i < 1000; i++ { + x := rng.Uint64() + y := p.At(x) + x2 := p.IndexOf(y) + if x != x2 { + t.Errorf("p.At(%d) = %d, p.IndexOf(%d) = %d, want %d", x, y, y, x2, x) + } + } +}