From 37edadedeca3fb1fd184ca9b9b53dde4fcdde917 Mon Sep 17 00:00:00 2001
From: Ryan Hileman <lunixbochs@gmail.com>
Date: Sat, 6 May 2017 07:32:35 -0700
Subject: [PATCH] go: add faster RegBatch type (#822)

---
 bindings/go/unicorn/reg_batch.go | 91 ++++++++++++++++++++++++++++++++
 bindings/go/unicorn/unicorn.go   |  5 ++
 2 files changed, 96 insertions(+)
 create mode 100644 bindings/go/unicorn/reg_batch.go

diff --git a/bindings/go/unicorn/reg_batch.go b/bindings/go/unicorn/reg_batch.go
new file mode 100644
index 00000000..120c9a0e
--- /dev/null
+++ b/bindings/go/unicorn/reg_batch.go
@@ -0,0 +1,91 @@
+package unicorn
+
+import (
+	"errors"
+	"runtime"
+	"unsafe"
+)
+
+/*
+#include <unicorn/unicorn.h>
+
+void *reg_batch_setup(int *regs, int count, uint64_t **vals, int **enums, void ***refs) {
+	size_t uvsz = sizeof(uint64_t) * count;
+	size_t ensz = sizeof(int) * count;
+	size_t ursz = sizeof(uintptr_t) * count;
+	int i;
+
+	uintptr_t buf = (uintptr_t)calloc(1, uvsz+ensz+ursz);
+	if (buf == 0) return NULL;
+
+	*vals = (uint64_t *)buf;
+	*enums = (int *)(buf + uvsz);
+	*refs = (void **)(buf + uvsz + ensz);
+	for (i = 0; i < count; i++) {
+		(*enums)[i] = regs[i];
+		(*refs)[i] = &(*vals)[i];
+	}
+	return (void *)buf;
+}
+*/
+import "C"
+
+type RegBatch struct {
+	// cast to local type
+	vals []uint64
+
+	// pass these to C
+	cenums *C.int
+	crefs  *unsafe.Pointer
+	ccount C.int
+}
+
+func regBatchSetup(regs []int) (buf unsafe.Pointer, vals []uint64, cenums *C.int, crefs *unsafe.Pointer) {
+	enums := make([]C.int, len(regs))
+	for i := 0; i < len(regs); i++ {
+		enums[i] = C.int(regs[i])
+	}
+	var cvals *C.uint64_t
+	buf = C.reg_batch_setup((*C.int)(unsafe.Pointer(&enums[0])), C.int(len(regs)), &cvals, &cenums, &crefs)
+	vals = (*[1 << 24]uint64)(unsafe.Pointer(cvals))[:len(regs)]
+	return
+}
+
+func NewRegBatch(regs []int) (*RegBatch, error) {
+	r := &RegBatch{}
+	var buf unsafe.Pointer
+	buf, r.vals, r.cenums, r.crefs = regBatchSetup(regs)
+	if buf == nil {
+		return nil, errors.New("failed to allocate RegBatch memory")
+	}
+	r.ccount = C.int(len(regs))
+	// when RegBatch is collected, free C-owned data
+	runtime.SetFinalizer(r, func(r *RegBatch) {
+		C.free(buf)
+	})
+	return r, nil
+}
+
+// ReadFast skips copying and returns the internal vals array
+func (r *RegBatch) ReadFast(u Unicorn) ([]uint64, error) {
+	ucerr := C.uc_reg_read_batch(u.Handle(), r.cenums, r.crefs, r.ccount)
+	if ucerr != ERR_OK {
+		return nil, errReturn(ucerr)
+	}
+	return r.vals, nil
+}
+
+func (r *RegBatch) Read(u Unicorn, vals []uint64) error {
+	tmp, err := r.ReadFast(u)
+	if err != nil {
+		return err
+	}
+	copy(vals, tmp[:len(vals)])
+	return nil
+}
+
+func (r *RegBatch) Write(u Unicorn, vals []uint64) error {
+	copy(r.vals[:len(vals)], vals)
+	ucerr := C.uc_reg_write_batch(u.Handle(), r.cenums, r.crefs, r.ccount)
+	return errReturn(ucerr)
+}
diff --git a/bindings/go/unicorn/unicorn.go b/bindings/go/unicorn/unicorn.go
index c2fc1425..9b3d6272 100644
--- a/bindings/go/unicorn/unicorn.go
+++ b/bindings/go/unicorn/unicorn.go
@@ -59,6 +59,7 @@ type Unicorn interface {
 
 	ContextSave(reuse Context) (Context, error)
 	ContextRestore(Context) error
+	Handle() *C.uc_engine
 }
 
 type uc struct {
@@ -226,3 +227,7 @@ func (u *uc) Query(queryType int) (uint64, error) {
 	ucerr := C.uc_query(u.handle, C.uc_query_type(queryType), &ret)
 	return uint64(ret), errReturn(ucerr)
 }
+
+func (u *uc) Handle() *C.uc_engine {
+	return u.handle
+}