cmd/compile: better optimization of type switches in instantiated generic functions

[ianlancetaylor]

Consider this test case:

package main

import (
	"fmt"
	"unsafe"
)

func start[S []byte | string](s S) *byte {
	switch s := any(s).(type) {
	case []byte:
		return unsafe.SliceData(s)
	case string:
		return unsafe.StringData(s)
	default:
		panic("can't happen")
	}
}

func F[S []byte | string](s S) {
	fmt.Println(start(s), len(s))
}

func main() {
	F("hi")
	F([]byte("bye"))
}

This runs fine. It instantiates F twice: main.F[go.shape.[]uint8] and main.F[go.shape.string]. Both versions take a dictionary, and both versions implement the type switch by loading the type descriptor from the dictionary and comparing it to the descriptors for []byte and string.

In these instantiations, however, the type descriptor is a constant. We know this because the constraint does not use a ~. It would be nice if the compiler were able to constant fold the dictionary value and eliminate the type switch.

I came up with this example because of #38776. It would be nice if we could change some of the hash implementations to use generic functions, which largely works until they try to call assembly code. At that point it's necessary to get the pointer to the slice or string data. It would be nice to be able to do that reasonably safely, which is what the above code does, and also efficiently, which the above code is not.

CC @randall77 @mdempsky

[randall77]

Note that this might be handled by my suggestion over at #48849, that if F and start inline, then we can see the load from a constant dictionary and can replace that load with the constant value we know is in that dictionary. (It would not help if inlining doesn't happen.)

This would require that all the |-separated types have a different shape. Or at least, it would only apply to shapes which correspond to just a single entry in the | list.

I think we might want to go further, doing unshaped instantiations where there is a non-~ type listed in the constraint. It would help not just with this issue (type switches), but with devirtualizing method invocations as well.

[gopherbot]

Change https://go.dev/cl/486895 mentions this issue: cmd/compile: constant-fold loads from constant dictionaries and types

[omerfirmak]

Depending on the solution implemented, resolving this could also help with the cases where the generic functions are called thru a function pointer. Consider the case below;

package main

import (
	"fmt"
	"unsafe"
)

func start[S []byte | string](s S) *byte {
	switch s := any(s).(type) {
	case []byte:
		return unsafe.SliceData(s)
	case string:
		return unsafe.StringData(s)
	default:
		panic("can't happen")
	}
}

func F[S []byte | string](s S) {
	fmt.Println(start(s), len(s))
}

var ByteSlicePrinter = F[[]byte]
var StringPrinter = F[string]

func main() {
	F("hi")
	F([]byte("bye"))

	ByteSlicePrinter([]byte("hi again"))
	StringPrinter("bye again")
}

alongside the main.F[go.shape.[]uint8] and main.F[go.shape.string] symbols, corresponding main.F[[]uint8] and main.F[string] wrapper symbols are also emitted. But for the [S []byte | string] constraint, compiler knows that there could only ever be one of such wrappers for a given shape, so it should be able to get rid of them.