Skip to content

Commit e655f68

Browse files
muayyad-alsadimuayyad-alsadi
authored
FIXES #81: 400x faster slice assignment (#134)
* FIXES #81: 400x faster slice assignment * FIXES #81: provide faster get too * FIXES #81: use slice.indices * rely on slice and add ptr_type once * add tests * add tests * add tests * add tests * add tests * add more tests * remove ctypes.memmove from __getitem__ * fix test * apply flake8 * fix test * fix test * fix test * introduce memorySlicer * introduce memorySlicer * introduce memorySlicer * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * replace memorySlicer with read/write * apply review feedback * apply review feedback * apply review feedback * apply review feedback * apply review feedback * apply review feedback * apply review feedback * remove stop from write * remove stop from write * remove stop from write * remove stop from write * remove stop from write * remove stop from write * types union, remove useless raise, remove () from pointer and add comment * fix typing * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it * write no longer accept array.array, make common code in src_ptr a function and document it
1 parent 4390dbc commit e655f68

File tree

2 files changed

+124
-0
lines changed

2 files changed

+124
-0
lines changed

‎tests/test_memory.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,54 @@ def test_large(self):
5353
MemoryType(Limits(0x100000000, None))
5454
with self.assertRaises(WasmtimeError):
5555
MemoryType(Limits(1, 0x100000000))
56+
57+
def test_slices(self):
58+
store = Store()
59+
ty = MemoryType(Limits(1, None))
60+
memory = Memory(store, ty)
61+
memory.grow(store, 2)
62+
data_ptr = memory.data_ptr(store)
63+
ba = bytearray([i for i in range(200)])
64+
size_bytes = memory.data_len(store)
65+
# happy cases
66+
offset = 2048
67+
ba_size = len(ba)
68+
# write with start and ommit stop
69+
memory.write(store, ba, offset)
70+
# check write success byte by byte, whole is asserted with read
71+
self.assertEqual(data_ptr[offset], 0)
72+
self.assertEqual(data_ptr[offset + 1], 1)
73+
self.assertEqual(data_ptr[offset + 199], 199)
74+
# read while and assert whole area
75+
out = memory.read(store, offset, offset + ba_size)
76+
self.assertEqual(ba, out)
77+
self.assertEqual(len(memory.read(store, -10)), 10)
78+
# write with start and stop
79+
memory.write(store, ba, offset + ba_size)
80+
out = memory.read(store, offset + ba_size, offset + ba_size + ba_size)
81+
self.assertEqual(ba, out)
82+
# assert old
83+
self.assertEqual(data_ptr[offset], 0)
84+
self.assertEqual(data_ptr[offset + 1], 1)
85+
self.assertEqual(data_ptr[offset + 199], 199)
86+
# assert new
87+
self.assertEqual(data_ptr[offset + ba_size], 0)
88+
self.assertEqual(data_ptr[offset + ba_size + 1], 1)
89+
self.assertEqual(data_ptr[offset + ba_size + 199], 199)
90+
# edge cases
91+
# empty slices
92+
self.assertEqual(len(memory.read(store, 0, 0)), 0)
93+
self.assertEqual(len(memory.read(store, offset, offset)), 0)
94+
self.assertEqual(len(memory.read(store, offset, offset - 1)), 0)
95+
# out of bound access returns empty array similar to list slice
96+
self.assertEqual(len(memory.read(store, size_bytes + 1)), 0)
97+
# write empty
98+
self.assertEqual(memory.write(store, bytearray(0), offset), 0)
99+
self.assertEqual(memory.write(store, bytearray(b""), offset), 0)
100+
with self.assertRaises(IndexError):
101+
memory.write(store, ba, size_bytes)
102+
with self.assertRaises(IndexError):
103+
memory.write(store, ba, size_bytes - ba_size + 1)
104+
self.assertEqual(memory.write(store, ba, -ba_size), ba_size)
105+
out = memory.read(store, -ba_size)
106+
self.assertEqual(ba, out)

‎wasmtime/_memory.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from . import _ffi as ffi
22
from ctypes import *
33
import ctypes
4+
import typing
45
from wasmtime import MemoryType, WasmtimeError
56
from ._store import Storelike
67

@@ -62,6 +63,78 @@ def data_ptr(self, store: Storelike) -> "ctypes._Pointer[c_ubyte]":
6263
"""
6364
return ffi.wasmtime_memory_data(store._context, byref(self._memory))
6465

66+
def get_buffer_ptr(self, store: Storelike,
67+
size: typing.Optional[int] = None,
68+
offset: int = 0) -> ctypes.Array:
69+
"""
70+
return raw pointer to buffer suitable for creating zero-copy writable NumPy Buffer Protocol
71+
this method is also used internally by `read()` and `write()`
72+
73+
np_mem = np.frombuffer(memory.get_buffer_ptr(store), dtype=np.uint8)
74+
np_mem[start:end] = A # write
75+
B = np_mem[start:end] # read
76+
"""
77+
if size is None:
78+
size = self.data_len(store)
79+
ptr_type = ctypes.c_ubyte * size
80+
return ptr_type.from_address(ctypes.addressof(self.data_ptr(store).contents) + offset)
81+
82+
def read(
83+
self,
84+
store: Storelike,
85+
start: typing.Optional[int] = 0,
86+
stop: typing.Optional[int] = None) -> bytearray:
87+
"""
88+
Reads this memory starting from `start` and up to `stop`
89+
and returns a copy of the contents as a `bytearray`.
90+
91+
The indexing behavior of this method is similar to `list[start:stop]`
92+
where negative starts can be used to read from the end, for example.
93+
"""
94+
size = self.data_len(store)
95+
key = slice(start, stop, None)
96+
start, stop, _ = key.indices(size)
97+
val_size = stop - start
98+
if val_size <= 0:
99+
# return bytearray of size zero
100+
return bytearray(0)
101+
src_ptr = self.get_buffer_ptr(store, val_size, start)
102+
return bytearray(src_ptr)
103+
104+
def write(
105+
self,
106+
store: Storelike,
107+
value: typing.Union[bytearray, bytes],
108+
start: typing.Optional[int] = None) -> int:
109+
"""
110+
write a bytearray value into a possibly large slice of memory
111+
negative start is allowed in a way similat to list slice mylist[-10:]
112+
if value is not bytearray it will be used to construct an intermediate bytearray (copyied twice)
113+
return number of bytes written
114+
raises IndexError when trying to write outside the memory range
115+
this happens when start offset is >= size or when end side of value is >= size
116+
"""
117+
size = self.data_len(store)
118+
key = slice(start, None)
119+
start = key.indices(size)[0]
120+
if start >= size:
121+
raise IndexError("index out of range")
122+
# value must be bytearray ex. cast bytes() to bytearray
123+
if not isinstance(value, bytearray):
124+
value = bytearray(value)
125+
val_size = len(value)
126+
if val_size == 0:
127+
return val_size
128+
# stop is exclusive
129+
stop = start + val_size
130+
if stop > size:
131+
raise IndexError("index out of range")
132+
ptr_type = ctypes.c_ubyte * val_size
133+
src_ptr = ptr_type.from_buffer(value)
134+
dst_ptr = self.get_buffer_ptr(store, val_size, start)
135+
ctypes.memmove(dst_ptr, src_ptr, val_size)
136+
return val_size
137+
65138
def data_len(self, store: Storelike) -> int:
66139
"""
67140
Returns the raw byte length of this memory.

0 commit comments

Comments
 (0)