runtime/jit: expose typed allocation for closure contexts

src/runtime/jit/jit.go

@@ -40,7 +40,7 @@ package jit
 
 import (
 	"sync"
-	_ "unsafe" // for go:linkname
+	"unsafe"
 )
 
 // UnwindMode controls how the Go runtime handles user frames during
@@ -215,6 +215,18 @@ func Preempt() bool {
 	return userFramePreempt()
 }
 
+// AllocTyped allocates one zeroed, garbage-collected object for JIT-generated
+// code. runtimeType must be the runtime type word of a reflect.Type. Keeping
+// size lookup and mallocgc inside the runtime/jit boundary isolates JIT clients
+// from the private runtime allocation ABI while preserving the exact GC bitmap
+// required for trailing closure context fields.
+func AllocTyped(runtimeType unsafe.Pointer) unsafe.Pointer {
+	return allocateTyped(runtimeType)
+}
+
+//go:linkname allocateTyped runtime/jit.allocateTyped
+func allocateTyped(runtimeType unsafe.Pointer) unsafe.Pointer
+
 //go:linkname userFramePreempt runtime/jit.userFramePreempt
 func userFramePreempt() bool
 

src/runtime/jit/jit_test.go

@@ -8,6 +8,7 @@ package jit_test
 
 import (
 	"fmt"
+	"reflect"
 	"runtime"
 	"runtime/jit"
 	"strings"
@@ -17,6 +18,32 @@ import (
 	"unsafe"
 )
 
+type interfaceWords struct {
+	tab  unsafe.Pointer
+	data unsafe.Pointer
+}
+
+type typedAllocationTestObject struct {
+	Pointer *int
+	Value   uintptr
+}
+
+func TestAllocTypedUsesRuntimeTypeSizeAndPointerMap(t *testing.T) {
+	typ := reflect.TypeOf(typedAllocationTestObject{})
+	runtimeType := (*interfaceWords)(unsafe.Pointer(&typ)).data
+	object := (*typedAllocationTestObject)(jit.AllocTyped(runtimeType))
+	if object.Pointer != nil || object.Value != 0 {
+		t.Fatalf("AllocTyped returned non-zero memory: %#v", object)
+	}
+	value := 42
+	object.Pointer = &value
+	runtime.GC()
+	if *object.Pointer != 42 {
+		t.Fatalf("typed allocation lost pointer field: %d", *object.Pointer)
+	}
+	runtime.KeepAlive(object)
+}
+
 // goFuncPtr returns the raw entry point of a Go function value.
 func goFuncPtr(fn func()) uintptr {
 	return **(**uintptr)(unsafe.Pointer(&fn))

src/runtime/userframe.go

@@ -522,3 +522,12 @@ func jit_userFramePreempt() bool {
 func jit_userFrameStackCheck() (guardOffset, stackSmall, moreStackPC uintptr) {
 	return unsafe.Offsetof(g{}.stackguard0), abi.StackSmall, abi.FuncPCABI0(morestack)
 }
+
+//go:linkname jit_allocateTyped runtime/jit.allocateTyped
+func jit_allocateTyped(runtimeType unsafe.Pointer) unsafe.Pointer {
+	typ := (*_type)(runtimeType)
+	if typ == nil {
+		throw("runtime/jit: nil allocation type")
+	}
+	return mallocgc(typ.Size_, typ, true)
+}