From a4fda8d32ac30b787cf2c648daed7f8fe5ef859d Mon Sep 17 00:00:00 2001
From: Jorropo <jorropo.pgm@gmail.com>
Date: Sun, 7 Dec 2025 03:01:17 +0100
Subject: [PATCH] cmd/compile: improve Ctz's limits modeling and add bruteforce
 tests

This make Ctz perfect within the limitations limits can represent.

Change-Id: I1e596d8d01892d1b70031cf03cecc487ce147b38
Reviewed-on: https://go-review.googlesource.com/c/go/+/727780
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Carlos Amedee <carlos@golang.org>
Reviewed-by: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Jorropo <jorropo.pgm@gmail.com>
---
 src/cmd/compile/internal/ssa/prove.go      | 48 +++++++++++++---------
 src/cmd/compile/internal/ssa/prove_test.go | 23 ++++++-----
 2 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/src/cmd/compile/internal/ssa/prove.go b/src/cmd/compile/internal/ssa/prove.go
index 93443a3d3cb..22de2dc7a55 100644
--- a/src/cmd/compile/internal/ssa/prove.go
+++ b/src/cmd/compile/internal/ssa/prove.go
@@ -269,6 +269,13 @@ func convertIntWithBitsize[Target uint64 | int64, Source uint64 | int64](x Sourc
 	}
 }
 
