wip

Author: szymonlopaciuk

tests/test_ref.py

@@ -164,7 +164,7 @@ class MyStruct2(xo.Struct):
 
 
 @for_all_test_contexts
-def no_test_unionref(test_context):
+def test_unionref(test_context):
     arr = xo.Float64[:]([1, 2, 3], _context=test_context)
     buf = arr._buffer
     string = xo.String("Test", _buffer=buf)
@@ -185,14 +185,16 @@ class MyStructRef(xo.Struct):
 
 
 @for_all_test_contexts
-def no_test_array_of_unionrefs(test_context):
+def test_array_of_unionrefs(test_context):
     class MyStructA(xo.Struct):
         a = xo.Float64
 
     class MyStructB(xo.Struct):
         a = xo.Int32
 
-    Element = xo.Ref[MyStructA, MyStructB]
+    class Element(xo.UnionRef):
+        _reftypes = [MyStructA, MyStructB]
+
     ArrOfUnionRefs = Element[:]
 
     aoref = ArrOfUnionRefs(10, _context=test_context)
@@ -338,3 +340,46 @@ class MyUnion(xo.UnionRef):
         _ref = [xo.Float64, xo.Int32]
 
     assert MyUnion._has_refs
+
+
+def test_move_arrays_of_unionref():
+    class Circle(xo.Struct):
+        radius = xo.UInt64
+
+    class Rectangle(xo.Struct):
+        width = xo.UInt64
+        height = xo.UInt64
+
+    class Shape(xo.UnionRef):
+        _reftypes = [Circle, Rectangle]
+
+    class StructHasRef(xo.Struct):
+        shapes = Shape[:]
+
+    class NestedHasRef(xo.Struct):
+        padding = xo.UInt64
+        nested = StructHasRef
+
+    common_buffer = xo.context_default.new_buffer()
+
+    circle = Circle(radius=5, _buffer=common_buffer)
+    rectangle = Rectangle(width=6, height=7, _buffer=common_buffer)
+
+    has_ref = StructHasRef(shapes=[circle, rectangle, circle], _buffer=common_buffer)
+    circle.radius = 10
+
+    assert has_ref.shapes[0].radius == 10
+    assert has_ref.shapes[1].width == 6
+    assert has_ref.shapes[1].height == 7
+    assert has_ref.shapes[2].radius == 10
+
+    nested_has_ref = NestedHasRef(padding=42, nested=has_ref, _buffer=common_buffer)
+    assert nested_has_ref.nested.shapes[0].radius == 10
+    assert nested_has_ref.nested.shapes[1].width == 6
+    assert nested_has_ref.nested.shapes[1].height == 7
+    assert nested_has_ref.nested.shapes[2].radius == 10
+
+    # Check that no copies were made
+    assert nested_has_ref.nested.shapes[0]._offset == has_ref.shapes[0]._offset == circle._offset
+    assert nested_has_ref.nested.shapes[1]._offset == has_ref.shapes[1]._offset == rectangle._offset
+    assert nested_has_ref.nested.shapes[2]._offset == has_ref.shapes[2]._offset == circle._offset

xobjects/__init__.py

@@ -19,6 +19,7 @@
 from .string import String
 from .struct import Struct, Field
 from .ref import Ref, UnionRef
+from .union import Union
 
 from .context_cpu import ContextCpu
 from .context_pyopencl import ContextPyopencl

xobjects/_patch_pyopencl_array.py

@@ -72,7 +72,7 @@ def _infer_fccont(arr):
     def copy_non_cont(src, dest, custom_itemsize=None, skip_typecheck=False):
         assert src.shape == dest.shape
 
-        # The case float -> complex just works (by using the src itemsize)
+        # The case float -> complex just works (by using the headers itemsize)
         if not (src.dtype == np.float64 and dest.dtype == np.complex128):
             if not skip_typecheck:
                 assert src.dtype == dest.dtype

xobjects/array.py

@@ -4,6 +4,7 @@
 # ########################################### #
 
 import logging
+from typing import TypeVar, Iterator, Protocol, Container, Iterable, Type, Self
 
 import numpy as np
 
@@ -175,6 +176,14 @@ def __init__(self, cls):
         self.cls = cls
 
 
+ArrayItemType = TypeVar("ArrayItemType")
+
+
+class ContainerInstance(Protocol[ArrayItemType]):
+    def __contains__(self, item: ArrayItemType) -> bool: ...
+    def __iter__(self) -> Iterator[ArrayItemType]: ...
+
+
 class MetaArray(type):
     def __new__(cls, name, bases, data):
         if "_itemtype" in data:  # specialized class
@@ -266,7 +275,7 @@ class Array(metaclass=MetaArray):
     _data_offset: int
 
     @classmethod
-    def mk_arrayclass(cls, itemtype, shape):
+    def mk_arrayclass(cls, itemtype: TypeVar, shape) -> Type[ContainerInstance[TypeVar] & Type[Self]]:
         if type(shape) in (int, slice):
             shape = (shape,)
         order = list(range(len(shape)))
@@ -478,13 +487,13 @@ def _to_buffer(cls, buffer, offset, value, info=None):
             if not isinstance(value, buffer.context.nplike_array_type):
                 value = buffer.context.nparray_to_context_array(value)
             buffer.update_from_nplike(coffset, cls._itemtype._dtype, value)
