[pypy-commit] pypy ffistruct: make it possible to allocate a struct, and set/get fields on it. The only supported type is 'long' so far

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: ffistruct
Changeset: r47133:63d9d0f04f6a
Date: 2011-09-07 11:48 +0200
http://bitbucket.org/pypy/pypy/changeset/63d9d0f04f6a/

Log:	make it possible to allocate a struct, and set/get fields on it.
	The only supported type is 'long' so far

diff --git a/pypy/module/_ffi/interp_struct.py b/pypy/module/_ffi/interp_struct.py
--- a/pypy/module/_ffi/interp_struct.py
+++ b/pypy/module/_ffi/interp_struct.py
@@ -1,10 +1,10 @@
-from pypy.rpython.lltypesystem import lltype
+from pypy.rpython.lltypesystem import lltype, rffi
 from pypy.rlib import clibffi
 from pypy.interpreter.baseobjspace import Wrappable
 from pypy.interpreter.typedef import TypeDef, interp_attrproperty
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.objspace.std.typetype import type_typedef
-from pypy.module._ffi.interp_ffitype import W_FFIType
+from pypy.module._ffi.interp_ffitype import W_FFIType, app_types
 
 class W_Field(Wrappable):
 
@@ -30,12 +30,15 @@
 
 # ==============================================================================
 
-
 class W__StructDescr(Wrappable):
 
-    def __init__(self, name, ffistruct):
+    def __init__(self, name, fields_w, ffistruct):
         self.ffistruct = ffistruct
-        self.ffitype = W_FFIType('struct %s' % name, ffistruct.ffistruct, 'fixme')
+        self.w_ffitype = W_FFIType('struct %s' % name, ffistruct.ffistruct, 'fixme')
+        self.fields_w = fields_w
+        self.name2w_field = {}
+        for w_field in fields_w:
+            self.name2w_field[w_field.name] = w_field
 
     @staticmethod
     @unwrap_spec(name=str)
@@ -51,7 +54,15 @@
             field_types.append(w_field.w_ffitype.ffitype)
         #
         ffistruct = clibffi.make_struct_ffitype_e(size, alignment, field_types)
-        return W__StructDescr(name, ffistruct)
+        return W__StructDescr(name, fields_w, ffistruct)
+
+    def allocate(self, space):
+        return W__StructInstance(self)
+
+    #@jit.elidable...
+    def get_type_and_offset_for_field(self, name):
+        w_field = self.name2w_field[name]
+        return w_field.w_ffitype, w_field.offset
 
     def __del__(self):
         if self.ffistruct:
@@ -61,6 +72,60 @@
 W__StructDescr.typedef = TypeDef(
     '_StructDescr',
     __new__ = interp2app(W__StructDescr.descr_new),
-    ffitype = interp_attrproperty('ffitype', W__StructDescr),
+    ffitype = interp_attrproperty('w_ffitype', W__StructDescr),
+    allocate = interp2app(W__StructDescr.allocate),
     )
 
+
+# ==============================================================================
+
+class W__StructInstance(Wrappable):
+
+    _immutable_fields_ = ['structdescr', 'rawmem']
+
+    def __init__(self, structdescr):
+        self.structdescr = structdescr
+        size = structdescr.w_ffitype.sizeof()
+        self.rawmem = lltype.malloc(rffi.VOIDP.TO, size, flavor='raw',
+                                    zero=True, add_memory_pressure=True)
+
+    def __del__(self):
+        if self.rawmem:
+            lltype.free(self.rawmem, flavor='raw')
+            self.rawmem = lltype.nullptr(rffi.VOIDP.TO)
+
+    def getaddr(self, space):
+        addr = rffi.cast(rffi.ULONG, self.rawmem)
+        return space.wrap(addr)
+
+    @unwrap_spec(name=str)
+    def getfield(self, space, name):
+        w_ffitype, offset = self.structdescr.get_type_and_offset_for_field(name)
+        assert w_ffitype is app_types.slong # XXX: handle all cases
+        FIELD_TYPE  = rffi.LONG
+        #
+        addr = rffi.ptradd(self.rawmem, offset)
+        PTR_FIELD = lltype.Ptr(rffi.CArray(FIELD_TYPE))
+        value = rffi.cast(PTR_FIELD, addr)[0]
+        #
+        return space.wrap(value)
+
+    @unwrap_spec(name=str)
+    def setfield(self, space, name, w_value):
+        w_ffitype, offset = self.structdescr.get_type_and_offset_for_field(name)
+        assert w_ffitype is app_types.slong # XXX: handle all cases
+        FIELD_TYPE  = rffi.LONG
+        value = space.int_w(w_value)
+        #
+        addr = rffi.ptradd(self.rawmem, offset)
+        PTR_FIELD = lltype.Ptr(rffi.CArray(FIELD_TYPE))
+        rffi.cast(PTR_FIELD, addr)[0] = value
+
+
+
+W__StructInstance.typedef = TypeDef(
+    '_StructInstance',
+    getaddr  = interp2app(W__StructInstance.getaddr),
+    getfield = interp2app(W__StructInstance.getfield),
+    setfield = interp2app(W__StructInstance.setfield),
+    )
diff --git a/pypy/module/_ffi/test/test_struct.py b/pypy/module/_ffi/test/test_struct.py
--- a/pypy/module/_ffi/test/test_struct.py
+++ b/pypy/module/_ffi/test/test_struct.py
@@ -2,6 +2,20 @@
 
 class AppTestStruct(BaseAppTestFFI):
 
+    def setup_class(cls):
+        BaseAppTestFFI.setup_class.im_func(cls)
+        #
+        def read_raw_mem(self, addr, typename, length):
+            import ctypes
+            addr = ctypes.cast(addr, ctypes.c_void_p)
+            c_type = getattr(ctypes, typename)
+            array_type = ctypes.POINTER(c_type * length)
+            ptr_array = ctypes.cast(addr, array_type)
+            array = ptr_array[0]
+            lst = [array[i] for i in range(length)]
+            return lst
+        cls.w_read_raw_mem = cls.space.wrap(read_raw_mem)
+        
     def test__StructDescr(self):
         from _ffi import _StructDescr, Field, types
         longsize = types.slong.sizeof()
@@ -28,3 +42,19 @@
         assert Point.y.offset == longsize
         assert Point._struct_.ffitype.sizeof() == longsize*2
         assert Point._struct_.ffitype.name == 'struct Point'
+
+    def test_getfield_setfield(self):
+        from _ffi import _StructDescr, Field, types
+        longsize = types.slong.sizeof()
+        fields = [
+            Field('x', types.slong),
+            Field('y', types.slong),
+            ]
+        descr = _StructDescr('foo', fields)
+        struct = descr.allocate()
+        struct.setfield('x', 42)
+        struct.setfield('y', 43)
+        assert struct.getfield('x') == 42
+        assert struct.getfield('y') == 43
+        mem = self.read_raw_mem(struct.getaddr(), 'c_long', 2)
+        assert mem == [42, 43]