+func (l limit) unsignedFixedLeadingBits() (fixed uint64, count uint) {
+	varying := uint(bits.Len64(l.umin ^ l.umax))
+	count = uint(bits.LeadingZeros64(l.umin ^ l.umax))
+	fixed = l.umin &^ (1<<varying - 1)
+	return
+}
+
 // add returns the limit obtained by adding a value with limit l
 // to a value with limit l2. The result must fit in b bits.
 func (l limit) add(l2 limit, b uint) limit {
@@ -404,6 +411,23 @@ func (l limit) neg(b uint) limit {
 	return l.com(b).add(limit{min: 1, max: 1, umin: 1, umax: 1}, b)
 }
 
+// Similar to add, but computes the TrailingZeros of the limit for bitsize b.
+func (l limit) ctz(b uint) limit {
+	fixed, fixedCount := l.unsignedFixedLeadingBits()
+	if fixedCount == 64 {
+		constResult := min(uint(bits.TrailingZeros64(fixed)), b)
+		return limit{min: int64(constResult), max: int64(constResult), umin: uint64(constResult), umax: uint64(constResult)}
+	}
+
+	varying := 64 - fixedCount
+	if l.umin&((1<<varying)-1) != 0 {
+		// there will always be at least one non-zero bit in the varying part
+		varying--
+		return noLimit.unsignedMax(uint64(varying))
+	}
+	return noLimit.unsignedMax(uint64(min(uint(bits.TrailingZeros64(fixed)), b)))
+}
+
 var noLimit = limit{math.MinInt64, math.MaxInt64, 0, math.MaxUint64}
 
 // a limitFact is a limit known for a particular value.
@@ -1904,26 +1928,10 @@ func (ft *factsTable) flowLimit(v *Value) {
 		}
 
 	// math/bits
-	case OpCtz64:
-		a := ft.limits[v.Args[0].ID]
-		if a.nonzero() {
-			ft.unsignedMax(v, uint64(bits.Len64(a.umax)-1))
-		}
-	case OpCtz32:
-		a := ft.limits[v.Args[0].ID]
-		if a.nonzero() {
-			ft.unsignedMax(v, uint64(bits.Len32(uint32(a.umax))-1))
-		}
-	case OpCtz16:
-		a := ft.limits[v.Args[0].ID]
-		if a.nonzero() {
-			ft.unsignedMax(v, uint64(bits.Len16(uint16(a.umax))-1))
-		}
-	case OpCtz8:
-		a := ft.limits[v.Args[0].ID]
-		if a.nonzero() {
-			ft.unsignedMax(v, uint64(bits.Len8(uint8(a.umax))-1))
-		}
+	case OpCtz64, OpCtz32, OpCtz16, OpCtz8:
+		a := v.Args[0]
+		al := ft.limits[a.ID]
+		ft.newLimit(v, al.ctz(uint(a.Type.Size())*8))
 
 	case OpPopCount64, OpPopCount32, OpPopCount16, OpPopCount8:
 		a := ft.limits[v.Args[0].ID]
diff --git a/src/cmd/compile/internal/ssa/prove_test.go b/src/cmd/compile/internal/ssa/prove_test.go
index 6315049870a..1d91736ff00 100644
--- a/src/cmd/compile/internal/ssa/prove_test.go
+++ b/src/cmd/compile/internal/ssa/prove_test.go
@@ -6,11 +6,11 @@ package ssa
 
 import (
 	"math"
+	"math/bits"
 	"testing"
 )
 
-func testLimitUnaryOpSigned8(t *testing.T, opName string, op func(l limit, bitsize uint) limit, opImpl func(int8) int8) {
-	sizeLimit := noLimitForBitsize(8)
+func testLimitUnaryOpSigned8(t *testing.T, opName string, initLimit limit, op func(l limit, bitsize uint) limit, opImpl func(int8) int8) {
 	for min := math.MinInt8; min <= math.MaxInt8; min++ {
 		for max := min; max <= math.MaxInt8; max++ {
 			realSmallest, realBiggest := int8(math.MaxInt8), int8(math.MinInt8)
@@ -26,7 +26,7 @@ func testLimitUnaryOpSigned8(t *testing.T, opName string, op func(l limit, bitsi
 
 			l := limit{int64(min), int64(max), 0, math.MaxUint64}
 			l = op(l, 8)
-			l = l.intersect(sizeLimit) // We assume this is gonna be used by newLimit which is seeded by the op size already.
+			l = l.intersect(initLimit) // We assume this is gonna be used by newLimit which is seeded by the op size already.
 
 			if l.min != int64(realSmallest) || l.max != int64(realBiggest) {
 				t.Errorf("%s(%d..%d) = %d..%d; want %d..%d", opName, min, max, l.min, l.max, realSmallest, realBiggest)
@@ -35,8 +35,7 @@ func testLimitUnaryOpSigned8(t *testing.T, opName string, op func(l limit, bitsi
 	}
 }
 
-func testLimitUnaryOpUnsigned8(t *testing.T, opName string, op func(l limit, bitsize uint) limit, opImpl func(uint8) uint8) {
-	sizeLimit := noLimitForBitsize(8)
+func testLimitUnaryOpUnsigned8(t *testing.T, opName string, initLimit limit, op func(l limit, bitsize uint) limit, opImpl func(uint8) uint8) {
 	for min := 0; min <= math.MaxUint8; min++ {
 		for max := min; max <= math.MaxUint8; max++ {
 			realSmallest, realBiggest := uint8(math.MaxUint8), uint8(0)
@@ -52,7 +51,7 @@ func testLimitUnaryOpUnsigned8(t *testing.T, opName string, op func(l limit, bit
 
 			l := limit{math.MinInt64, math.MaxInt64, uint64(min), uint64(max)}
 			l = op(l, 8)
-			l = l.intersect(sizeLimit) // We assume this is gonna be used by newLimit which is seeded by the op size already.
+			l = l.intersect(initLimit) // We assume this is gonna be used by newLimit which is seeded by the op size already.
 
 			if l.umin != uint64(realSmallest) || l.umax != uint64(realBiggest) {
 				t.Errorf("%s(%d..%d) = %d..%d; want %d..%d", opName, min, max, l.umin, l.umax, realSmallest, realBiggest)
@@ -62,15 +61,19 @@ func testLimitUnaryOpUnsigned8(t *testing.T, opName string, op func(l limit, bit
 }
 
 func TestLimitNegSigned(t *testing.T) {
-	testLimitUnaryOpSigned8(t, "neg", limit.neg, func(x int8) int8 { return -x })
+	testLimitUnaryOpSigned8(t, "neg", noLimitForBitsize(8), limit.neg, func(x int8) int8 { return -x })
 }
 func TestLimitNegUnsigned(t *testing.T) {
-	testLimitUnaryOpUnsigned8(t, "neg", limit.neg, func(x uint8) uint8 { return -x })
+	testLimitUnaryOpUnsigned8(t, "neg", noLimitForBitsize(8), limit.neg, func(x uint8) uint8 { return -x })
 }
 
 func TestLimitComSigned(t *testing.T) {
-	testLimitUnaryOpSigned8(t, "com", limit.com, func(x int8) int8 { return ^x })
+	testLimitUnaryOpSigned8(t, "com", noLimitForBitsize(8), limit.com, func(x int8) int8 { return ^x })
 }
 func TestLimitComUnsigned(t *testing.T) {
-	testLimitUnaryOpUnsigned8(t, "com", limit.com, func(x uint8) uint8 { return ^x })
+	testLimitUnaryOpUnsigned8(t, "com", noLimitForBitsize(8), limit.com, func(x uint8) uint8 { return ^x })
+}
+
+func TestLimitCtzUnsigned(t *testing.T) {
+	testLimitUnaryOpUnsigned8(t, "ctz", limit{-128, 127, 0, 8}, limit.ctz, func(x uint8) uint8 { return uint8(bits.TrailingZeros8(x)) })
 }