bits: use normal Go generics

FUTURE_COPYBARA_INTEGRATE_REVIEW=#12482 from tamird:bits-generics 23c4f2e
PiperOrigin-RevId: 855373016

Author: tamird

pkg/abi/linux/signal.go

@@ -130,12 +130,12 @@ func MakeSignalSet(sigs ...Signal) SignalSet {
 	for i, sig := range sigs {
 		indices[i] = sig.Index()
 	}
-	return SignalSet(bits.Mask64(indices...))
+	return bits.Mask[SignalSet](indices...)
 }
 
 // SignalSetOf returns a SignalSet with a single signal set.
 func SignalSetOf(sig Signal) SignalSet {
-	return SignalSet(bits.MaskOf64(sig.Index()))
+	return bits.MaskOf[SignalSet](sig.Index())
 }
 
 // ForEachSignal invokes f for each signal set in the given mask.

pkg/bits/BUILD

@@ -1,5 +1,4 @@
 load("//tools:defs.bzl", "go_library", "go_test")
-load("//tools/go_generics:defs.bzl", "go_template", "go_template_instance")
 
 package(
     default_applicable_licenses = ["//:license"],
@@ -10,46 +9,18 @@ go_library(
     name = "bits",
     srcs = [
         "bits.go",
-        "bits32.go",
-        "bits64.go",
+        "bits_template.go",
         "uint64_arch.go",
         "uint64_arch_amd64_asm.s",
         "uint64_arch_arm64_asm.s",
         "uint64_arch_generic.go",
     ],
     visibility = ["//:sandbox"],
-)
-
-go_template(
-    name = "bits_template",
-    srcs = ["bits_template.go"],
-    types = [
-        "T",
+    deps = [
+        "@org_golang_x_exp//constraints:go_default_library",
     ],
 )
 
-go_template_instance(
-    name = "bits64",
-    out = "bits64.go",
-    package = "bits",
-    suffix = "64",
-    template = ":bits_template",
-    types = {
-        "T": "uint64",
-    },
-)
-
-go_template_instance(
-    name = "bits32",
-    out = "bits32.go",
-    package = "bits",
-    suffix = "32",
-    template = ":bits_template",
-    types = {
-        "T": "uint32",
-    },
-)
-
 go_test(
     name = "bits_test",
     size = "small",

pkg/bits/bits_template.go

@@ -14,37 +14,36 @@
 
 package bits
 
-// Non-atomic bit operations on a template type T.
+import "golang.org/x/exp/constraints"
 
-// T is a required type parameter that must be an integral type.
-type T uint64
+// Non-atomic bit operations on integral types.
 
 // IsOn returns true if *all* bits set in 'bits' are set in 'mask'.
-func IsOn(mask, bits T) bool {
+func IsOn[T constraints.Integer](mask, bits T) bool {
 	return mask&bits == bits
 }
 
 // IsAnyOn returns true if *any* bit set in 'bits' is set in 'mask'.
-func IsAnyOn(mask, bits T) bool {
+func IsAnyOn[T constraints.Integer](mask, bits T) bool {
 	return mask&bits != 0
 }
 
 // Mask returns a T with all of the given bits set.
-func Mask(is ...int) T {
+func Mask[T constraints.Integer](is ...int) T {
 	ret := T(0)
 	for _, i := range is {
-		ret |= MaskOf(i)
+		ret |= MaskOf[T](i)
 	}
 	return ret
 }
 
 // MaskOf is like Mask, but sets only a single bit (more efficiently).
-func MaskOf(i int) T {
+func MaskOf[T constraints.Integer](i int) T {
 	return T(1) << T(i)
 }
 
 // IsPowerOfTwo returns true if v is power of 2.
-func IsPowerOfTwo(v T) bool {
+func IsPowerOfTwo[T constraints.Integer](v T) bool {
 	if v == 0 {
 		return false
 	}

pkg/bits/uint64_arch.go

@@ -32,6 +32,6 @@ func ForEachSetBit64(x uint64, f func(i int)) {
 	for x != 0 {
 		i := TrailingZeros64(x)
 		f(i)
-		x &^= MaskOf64(i)
+		x &^= MaskOf[uint64](i)
 	}
 }

pkg/bits/uint64_test.go

@@ -65,6 +65,11 @@ func TestMostSignificantOne64(t *testing.T) {
 	}
 }
 
+// Mask64 returns a uint64 with all of the given bits set.
+func Mask64(is ...int) uint64 {
+	return Mask[uint64](is...)
+}
+
 func TestForEachSetBit64(t *testing.T) {
 	for _, want := range [][]int{
 		{},
@@ -106,10 +111,10 @@ func TestIsOn(t *testing.T) {
 		{Mask64(1, 63), Mask64(0, 1, 63), true, false},
 		{Mask64(1, 63), Mask64(0, 62), false, false},
 	} {
-		if ok := IsAnyOn64(s.mask, s.bits); ok != s.any {
+		if ok := IsAnyOn(s.mask, s.bits); ok != s.any {
 			t.Errorf("IsAnyOn(%#x, %#x) = %v, wanted: %v", s.mask, s.bits, ok, s.any)
 		}
-		if ok := IsOn64(s.mask, s.bits); ok != s.all {
+		if ok := IsOn(s.mask, s.bits); ok != s.all {
 			t.Errorf("IsOn(%#x, %#x) = %v, wanted: %v", s.mask, s.bits, ok, s.all)
 		}
 	}
@@ -127,7 +132,7 @@ func TestIsPowerOfTwo(t *testing.T) {
 		{v: 4, want: true},
 		{v: 5, want: false},
 	} {
-		if got := IsPowerOfTwo64(tc.v); got != tc.want {
+		if got := IsPowerOfTwo(tc.v); got != tc.want {
 			t.Errorf("IsPowerOfTwo(%d) = %t, want: %t", tc.v, got, tc.want)
 		}
 	}

pkg/sentry/kernel/auth/capability_set.go

@@ -31,16 +31,16 @@ var AllCapabilities = CapabilitySetOf(linux.CAP_LAST_CAP+1) - 1
 // CapabilitySetOf returns a CapabilitySet containing only the given
 // capability.
 func CapabilitySetOf(cp linux.Capability) CapabilitySet {
-	return CapabilitySet(bits.MaskOf64(int(cp)))
+	return bits.MaskOf[CapabilitySet](int(cp))
 }
 
 // CapabilitySetOfMany returns a CapabilitySet containing the given capabilities.
 func CapabilitySetOfMany(cps []linux.Capability) CapabilitySet {
-	var cs uint64
+	var cs CapabilitySet
 	for _, cp := range cps {
-		cs |= bits.MaskOf64(int(cp))
+		cs |= bits.MaskOf[CapabilitySet](int(cp))
 	}
-	return CapabilitySet(cs)
+	return cs
 }
 
 // Add adds the given capability to the CapabilitySet.

pkg/sentry/kernel/syscalls.go

@@ -165,7 +165,7 @@ func (e *SyscallFlagsTable) UpdateSecCheck(state *seccheck.State) {
 	defer e.mu.Unlock()
 	for sysno := uintptr(0); sysno <= sentry.MaxSyscallNum; sysno++ {
 		oldFlags := e.enable[sysno].Load()
-		if !bits.IsOn32(oldFlags, syscallPresent) {
+		if !bits.IsOn(oldFlags, syscallPresent) {
 			continue
 		}
 		flags := oldFlags
@@ -226,7 +226,7 @@ func (e *SyscallFlagsTable) Enable(bit uint32, s map[uintptr]bool, missingEnable
 
 	for num := range e.enable {
 		val := e.enable[num].Load()
-		if !bits.IsOn32(val, syscallPresent) {
+		if !bits.IsOn(val, syscallPresent) {
 			// Missing.
 			e.enable[num].Store(missingVal)
 			continue
@@ -252,7 +252,7 @@ func (e *SyscallFlagsTable) EnableAll(bit uint32) {
 
 	for num := range e.enable {
 		val := e.enable[num].Load()
-		if !bits.IsOn32(val, syscallPresent) {
+		if !bits.IsOn(val, syscallPresent) {
 			// Missing.
 			e.enable[num].Store(missingVal)
 			continue

pkg/sentry/kernel/task_syscall.go

@@ -87,11 +87,11 @@ func (t *Task) executeSyscall(sysno uintptr, args arch.SyscallArguments) (rval u
 	fe := s.FeatureEnable.Word(sysno)
 
 	var straceContext any
-	if bits.IsAnyOn32(fe, StraceEnableBits) {
+	if bits.IsAnyOn(fe, StraceEnableBits) {
 		straceContext = s.Stracer.SyscallEnter(t, sysno, args, fe)
 	}
 
-	if bits.IsAnyOn32(fe, SecCheckRawEnter) {
+	if bits.IsAnyOn(fe, SecCheckRawEnter) {
 		info := pb.Syscall{
 			Sysno: uint64(sysno),
 			Arg1:  args[0].Uint64(),
@@ -110,7 +110,7 @@ func (t *Task) executeSyscall(sysno uintptr, args arch.SyscallArguments) (rval u
 			return c.RawSyscall(t, fields, &info)
 		})
 	}
-	if bits.IsAnyOn32(fe, SecCheckEnter) {
+	if bits.IsAnyOn(fe, SecCheckEnter) {
 		fields := seccheck.Global.GetFieldSet(seccheck.GetPointForSyscall(seccheck.SyscallEnter, sysno))
 		var ctxData *pb.ContextData
 		if !fields.Context.Empty() {
@@ -128,7 +128,7 @@ func (t *Task) executeSyscall(sysno uintptr, args arch.SyscallArguments) (rval u
 		})
 	}
 
-	if bits.IsOn32(fe, ExternalBeforeEnable) && (s.ExternalFilterBefore == nil || s.ExternalFilterBefore(t, sysno, args)) {
+	if bits.IsOn(fe, ExternalBeforeEnable) && (s.ExternalFilterBefore == nil || s.ExternalFilterBefore(t, sysno, args)) {
 		t.invokeExternal()
 		// Ensure we check for stops, then invoke the syscall again.
 		ctrl = ctrlStopAndReinvokeSyscall
@@ -150,16 +150,16 @@ func (t *Task) executeSyscall(sysno uintptr, args arch.SyscallArguments) (rval u
 		}
 	}
 
-	if bits.IsOn32(fe, ExternalAfterEnable) && (s.ExternalFilterAfter == nil || s.ExternalFilterAfter(t, sysno, args)) {
+	if bits.IsOn(fe, ExternalAfterEnable) && (s.ExternalFilterAfter == nil || s.ExternalFilterAfter(t, sysno, args)) {
 		t.invokeExternal()
 		// Don't reinvoke the unix.
 	}
 
-	if bits.IsAnyOn32(fe, StraceEnableBits) {
+	if bits.IsAnyOn(fe, StraceEnableBits) {
 		s.Stracer.SyscallExit(straceContext, t, sysno, rval, err)
 	}
 
-	if bits.IsAnyOn32(fe, SecCheckRawExit) {
+	if bits.IsAnyOn(fe, SecCheckRawExit) {
 		info := pb.Syscall{
 			Sysno: uint64(sysno),
 			Arg1:  args[0].Uint64(),
@@ -182,7 +182,7 @@ func (t *Task) executeSyscall(sysno uintptr, args arch.SyscallArguments) (rval u
 			return c.RawSyscall(t, fields, &info)
 		})
 	}
-	if bits.IsAnyOn32(fe, SecCheckExit) {
+	if bits.IsAnyOn(fe, SecCheckExit) {
 		fields := seccheck.Global.GetFieldSet(seccheck.GetPointForSyscall(seccheck.SyscallExit, sysno))
 		var ctxData *pb.ContextData
 		if !fields.Context.Empty() {

pkg/sentry/strace/strace.go

@@ -682,10 +682,10 @@ func (s SyscallMap) SyscallEnter(t *kernel.Task, sysno uintptr, args arch.Syscal
 	}
 
 	var output, eventOutput []string
-	if bits.IsOn32(flags, kernel.StraceEnableLog) {
+	if bits.IsOn(flags, kernel.StraceEnableLog) {
 		output = info.printEnter(t, args)
 	}
-	if bits.IsOn32(flags, kernel.StraceEnableEvent) {
+	if bits.IsOn(flags, kernel.StraceEnableEvent) {
 		eventOutput = info.sendEnter(t, args)
 	}
 
@@ -706,10 +706,10 @@ func (s SyscallMap) SyscallExit(context any, t *kernel.Task, sysno, rval uintptr
 	c := context.(*syscallContext)
 
 	elapsed := time.Since(c.start)
-	if bits.IsOn32(c.flags, kernel.StraceEnableLog) {
+	if bits.IsOn(c.flags, kernel.StraceEnableLog) {
 		c.info.printExit(t, elapsed, c.logOutput, c.args, rval, err, errno)
 	}
-	if bits.IsOn32(c.flags, kernel.StraceEnableEvent) {
+	if bits.IsOn(c.flags, kernel.StraceEnableEvent) {
 		c.info.sendExit(t, elapsed, c.eventOutput, c.args, rval, err, errno)
 	}
 }
@@ -786,10 +786,10 @@ const (
 
 func convertToSyscallFlag(sinks SinkType) uint32 {
 	ret := uint32(0)
-	if bits.IsOn32(uint32(sinks), uint32(SinkTypeLog)) {
+	if bits.IsOn(uint32(sinks), uint32(SinkTypeLog)) {
 		ret |= kernel.StraceEnableLog
 	}
-	if bits.IsOn32(uint32(sinks), uint32(SinkTypeEvent)) {
+	if bits.IsOn(uint32(sinks), uint32(SinkTypeEvent)) {
 		ret |= kernel.StraceEnableEvent
 	}
 	return ret

pkg/sentry/syscalls/linux/sys_stat.go

@@ -165,7 +165,7 @@ func Statx(t *kernel.Task, sysno uintptr, args arch.SyscallArguments) (uintptr,
 	}
 	// Make sure that only one sync type option is set.
 	syncType := uint32(flags & linux.AT_STATX_SYNC_TYPE)
-	if syncType != 0 && !bits.IsPowerOfTwo32(syncType) {
+	if syncType != 0 && !bits.IsPowerOfTwo(syncType) {
 		return 0, nil, linuxerr.EINVAL
 	}
 	if mask&linux.STATX__RESERVED != 0 {