Merge pull request #73 from iodide-project/reduce-data-copies

Don't copy data if already on the WASM heap

src/hiwire.c

@@ -212,6 +212,18 @@ EM_JS(int, hiwire_is_typedarray, (int idobj), {
   // clang-format on
 });
 
+EM_JS(int, hiwire_is_on_wasm_heap, (int idobj), {
+  var jsobj = Module.hiwire_get_value(idobj);
+  // clang-format off
+  return (jsobj.buffer === Module.HEAPU8.buffer) ? 1 : 0;
+  // clang-format on
+});
+
+EM_JS(int, hiwire_get_byteOffset, (int idobj), {
+  var jsobj = Module.hiwire_get_value(idobj);
+  return jsobj['byteOffset'];
+});
+
 EM_JS(int, hiwire_get_byteLength, (int idobj), {
   var jsobj = Module.hiwire_get_value(idobj);
   return jsobj['byteLength'];

src/hiwire.h

@@ -349,6 +349,12 @@ hiwire_nonzero(int idobj);
 int
 hiwire_is_typedarray(int idobj);
 
+/**
+ * Returns 1 if the value is a typedarray whose buffer is part of the WASM heap.
+ */
+int
+hiwire_is_on_wasm_heap(int idobj);
+
 /**
  * Returns the value of obj.byteLength.
  *
@@ -358,6 +364,15 @@ hiwire_is_typedarray(int idobj);
 int
 hiwire_get_byteLength(int idobj);
 
+/**
+ * Returns the value of obj.byteOffset.
+ *
+ * There is no error checking. Caller must ensure that hiwire_is_typedarray is
+ * true and hiwire_is_on_wasm_heap is true.
+ */
+int
+hiwire_get_byteOffset(int idobj);
+
 /**
  * Copies the buffer contents of a given typed array or buffer into the memory
  * at ptr.

src/jsproxy.c

@@ -48,7 +48,7 @@ JsProxy_GetAttr(PyObject* o, PyObject* attr_name)
 
   char* key = PyUnicode_AsUTF8(str);
 
-  if (strncmp(key, "new", 4) == 0) {
+  if (strncmp(key, "new", 4) == 0 || strncmp(key, "_has_bytes", 11) == 0) {
     Py_DECREF(str);
     return PyObject_GenericGetAttr(o, attr_name);
   } else if (strncmp(key, "typeof", 7) == 0) {
@@ -268,15 +268,20 @@ JsProxy_GetBuffer(PyObject* o, Py_buffer* view, int flags)
 
   Py_ssize_t byteLength = hiwire_get_byteLength(self->js);
 
-  if (self->bytes == NULL) {
-    self->bytes = PyBytes_FromStringAndSize(NULL, byteLength);
+  void *ptr;
+  if (hiwire_is_on_wasm_heap(self->js)) {
+    ptr = (void *)hiwire_get_byteOffset(self->js);
+  } else {
     if (self->bytes == NULL) {
-      return -1;
+      self->bytes = PyBytes_FromStringAndSize(NULL, byteLength);
+      if (self->bytes == NULL) {
+        return -1;
+      }
     }
-  }
 
-  void* ptr = PyBytes_AsString(self->bytes);
-  hiwire_copy_to_ptr(self->js, (int)ptr);
+    ptr = PyBytes_AsString(self->bytes);
+    hiwire_copy_to_ptr(self->js, (int)ptr);
+  }
 
   int dtype = hiwire_get_dtype(self->js);
 
@@ -341,6 +346,17 @@ JsProxy_GetBuffer(PyObject* o, Py_buffer* view, int flags)
   return 0;
 }
 
+static PyObject*
+JsProxy_HasBytes(PyObject *o) {
+  JsProxy* self = (JsProxy *)o;
+
+  if (self->bytes == NULL) {
+    Py_RETURN_FALSE;
+  } else {
+    Py_RETURN_TRUE;
+  }
+}
+
 // clang-format off
 static PyMappingMethods JsProxy_MappingMethods = {
   JsProxy_length,
@@ -358,6 +374,10 @@ static PyMethodDef JsProxy_Methods[] = {
     (PyCFunction)JsProxy_New,
     METH_VARARGS | METH_KEYWORDS,
     "Construct a new instance" },
+  { "_has_bytes",
+    (PyCFunction)JsProxy_HasBytes,
+    METH_NOARGS,
+    "Returns true if instance has buffer memory. For testing only." },
   { NULL }
 };
 // clang-format on

test/test_python.py

@@ -105,27 +105,40 @@ def test_js2python(selenium):
 
 
 def test_typed_arrays(selenium):
-    for (jstype, pytype) in (
-            ('Int8Array', 'b'),
-            ('Uint8Array', 'B'),
-            ('Uint8ClampedArray', 'B'),
-            ('Int16Array', 'h'),
-            ('Uint16Array', 'H'),
-            ('Int32Array', 'i'),
-            ('Uint32Array', 'I'),
-            ('Float32Array', 'f'),
-            ('Float64Array', 'd')):
-        print(jstype, pytype)
-        selenium.run_js(
-            f'window.array = new {jstype}([1, 2, 3, 4]);\n')
-        assert selenium.run(
-            'from js import array\n'
-            'import struct\n'
-            f'expected = struct.pack("{pytype*4}", 1, 2, 3, 4)\n'
-            'print(array.format, array.tolist(), array.tobytes())\n'
-            f'array.format == "{pytype}" '
-            'and array.tolist() == [1, 2, 3, 4] '
-            'and array.tobytes() == expected')
+    for wasm_heap in (False, True):
+        for (jstype, pytype) in (
+                ('Int8Array', 'b'),
+                ('Uint8Array', 'B'),
+                ('Uint8ClampedArray', 'B'),
+                ('Int16Array', 'h'),
+                ('Uint16Array', 'H'),
+                ('Int32Array', 'i'),
+                ('Uint32Array', 'I'),
+                ('Float32Array', 'f'),
+                ('Float64Array', 'd')):
+            print(wasm_heap, jstype, pytype)
+            if not wasm_heap:
+                selenium.run_js(
+                    f'window.array = new {jstype}([1, 2, 3, 4]);\n')
+            else:
+                selenium.run_js(
+                    'var buffer = pyodide._malloc('
+                    f'4 * {jstype}.BYTES_PER_ELEMENT);\n'
+                    f'window.array = new {jstype}('
+                    'pyodide.HEAPU8.buffer, buffer, 4);\n'
+                    'window.array[0] = 1;\n'
+                    'window.array[1] = 2;\n'
+                    'window.array[2] = 3;\n'
+                    'window.array[3] = 4;\n')
+            assert selenium.run(
+                'from js import array\n'
+                'import struct\n'
+                f'expected = struct.pack("{pytype*4}", 1, 2, 3, 4)\n'
+                'print(array.format, array.tolist(), array.tobytes())\n'
+                f'array.format == "{pytype}" '
+                'and array.tolist() == [1, 2, 3, 4] '
+                'and array.tobytes() == expected '
+                f'and array.obj._has_bytes() is {not wasm_heap}')
 
 
 def test_import_js(selenium):