Support native tuple subclasses

Author: msimacek

graalpython/com.oracle.graal.python.cext/src/tupleobject.c

@@ -95,15 +95,6 @@ PyObject * tuple_subtype_new(PyTypeObject *type, PyObject *iterable) {
     return (PyObject*) newobj;
 }
 
-int PyTruffle_Tuple_SetItem(PyObject* tuple, Py_ssize_t position, PyObject* item) {
-    PyTuple_SET_ITEM(tuple, position, item);
-    return 0;
-}
-
-PyObject* PyTruffle_Tuple_GetItem(PyObject* tuple, Py_ssize_t position) {
-    return PyTuple_GET_ITEM(tuple, position);
-}
-
 PyObject* PyTruffle_Tuple_Alloc(PyTypeObject* cls, Py_ssize_t nitems) {
 	/*
 	 * TODO(fa): For 'PyVarObjects' (i.e. 'nitems > 0') we increase the size by 'sizeof(void *)'
@@ -121,3 +112,7 @@ PyObject* PyTruffle_Tuple_Alloc(PyTypeObject* cls, Py_ssize_t nitems) {
     return newObj;
 }
 
+void* PyTruffle_NativeTupleItems(PyTupleObject* tuple) {
+    return polyglot_from_PyObjectPtr_array(tuple->ob_item, tuple->ob_base.ob_size);
+}
+

graalpython/com.oracle.graal.python.test/src/tests/cpyext/test_tuple.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2018, 2021, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2018, 2023, Oracle and/or its affiliates. All rights reserved.
 # DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 #
 # The Universal Permissive License (UPL), Version 1.0
@@ -36,8 +36,10 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
+import unittest
+
+from . import CPyExtTestCase, CPyExtFunction, CPyExtFunctionOutVars, unhandled_error_compare, GRAALPYTHON, CPyExtType
 
-from . import CPyExtTestCase, CPyExtFunction, CPyExtFunctionOutVars, unhandled_error_compare, GRAALPYTHON
 __dir__ = __file__.rpartition("/")[0]
 
 
@@ -48,6 +50,14 @@ def _reference_getslice(args):
     return t[start:end]
 
 
+def _reference_getitem(args):
+    t = args[0]
+    idx = args[1]
+    if idx < 0 or idx >= len(t):
+        raise IndexError('tuple index out of range')
+    return t[idx]
+
+
 class MyStr(str):
 
     def __init__(self, s):
@@ -57,6 +67,24 @@ def __repr__(self):
         return self.s
 
 
+TupleSubclass = CPyExtType(
+    "TupleSubclass",
+    """
+        static PyObject* tuple_subclass_new(PyTypeObject* type, PyObject* args, PyObject* kwargs) {
+            args = Py_BuildValue("(O)", args);
+            if (args == NULL)
+                return NULL;
+            PyObject* result = PyTuple_Type.tp_new(type, args, kwargs);
+            return result;
+       }
+    """,
+    tp_base='&PyTuple_Type',
+    tp_new='tuple_subclass_new',
+    tp_alloc='0',
+    tp_free='0',
+)
+
+
 class TestPyTuple(CPyExtTestCase):
 
     def compile_module(self, name):
@@ -70,6 +98,7 @@ def compile_module(self, name):
             (tuple(),),
             ((1, 2, 3),),
             (("a", "b"),),
+            (TupleSubclass(1, 2, 3),),
         ),
         resultspec="n",
         argspec='O',
@@ -83,6 +112,7 @@ def compile_module(self, name):
             (tuple(),),
             ((1, 2, 3),),
             (("a", "b"),),
+            (TupleSubclass(1, 2, 3),),
             # no type checking, also accepts different objects
             ([1, 2, 3, 4],),
             ({"a": 1, "b":2},),
@@ -92,6 +122,21 @@ def compile_module(self, name):
         arguments=["PyObject* tuple"],
     )
 
+    # PyTuple_GetItem
+    test_PyTuple_GetItem = CPyExtFunctionOutVars(
+        _reference_getitem,
+        lambda: (
+            ((1, 2, 3), 1),
+            (TupleSubclass(1, 2, 3), 1),
+            ((1, 2, 3), -1),
+            ((1, 2, 3), 3),
+        ),
+        resultspec="O",
+        argspec='On',
+        arguments=["PyObject* tuple", "Py_ssize_t index"],
+        resulttype="PyObject*",
+    )
+
     # PyTuple_GetSlice
     test_PyTuple_GetSlice = CPyExtFunctionOutVars(
         _reference_getslice,
@@ -100,6 +145,7 @@ def compile_module(self, name):
             ((1, 2, 3), 0, 2),
             ((4, 5, 6), 1, 2),
             ((7, 8, 9), 2, 2),
+            (TupleSubclass(1, 2, 3), 1, 2),
         ),
         resultspec="O",
         argspec='Onn',
@@ -123,6 +169,7 @@ def compile_module(self, name):
         lambda: (
             (tuple(),),
             (("hello", "world"),),
+            (TupleSubclass(1, 2, 3),),
             ((None,),),
             ([],),
             ({},),
@@ -139,6 +186,7 @@ def compile_module(self, name):
         lambda: (
             (tuple(),),
             (("hello", "world"),),
+            (TupleSubclass(1, 2, 3),),
             ((None,),),
             ([],),
             ({},),
@@ -148,3 +196,22 @@ def compile_module(self, name):
         arguments=["PyObject* o"],
         cmpfunc=unhandled_error_compare
     )
+
+
+class TestNativeSubclass(unittest.TestCase):
+    def test_builtins(self):
+        t = TupleSubclass(1, 2, 3)
+        assert t
+        assert len(t) == 3
+        assert t[1] == 2
+        assert t[1:] == (2, 3)
+        assert t == (1, 2, 3)
+        assert t > (1, 2, 2)
+        assert t < (1, 2, 4)
+        assert 2 in t
+        assert t + (4, 5) == (1, 2, 3, 4, 5)
+        assert t * 2 == (1, 2, 3, 1, 2, 3)
+        assert list(t) == [1, 2, 3]
+        assert repr(t) == "(1, 2, 3)"
+        assert t.index(2) == 1
+        assert t.count(2) == 1

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/cext/PythonCextTupleBuiltins.java

@@ -55,16 +55,21 @@
 import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiUnaryBuiltinNode;
 import com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.PromoteBorrowedValue;
 import com.oracle.graal.python.builtins.objects.PNone;
+import com.oracle.graal.python.builtins.objects.cext.PythonAbstractNativeObject;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetItemNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.GetItemScalarNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.ListGeneralizationNode;
+import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.SetItemNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes.SetItemScalarNode;
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
-import com.oracle.graal.python.builtins.objects.tuple.TupleBuiltins.GetItemNode;
 import com.oracle.graal.python.lib.PySliceNew;
+import com.oracle.graal.python.lib.PyTupleSizeNode;
 import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.builtins.TupleNodes.GetNativeTupleStorage;
 import com.oracle.graal.python.runtime.sequence.storage.SequenceStorage;
 import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Cached.Shared;
 import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.profiles.ConditionProfile;
@@ -83,16 +88,41 @@ PTuple doGeneric(long size) {
     @CApiBuiltin(ret = Int, args = {PyObject, Py_ssize_t, PyObjectTransfer}, call = Direct)
     public abstract static class PyTuple_SetItem extends CApiTernaryBuiltinNode {
         @Specialization
-        static int doManaged(PTuple tuple, Object position, Object element,
-                        @Cached("createForList()") SequenceStorageNodes.SetItemNode setItemNode,
+        int doManaged(PTuple tuple, long index, Object element,
+                        @Shared("setItem") @Cached("createSetItem()") SequenceStorageNodes.SetItemNode setItemNode,
                         @Cached ConditionProfile generalizedProfile) {
-            setItemNode.execute(null, tuple.getSequenceStorage(), position, element);
-            SequenceStorage newStorage = setItemNode.execute(null, tuple.getSequenceStorage(), position, element);
+            SequenceStorage newStorage = setItem(tuple.getSequenceStorage(), index, element, setItemNode);
             if (generalizedProfile.profile(tuple.getSequenceStorage() != newStorage)) {
                 tuple.setSequenceStorage(newStorage);
             }
             return 0;
         }
+
+        @Specialization
+        int doNative(PythonAbstractNativeObject tuple, long index, Object element,
+                        @Cached GetNativeTupleStorage asNativeStorage,
+                        @Shared("setItem") @Cached("createSetItem()") SequenceStorageNodes.SetItemNode setItemNode) {
+            setItem(asNativeStorage.execute(tuple), index, element, setItemNode);
+            return 0;
+        }
+
+        @Fallback
+        @SuppressWarnings("unused")
+        Object fallback(Object tuple, Object index, Object element) {
+            throw raiseFallback(tuple, PythonBuiltinClassType.PTuple);
+        }
+
+        private SequenceStorage setItem(SequenceStorage sequenceStorage, long index, Object element, SetItemNode setItemNode) {
+            // we must do a bounds-check but we must not normalize the index
+            if (index < 0 || index >= sequenceStorage.length()) {
+                throw raise(IndexError, ErrorMessages.TUPLE_OUT_OF_BOUNDS);
+            }
+            return setItemNode.execute(null, sequenceStorage, (int) index, element);
+        }
+
+        protected static SetItemNode createSetItem() {
+            return SetItemNode.create(null, ListGeneralizationNode::create);
+        }
     }
 
     @CApiBuiltin(ret = PyObjectBorrowed, args = {PyObject, Py_ssize_t}, call = Direct)
@@ -102,21 +132,35 @@ public abstract static class PyTuple_GetItem extends CApiBinaryBuiltinNode {
 
         @Specialization
         Object doPTuple(PTuple tuple, long key,
-                        @Cached PromoteBorrowedValue promoteNode,
+                        @Shared("promote") @Cached PromoteBorrowedValue promoteNode,
                         @Cached ListGeneralizationNode generalizationNode,
-                        @Cached SetItemScalarNode setItemNode,
-                        @Cached GetItemScalarNode getItemNode) {
+                        @Shared("setItem") @Cached SetItemScalarNode setItemNode,
+                        @Shared("getItem") @Cached GetItemScalarNode getItemNode) {
             SequenceStorage sequenceStorage = tuple.getSequenceStorage();
-            // we must do a bounds-check but we must not normalize the index
-            if (key < 0 || key >= sequenceStorage.length()) {
-                throw raise(IndexError, ErrorMessages.TUPLE_OUT_OF_BOUNDS);
-            }
-            Object result = getItemNode.execute(sequenceStorage, (int) key);
+            int index = checkIndex(key, sequenceStorage);
+            Object result = getItemNode.execute(sequenceStorage, index);
             Object promotedValue = promoteNode.execute(result);
             if (promotedValue != null) {
                 sequenceStorage = generalizationNode.execute(sequenceStorage, promotedValue);
                 tuple.setSequenceStorage(sequenceStorage);
-                setItemNode.execute(sequenceStorage, (int) key, promotedValue);
+                setItemNode.execute(sequenceStorage, index, promotedValue);
+                return promotedValue;
+            }
+            return result;
+        }
+
+        @Specialization
+        Object doNative(PythonAbstractNativeObject tuple, long key,
+                        @Cached GetNativeTupleStorage asNativeStorage,
+                        @Shared("promote") @Cached PromoteBorrowedValue promoteNode,
+                        @Shared("setItem") @Cached SetItemScalarNode setItemNode,
+                        @Shared("getItem") @Cached GetItemScalarNode getItemNode) {
+            SequenceStorage sequenceStorage = asNativeStorage.execute(tuple);
+            int index = checkIndex(key, sequenceStorage);
+            Object result = getItemNode.execute(sequenceStorage, index);
+            Object promotedValue = promoteNode.execute(result);
+            if (promotedValue != null) {
+                setItemNode.execute(sequenceStorage, index, promotedValue);
                 return promotedValue;
             }
             return result;
@@ -126,13 +170,21 @@ Object doPTuple(PTuple tuple, long key,
         Object fallback(Object tuple, @SuppressWarnings("unused") Object pos) {
             throw raiseFallback(tuple, PythonBuiltinClassType.PTuple);
         }
+
+        private int checkIndex(long key, SequenceStorage sequenceStorage) {
+            // we must do a bounds-check but we must not normalize the index
+            if (key < 0 || key >= sequenceStorage.length()) {
+                throw raise(IndexError, ErrorMessages.TUPLE_OUT_OF_BOUNDS);
+            }
+            return (int) key;
+        }
     }
 
     @CApiBuiltin(ret = Py_ssize_t, args = {PyObject}, call = Direct)
     public abstract static class PyTuple_Size extends CApiUnaryBuiltinNode {
         @Specialization
         public static int size(Object tuple,
-                        @Cached com.oracle.graal.python.lib.PyTupleSizeNode pyTupleSizeNode) {
+                        @Cached PyTupleSizeNode pyTupleSizeNode) {
             return pyTupleSizeNode.execute(tuple);
         }
     }
@@ -141,15 +193,27 @@ public static int size(Object tuple,
     abstract static class PyTuple_GetSlice extends CApiTernaryBuiltinNode {
         @Specialization
         Object getSlice(PTuple tuple, Object iLow, Object iHigh,
-                        @Cached GetItemNode getItemNode,
-                        @Cached PySliceNew sliceNode) {
-            return getItemNode.execute(null, tuple, sliceNode.execute(iLow, iHigh, PNone.NONE));
+                        @Shared("getItem") @Cached("createForTuple()") SequenceStorageNodes.GetItemNode getItemNode,
+                        @Shared("newSlice") @Cached PySliceNew sliceNode) {
+            return doGetSlice(tuple.getSequenceStorage(), iLow, iHigh, getItemNode, sliceNode);
+        }
+
+        @Specialization
+        Object doNative(PythonAbstractNativeObject tuple, Object iLow, Object iHigh,
+                        @Shared("getItem") @Cached("createForTuple()") SequenceStorageNodes.GetItemNode getItemNode,
+                        @Shared("newSlice") @Cached PySliceNew sliceNode,
+                        @Cached GetNativeTupleStorage asNativeStorage) {
+            return doGetSlice(asNativeStorage.execute(tuple), iLow, iHigh, getItemNode, sliceNode);
         }
 
         @SuppressWarnings("unused")
         @Fallback
         Object fallback(Object tuple, Object iLow, Object iHigh) {
             throw raiseFallback(tuple, PythonBuiltinClassType.PTuple);
         }
+
+        private static Object doGetSlice(SequenceStorage storage, Object iLow, Object iHigh, GetItemNode getItemNode, PySliceNew sliceNode) {
+            return getItemNode.execute(null, storage, sliceNode.execute(iLow, iHigh, PNone.NONE));
+        }
     }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/NativeCAPISymbol.java

@@ -157,13 +157,12 @@ public enum NativeCAPISymbol implements NativeCExtSymbol {
     FUN_PY_TRUFFLE_OBJECT_ARRAY_REALLOC("PyTruffle_ObjectArrayRealloc"),
     FUN_PY_TRUFFLE_PRIMITIVE_ARRAY_FREE("PyTruffle_PrimitiveArrayFree"),
     FUN_PY_TRUFFLE_OBJECT_ARRAY_FREE("PyTruffle_ObjectArrayFree"),
+    FUN_PY_TRUFFLE_NATIVE_TUPLE_ITEMS("PyTruffle_NativeTupleItems"),
     FUN_PY_OBJECT_GENERIC_GET_DICT("_PyObject_GenericGetDict"),
     FUN_PY_OBJECT_NEW("PyTruffle_Object_New"),
     FUN_GET_THREAD_STATE_TYPE_ID("get_thread_state_typeid"),
     FUN_GET_PY_BUFFER_TYPEID("get_Py_buffer_typeid"),
     FUN_ADD_NATIVE_SLOTS("PyTruffle_Type_AddSlots"),
-    FUN_PY_TRUFFLE_TUPLE_SET_ITEM("PyTruffle_Tuple_SetItem"),
-    FUN_PY_TRUFFLE_TUPLE_GET_ITEM("PyTruffle_Tuple_GetItem"),
     FUN_PY_TRUFFLE_OBJECT_SIZE("PyTruffle_Object_Size"),
     FUN_PY_TYPE_READY("PyType_Ready"),
     FUN_GET_NEWFUNC_TYPE_ID("get_newfunc_typeid"),