feat: redefine.Func on arm64

Implement redefine.Func on arm64. There's nothing to handle cloning
functions right now, so this also re-organizes the code so Original
isn't compiled on arm64.

Author: pboyd

.github/workflows/test.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       matrix:
         go-version: ['1.25', '1.26']
-        os: [ubuntu-latest, macos-15-intel, windows-latest]
+        os: [ubuntu-latest, macos-15-intel, windows-latest, ubuntu-24.04-arm, windows-11-arm, macos-latest]
 
     steps:
       - name: Checkout code

asm_amd64.go

@@ -36,7 +36,7 @@ func insertJump(buf []byte, dest uintptr) error {
 	diff32 := int32(dest - src)
 	binary.LittleEndian.PutUint32(buf[1:], uint32(diff32))
 
-	// Pad the rest of the buffer INT3 opcodes to match what the compiler does
+	// Pad the rest of the buffer with INT3 opcodes to match what the compiler does
 	for i := instructionSize; i < len(buf); i++ {
 		buf[i] = opcodeINT3
 	}

asm_arm64.go

@@ -0,0 +1,30 @@
+package redefine
+
+import (
+	"encoding/binary"
+	"errors"
+	"unsafe"
+)
+
+func insertJump(buf []byte, dest uintptr) error {
+	if len(buf) < 4 {
+		return errors.New("buffer too small")
+	}
+
+	addr := uintptr(unsafe.Pointer(unsafe.SliceData(buf)))
+	offset := int32(dest - addr)
+
+	// Encode the instruction:
+	// -----------------------------------
+	// | 000101 | ... 26 bit address ... |
+	// -----------------------------------
+	inst := (5 << 26) | (uint32(offset>>2) & (1<<26 - 1))
+	binary.LittleEndian.PutUint32(buf, inst)
+
+	// Pad the rest of the buffer with nulls
+	for i := 4; i < len(buf); i++ {
+		buf[i] = 0
+	}
+
+	return nil
+}

cacheflush_arm64.go

@@ -0,0 +1,13 @@
+//go:build arm64
+
+package redefine
+
+import "unsafe"
+
+import "C"
+
+func cacheflush(buf []byte) {
+	start := unsafe.Pointer(unsafe.SliceData(buf))
+	end := unsafe.Pointer(uintptr(len(buf)) + uintptr(start))
+	C.__builtin___clear_cache(start, end)
+}

cacheflush_fallback.go

@@ -0,0 +1,8 @@
+//go:build !arm64
+
+package redefine
+
+// This isn't needed on amd64. The arm64 version uses the C builtin which is a
+// no-op, but avoiding cgo makes cross-compiling much easier for non-Linux
+// OS's.
+func cacheflush(buf []byte) {}

clone.go

@@ -8,7 +8,6 @@ import (
 	"runtime"
 	"sync"
 	"syscall"
-	"unsafe"
 
 	"github.com/pboyd/malloc"
 )
@@ -26,39 +25,16 @@ func cloneFunc[T any](fn T) (*clonedFunc[T], error) {
 		return nil, err
 	}
 
-	//fmt.Println(disassemble(originalCode))
-
-	cloneAllocator.BeginMutate()
-	defer cloneAllocator.EndMutate()
-
-	newCode, err := cloneAllocator.Allocate(len(originalCode))
+	cf, err := _cloneFunc(fn, originalCode)
 	if err != nil {
 		return nil, err
 	}
 
-	newCode, err = relocateFunc(originalCode, newCode)
-	if err != nil {
-		return nil, err
-	}
-
-	//fmt.Println(disassemble(newCode))
-
-	// This seems too complicated. The idea is to take our newly allocated
-	// buffer of machine instructions and convince Go that it's really a
-	// function pointer of type T.
-	codeData := unsafe.SliceData(newCode)
-	cf := clonedFunc[T]{
-		clonedCode: newCode,
-		// Keep a reference to codeData so it stays around.
-		ref: &codeData,
-	}
-	cf.Func = *(*T)(unsafe.Pointer(uintptr(unsafe.Pointer(&cf.ref))))
-
 	// Make a copy of the code so that no matter what it can be restored.
 	cf.originalCode = make([]byte, len(originalCode))
 	copy(cf.originalCode, originalCode)
 
-	return &cf, nil
+	return cf, nil
 }
 
 type allocator struct {
@@ -240,10 +216,14 @@ func (cf *clonedFunc[T]) Free() {
 	cloneAllocator.BeginMutate()
 	defer cloneAllocator.EndMutate()
 
-	cloneAllocator.Free(cf.clonedCode)
+	if cf.clonedCode != nil {
+		cloneAllocator.Free(cf.clonedCode)
+	}
 
 	cf.clonedCode = nil
-	*cf.ref = nil
-	cf.ref = nil
+	if cf.ref != nil {
+		*cf.ref = nil
+		cf.ref = nil
+	}
 	cf.originalCode = nil
 }

clone_amd64.go

@@ -0,0 +1,42 @@
+//go:build amd64
+
+package redefine
+
+import "unsafe"
+
+func _cloneFunc[T any](fn T, originalCode []byte) (*clonedFunc[T], error) {
+	originalCode, err := funcSlice(fn)
+	if err != nil {
+		return nil, err
+	}
+
+	//fmt.Println(disassemble(originalCode))
+
+	cloneAllocator.BeginMutate()
+	defer cloneAllocator.EndMutate()
+
+	newCode, err := cloneAllocator.Allocate(len(originalCode))
+	if err != nil {
+		return nil, err
+	}
+
+	newCode, err = relocateFunc(originalCode, newCode)
+	if err != nil {
+		return nil, err
+	}
+
+	//fmt.Println(disassemble(newCode))
+
+	// This seems too complicated. The idea is to take our newly allocated
+	// buffer of machine instructions and convince Go that it's really a
+	// function pointer of type T.
+	codeData := unsafe.SliceData(newCode)
+	cf := clonedFunc[T]{
+		clonedCode: newCode,
+		// Keep a reference to codeData so it stays around.
+		ref: &codeData,
+	}
+	cf.Func = *(*T)(unsafe.Pointer(uintptr(unsafe.Pointer(&cf.ref))))
+
+	return &cf, nil
+}

clone_fallback.go

@@ -0,0 +1,7 @@
+//go:build !amd64
+
+package redefine
+
+func _cloneFunc[T any](fn T, originalCode []byte) (*clonedFunc[T], error) {
+	return &clonedFunc[T]{}, nil
+}

clone_test.go

@@ -1,3 +1,5 @@
+//go:build amd64
+
 package redefine
 
 import (

example_amd64_test.go

@@ -0,0 +1,31 @@
+//go:build amd64
+
+package redefine_test
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/pboyd/redefine"
+)
+
+func ExampleOriginal() {
+	redefine.Func(json.Marshal, func(v any) ([]byte, error) {
+		// Pass strings through
+		if _, ok := v.(string); ok {
+			return redefine.Original(json.Marshal)(v)
+		}
+
+		return []byte(`{"nah": true}`), nil
+	})
+	defer redefine.Restore(json.Marshal)
+
+	buf, _ := json.Marshal("A string")
+	fmt.Println(string(buf))
+
+	buf, _ = json.Marshal(123)
+	fmt.Println(string(buf))
+	// Output:
+	// "A string"
+	// {"nah": true}
+}