diff --git a/crmath/scale.go b/crmath/scale.go new file mode 100644 index 0000000..047b138 --- /dev/null +++ b/crmath/scale.go @@ -0,0 +1,45 @@ +// 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 crmath + +import ( + "math" + "math/bits" +) + +// ScaleUint64 returns the ceiling of ax/b, clamped to the uint64 range. Panics +// if b is zero. +// +// More precisely, returns: +// +// min( ⌈ x*a / b ⌉, math.MaxUint64 ) +func ScaleUint64(x uint64, a uint64, b uint64) uint64 { + var quo, rem uint64 + if x < math.MaxUint32 && a < math.MaxUint32 { + prod := a * x + quo = prod / b + rem = prod % b + } else { + hi, lo := bits.Mul64(x, a) + if hi >= b && b != 0 { + return math.MaxUint64 + } + quo, rem = bits.Div64(hi, lo, b) + } + if rem == 0 { + return quo + } + return quo + 1 +} diff --git a/crmath/scale_test.go b/crmath/scale_test.go new file mode 100644 index 0000000..2d078a1 --- /dev/null +++ b/crmath/scale_test.go @@ -0,0 +1,90 @@ +// 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 crmath + +import ( + "math" + "math/big" + "math/rand/v2" + "testing" +) + +func TestScaleUint64(t *testing.T) { + expect := func(x, a, b, expected uint64) { + if got := ScaleUint64(x, a, b); got != expected { + t.Helper() + t.Fatalf("ScaleUint64(%d, %d, %d) = %d; want %d", x, a, b, got, expected) + } + } + + expect(0, 1, 1, 0) + + expect(1, 1, 123456, 1) + expect(1, 1, math.MaxUint64, 1) + + expect(1000, 1, 10, 100) + + expect(1<<52, 1<<40, 1<<60, 1<<32) + expect(1<<52, 1<<60, 1<<50, 1<<62) + expect(1<<52, 1<<60, 1<<40, math.MaxUint64) + + for range 100 { + n := rand.Uint64() + expect(n, 1, 1, n) + expect(n, math.MaxUint64, math.MaxUint64, n) + m := rand.Uint64() + expect(n, m, m, n) + } + + // calc is an alternative implementation using big.Int. + calc := func(xx, aa, bb uint64) uint64 { + var x, a, b, numerator, res big.Int + x.SetUint64(xx) + a.SetUint64(aa) + b.SetUint64(bb) + numerator.Mul(&a, &x) + numerator.Add(&numerator, &b) + numerator.Add(&numerator, big.NewInt(-1)) + res.Quo(&numerator, &b) + if res.BitLen() > 64 { + return math.MaxUint64 + } + return res.Uint64() + } + + checkDivByZero := func(x, a uint64) { + defer func() { + if r := recover(); r == nil { + t.Fatalf("ScaleUint64(%d, %d, 0) did not panic", x, a) + } + }() + _ = ScaleUint64(x, a, 0) + } + + for range 1000 { + r := func() uint64 { + if rand.IntN(2) == 0 { + // Return values close to the fast path cutoff. + return math.MaxUint32 + rand.Uint64N(5) - 2 + } + return rand.Uint64() + } + x, a, b := r(), r(), r() + expect(x, a, b, calc(x, a, b)) + + checkDivByZero(x, a) + } + +}