-        elif isinstance(value, cls):
+        elif isinstance(value, cls) and not cls._has_refs:
             if value._size == info.size:
                 buffer.update_from_xbuffer(
                     offset, value._buffer, value._offset, value._size
                 )
             else:
-                raise ValueError("Value {value} not compatible size")
+                raise ValueError(f"Value {value} not compatible size")
         elif value is None:  # no value to initialize
             if is_scalar(cls._itemtype):
                 pass  # leave uninitialized

xobjects/capi.py

@@ -208,10 +208,14 @@ def gen_c_pointed(target: Arg, conf):
     size = gen_c_size_from_arg(target, conf)
     ret = gen_c_type_from_arg(target, conf)
 
-    if target.pointer or is_compound(target.atype) or is_string(target.atype):
+    if target.pointer or is_compound(target.atype):
         chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
         return f"({ret})(({chartype}) obj+offset)"
 
+    if is_string(target.atype):
+        chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
+        return f"({chartype}) obj+offset+8"  # first entry in the string is it's size
+
     rettype = gen_pointer(ret + "*", conf)
     if size == 1:
         return f"*(({rettype}) obj+offset)"

xobjects/context_cpu.py

@@ -464,9 +464,10 @@ def compile_kernel(
                 module_name + ".o",
             ]
 
-            for ff in files_to_remove:
-                if os.path.exists(ff):
-                    os.remove(ff)
+            if 'XOBJECTS_RETAIN_BUILD_FILES' not in os.environ:
+                for ff in files_to_remove:
+                    if os.path.exists(ff):
+                        os.remove(ff)
 
     def _build_sources(
         self,

xobjects/ref.py

@@ -4,10 +4,12 @@
 # ########################################### #
 
 import logging
+from typing import Dict
 
 import numpy as np
 
-from .typeutils import Info, dispatch_arg, allocate_on_buffer, default_conf
+from .context import Source, Kernel
+from .typeutils import Info, allocate_on_buffer, default_conf
 from .scalar import Int64
 from .array import Array
 
@@ -269,21 +271,24 @@ def _gen_data_paths(cls, base=None):
         return paths
 
     @classmethod
-    def _gen_c_decl(cls, conf=default_conf):
+    def _gen_c_decl(cls, conf=default_conf) -> str:
+        """Generate method declarations for ``cls``."""
         from . import capi
 
         paths = cls._gen_data_paths()
         return capi.gen_cdefs(cls, paths, conf)
 
     @classmethod
-    def _gen_c_api(cls, conf=default_conf):
+    def _gen_c_api(cls, conf=default_conf) -> Source:
+        """Generate method definitions for ``cls``."""
         from . import capi
 
         paths = cls._gen_data_paths()
         return capi.gen_code(cls, paths, conf)
 
     @classmethod
-    def _gen_kernels(cls, conf=default_conf):
+    def _gen_kernels(cls, conf=default_conf) -> Dict[str, Kernel]:
+        """Generate kernel definitions for ``cls``."""
         from . import capi
 
         paths = cls._gen_data_paths()

xobjects/struct.py

@@ -48,20 +48,11 @@
 import logging
 from typing import Callable, Optional
 
-from .typeutils import (
-    allocate_on_buffer,
-    dispatch_arg,
-    Info,
-    _to_slot_size,
-    _is_dynamic,
-    default_conf,
-)
-
-from .general import Print
-from .scalar import Int64
 from .array import Array
-from .context import Source, Arg, Kernel
-from .context_cpu import ContextCpu
+from .context import Source
+from .scalar import Int64
+from .typeutils import (Info, _to_slot_size, allocate_on_buffer, default_conf,
+                        dispatch_arg)
 
 log = logging.getLogger(__name__)
 
@@ -307,24 +298,28 @@ def _to_buffer(cls, buffer, offset, value, info=None):
             buffer.update_from_xbuffer(
                 offset, value._buffer, value._offset, value._size
             )
-        else:  # value must be a dict, again potential disctructive
-            if info is None:
-                info = cls._inspect_args(value)
-            if cls._size is None:
-                Int64._to_buffer(buffer, offset, info.size)
-            if hasattr(
-                info, "_offsets"
-            ):  # if it has a least two dynamic fields
-                cls._set_offsets(buffer, offset, info._offsets)
-            extra = getattr(info, "extra", {})
-            for field in cls._fields:
-                fvalue = field.value_from_args(value)
-                if field.is_reference:
-                    foffset = offset + info._offsets[field.index]
-                else:
-                    foffset = offset + field.offset
-                finfo = extra.get(field.index)
-                field.ftype._to_buffer(buffer, foffset, fvalue, finfo)
+            return
+
+        if info is None:
+            info = cls._inspect_args(value)
+
+        if cls._size is None:
+            Int64._to_buffer(buffer, offset, info.size)
+
+        if hasattr(
+            info, "_offsets"
+        ):  # if it has at least two dynamic fields
+            cls._set_offsets(buffer, offset, info._offsets)
+
+        extra = getattr(info, "extra", {})
+        for field in cls._fields:
+            fvalue = field.value_from_args(value)
+            if field.is_reference:
+                foffset = offset + info._offsets[field.index]
+            else:
+                foffset = offset + field.offset
+            finfo = extra.get(field.index)
+            field.ftype._to_buffer(buffer, foffset, fvalue, finfo)
 
     def _update(self, value):
         # check if direct copy is possible