[GR-77959] Cache utf8 bytes in native.

graalpython/com.oracle.graal.python.cext/src/capi.h

@@ -144,10 +144,12 @@ typedef struct {
 typedef struct {
     GraalPyObject ob_base;
     Py_ssize_t length;
-    Py_ssize_t byte_length;
     Py_hash_t hash;
     /* Bits 0-2: kind; bit 3: is_ascii; bits 4-5: interned state. */
     uint64_t state;
+    Py_ssize_t utf8_length;     /* Number of bytes in utf8, excluding the
+                                 * terminating \0. */
+    char *utf8;                 /* UTF-8 representation (null-terminated) */
     void *data;
 } GraalPyUnicodeObject;
 

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

@@ -1247,7 +1247,6 @@ graalpy_resize_compact(PyObject *unicode, Py_ssize_t length)
     // TODO: Actually shrink the allocation when capacity exceeds length by a large amount.
     kind = PyUnicode_KIND(unicode);
     native_unicode->length = length;
-    native_unicode->byte_length = length * kind;
     PyUnicode_WRITE(kind, native_unicode->data, length, 0);
     assert(_PyUnicode_CheckConsistency(unicode, 0));
     return unicode;
@@ -3851,6 +3850,9 @@ PyUnicode_FSDecoder(PyObject* arg, void* addr)
 static int unicode_fill_utf8(PyObject *unicode);
 #endif // GraalPy change
 
+// GraalPy change
+static const char *graalpy_unicode_as_utf8_and_size(PyObject *unicode, Py_ssize_t *size);
+
 const char *
 PyUnicode_AsUTF8AndSize(PyObject *unicode, Py_ssize_t *psize)
 {
@@ -3860,7 +3862,7 @@ PyUnicode_AsUTF8AndSize(PyObject *unicode, Py_ssize_t *psize)
     }
     // GraalPy change: upcall for managed objects
     if (points_to_py_handle_space(unicode)) {
-        return GraalPyPrivate_Unicode_AsUTF8AndSize(unicode, psize);
+        return graalpy_unicode_as_utf8_and_size(unicode, psize);
     }
 
     if (PyUnicode_UTF8(unicode) == NULL) {
@@ -15403,12 +15405,18 @@ GraalPyUnicodeObject_IsCompact(GraalPyUnicodeObject *unicode)
     return GraalPyUnicodeObject_IsCompactFromState(unicode->state);
 }
 
+static inline int
+GraalPyUnicodeObject_HasNativeData(GraalPyUnicodeObject *raw) {
+    assert (!points_to_py_handle_space(raw));
+    return GraalPyUnicodeObject_GetKind(raw);
+}
+
 static inline GraalPyUnicodeObject *
 GraalpyUnicodeObject_EnsureNativeData(PyObject *op) {
     GraalPyUnicodeObject *raw = (GraalPyUnicodeObject *) pointer_to_stub(op);
     /* 'kind == 0' is not a valid kind for any unicode object. We use it to indicate that
      * the native data was not yet initialized. */
-    if (GraalPyUnicodeObject_GetKind(raw) == 0) {
+    if (!GraalPyUnicodeObject_HasNativeData(raw)) {
         GraalPyPrivate_Unicode_FillNativeData(op);
     }
     return raw;
@@ -15474,6 +15482,27 @@ void* GraalPyUnicode_COMPACT_DATA(PyObject* op) {
     return _Py_STATIC_CAST(void*, (_PyCompactUnicodeObject_CAST(op) + 1));
 }
 
+static const char * graalpy_unicode_as_utf8_and_size(PyObject *unicode, Py_ssize_t *psize) {
+    assert (points_to_py_handle_space(unicode));
+    GraalPyUnicodeObject *raw = (GraalPyUnicodeObject *) pointer_to_stub(unicode);
+
+    /* If native data is already available and is ASCII, then just use it. */
+    if (GraalPyUnicodeObject_HasNativeData(raw) && GraalPyUnicodeObject_IsAscii(raw)) {
+        assert (!raw->utf8);
+        if (psize) {
+            *psize = raw->length;
+        }
+        return PyUnicode_DATA(unicode);
+    }
+    if (raw->utf8) {
+        if (psize) {
+            *psize = raw->utf8_length;
+        }
+        return raw->utf8;
+    }
+    return GraalPyPrivate_Unicode_AsUTF8AndSize(unicode, psize);
+}
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -1256,10 +1256,50 @@ def test_unicode_data(self):
                 Py_ssize_t n = PyUnicode_GET_LENGTH(arg);
                 return PyBytes_FromStringAndSize(data, n);
             }
+
+            static PyObject* compact_unicode_as_utf8(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(arg)) {
+                PyObject *compact = create_compact();
+                if (compact == NULL) {
+                    return NULL;
+                }
+                if (!PyUnicode_IS_COMPACT(compact) || PyUnicode_IS_ASCII(compact)) {
+                    Py_DECREF(compact);
+                    PyErr_SetString(PyExc_TypeError, "expected compact non-ASCII layout");
+                    return NULL;
+                }
+
+                Py_ssize_t size1;
+                const char *utf8_1 = PyUnicode_AsUTF8AndSize(compact, &size1);
+                if (utf8_1 == NULL) {
+                    Py_DECREF(compact);
+                    return NULL;
+                }
+                Py_ssize_t size2;
+                const char *utf8_2 = PyUnicode_AsUTF8AndSize(compact, &size2);
+                if (utf8_2 == NULL) {
+                    Py_DECREF(compact);
+                    return NULL;
+                }
+                if (utf8_1 != utf8_2 || size1 != size2) {
+                    Py_DECREF(compact);
+                    PyErr_SetString(PyExc_RuntimeError, "UTF-8 cache was not reused");
+                    return NULL;
+                }
+                if (!PyUnicode_IS_COMPACT(compact) || PyUnicode_IS_ASCII(compact)) {
+                    Py_DECREF(compact);
+                    PyErr_SetString(PyExc_RuntimeError, "UTF-8 conversion changed the Unicode layout");
+                    return NULL;
+                }
+
+                PyObject *result = PyBytes_FromStringAndSize(utf8_1, size1);
+                Py_DECREF(compact);
+                return result;
+            }
             ''',
             tp_methods='''
             {"compact_unicode_as_bytes", (PyCFunction)compact_unicode_as_bytes, METH_NOARGS, ""},
-            {"noncompact_unicode_as_bytes", (PyCFunction)noncompact_unicode_as_bytes, METH_O, ""}
+            {"noncompact_unicode_as_bytes", (PyCFunction)noncompact_unicode_as_bytes, METH_O, ""},
+            {"compact_unicode_as_utf8", (PyCFunction)compact_unicode_as_utf8, METH_NOARGS, ""}
             ''',
         )
         tester = TestUnicodeData()
@@ -1270,6 +1310,7 @@ def test_unicode_data(self):
 
         assert tester.compact_unicode_as_bytes() == b'hello'
         assert tester.noncompact_unicode_as_bytes(noncompact_unicode) == b'hello'
+        assert tester.compact_unicode_as_utf8() == b'hello'
 
 
 class TestNativeUnicodeSubclass(unittest.TestCase):

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

@@ -90,6 +90,7 @@
 import static com.oracle.graal.python.runtime.nativeaccess.NativeMemory.NULLPTR;
 import static com.oracle.graal.python.runtime.nativeaccess.NativeMemory.readByteArrayElement;
 import static com.oracle.graal.python.util.PythonUtils.TS_ENCODING;
+import static com.oracle.truffle.api.strings.TruffleString.Encoding.US_ASCII;
 import static com.oracle.truffle.api.strings.TruffleString.Encoding.UTF_16LE;
 import static com.oracle.truffle.api.strings.TruffleString.Encoding.UTF_32LE;
 import static com.oracle.truffle.api.strings.TruffleString.Encoding.UTF_8;
@@ -142,6 +143,7 @@
 import com.oracle.graal.python.builtins.objects.str.StringBuiltins.RFindNode;
 import com.oracle.graal.python.builtins.objects.str.StringBuiltins.ReplaceNode;
 import com.oracle.graal.python.builtins.objects.str.StringBuiltins.StartsWithNode;
+import com.oracle.graal.python.builtins.objects.str.StringNodes.StringMaterializeNode;
 import com.oracle.graal.python.lib.PyNumberIndexNode;
 import com.oracle.graal.python.lib.PyObjectIsTrueNode;
 import com.oracle.graal.python.lib.PyObjectLookupAttr;
@@ -178,6 +180,8 @@
 import com.oracle.truffle.api.dsl.Cached.Exclusive;
 import com.oracle.truffle.api.dsl.Cached.Shared;
 import com.oracle.truffle.api.dsl.Fallback;
+import com.oracle.truffle.api.dsl.GenerateCached;
+import com.oracle.truffle.api.dsl.GenerateInline;
 import com.oracle.truffle.api.dsl.ImportStatic;
 import com.oracle.truffle.api.dsl.NeverDefault;
 import com.oracle.truffle.api.dsl.Specialization;
@@ -1185,20 +1189,56 @@ abstract static class GraalPyPrivate_Unicode_AsUTF8AndSize extends CApiBinaryBui
         @Specialization
         static long doUnicode(PString s, long sizePtr,
                         @Bind Node inliningTarget,
-                        @Cached InlinedConditionProfile hasSizeProfile,
-                        @Cached InlinedConditionProfile hasUtf8Profile,
-                        @Cached _PyUnicode_AsUTF8String asUTF8String,
-                        @Cached HiddenAttr.ReadNode readAttrNode,
-                        @Cached HiddenAttr.WriteNode writeAttrNode) {
-            PBytes utf8bytes = s.getUtf8Bytes(inliningTarget, readAttrNode);
-            if (hasUtf8Profile.profile(inliningTarget, utf8bytes == null)) {
-                utf8bytes = (PBytes) asUTF8String.execute(s, T_STRICT);
-                s.setUtf8Bytes(inliningTarget, writeAttrNode, utf8bytes);
+                        @Cached StringMaterializeNode stringMaterializeNode,
+                        @Cached TruffleString.GetCodeRangeNode getCodeRangeNode,
+                        @Cached FillNativeDataWithKind fillNativeDataWithKind,
+                        @Cached EncodeNativeStringNode encodeNativeStringNode,
+                        @Cached CStructAccess.WriteTruffleStringNode writeTruffleStringNode,
+                        @Cached TruffleString.CopyToNativeMemoryNode copyToNativeMemoryNode) {
+
+            TruffleString materialized;
+            if (s.isMaterialized()) {
+                materialized = s.getMaterialized();
+            } else {
+                materialized = stringMaterializeNode.execute(inliningTarget, s);
             }
-            if (hasSizeProfile.profile(inliningTarget, sizePtr != NULLPTR)) {
-                NativeMemory.writeLong(sizePtr, utf8bytes.getSequenceStorage().length());
+
+            long len;
+            long mem;
+
+            long raw = HandlePointerConverter.pointerToStub(s.getNativePointer());
+            boolean hasNativeData = GraalPyUnicodeObjectUtil.isStateInitialized(raw);
+            TruffleString.CodeRange range = getCodeRangeNode.execute(materialized, PythonUtils.TS_ENCODING);
+            if (!hasNativeData && range == TruffleString.CodeRange.ASCII) {
+                /*
+                 * Initialize native data for ordinary managed ASCII strings so that their character data can also be used as
+                 * UTF-8. Existing native layouts must be preserved: their actual contents may be ASCII even when their layout
+                 * was created for a larger maximum character.
+                 */
+                mem = fillNativeDataWithKind.execute(inliningTarget, s, 1, true, US_ASCII);
+                len = CStructAccess.readLongField(raw, CFields.GraalPyUnicodeObject__length);
+            } else {
+
+                // should never be called if utf8 data is already available
+                assert CStructAccess.readPtrField(raw, CFields.GraalPyUnicodeObject__utf8) == NULLPTR;
+
+                TruffleString utf8Str = encodeNativeStringNode.execute(UTF_8, s, T_STRICT);
+                int iLen = utf8Str.byteLength(UTF_8);
+                len = iLen;
+                mem = NativeMemory.malloc(iLen + 1);
+                NativeMemory.writeByte(mem + len, (byte) 0);
+                copyToNativeMemoryNode.execute(utf8Str, 0, mem, 0, iLen, UTF_8);
+                writeTruffleStringNode.write(mem, utf8Str, UTF_8);
+
+                // populate to native data structure
+                CStructAccess.writePtrField(raw, CFields.GraalPyUnicodeObject__utf8, mem);
+                CStructAccess.writeLongField(raw, CFields.GraalPyUnicodeObject__utf8_length, len);
             }
-            return PySequenceArrayWrapper.ensureNativeSequence(utf8bytes);
+
+            if (sizePtr != NULLPTR) {
+                NativeMemory.writeLong(sizePtr, len);
+            }
+            return mem;
         }
 
         @Fallback
@@ -1373,20 +1413,55 @@ long count(Object string, Object sub, long start, long end,
         }
     }
 
+    @GenerateInline
+    @GenerateCached(false)
+    abstract static class FillNativeDataWithKind extends Node {
+
+        abstract long execute(Node inliningTarget, PString stringObject, int unicodeCharSize, boolean isAscii, TruffleString.Encoding unicodeEncoding);
+
+        @Specialization
+        static long doGeneric(PString stringObject, int charSize, boolean isAscii, TruffleString.Encoding unicodeEncoding,
+                        @Cached TruffleString.SwitchEncodingNode switchEncodingNode,
+                        @Cached(inline = false) CStructAccess.WriteTruffleStringNode writeTruffleStringNode) {
+
+            long byteLength;
+            TruffleString unicodeString = stringObject.getMaterialized();
+            unicodeString = switchEncodingNode.execute(unicodeString, unicodeEncoding);
+            byteLength = unicodeString.byteLength(unicodeEncoding);
+            long dataSize = byteLength + charSize;
+
+            long taggedPointer = stringObject.getNativePointer();
+            assert HandlePointerConverter.pointsToPyHandleSpace(taggedPointer);
+            long rawPointer = HandlePointerConverter.pointerToStub(taggedPointer);
+            long data = NativeMemory.malloc(dataSize);
+
+            // unicode object may have been interned already
+            int interned = GraalPyUnicodeObjectUtil.getInterned(rawPointer);
+            if (interned == GRAALPY_UNICODE_INTERN_STATE_UNDETERMINED) {
+                interned = GRAALPY_UNICODE_INTERN_STATE_NOT_INTERNED;
+            }
+
+            assert !GraalPyUnicodeObjectUtil.isCompact(rawPointer);
+            GraalPyUnicodeObjectUtil.initializeGraalPyUnicodeObject(rawPointer, data, byteLength / charSize, byteLength, charSize, isAscii, interned, false);
+            writeTruffleStringNode.write(data, unicodeString, unicodeEncoding);
+            return data;
+        }
+
+    }
+
     @CApiBuiltin(ret = Void, args = {PyObject}, call = Ignored)
     abstract static class GraalPyPrivate_Unicode_FillNativeData extends CApiUnaryBuiltinNode {
         @Specialization
         static Object doUnicode(PString stringObject,
+                        @Bind Node inliningTarget,
                         @Cached TruffleString.GetCodeRangeNode getCodeRangeNode,
-                        @Cached TruffleString.SwitchEncodingNode switchEncodingNode,
-                        @Cached CStructAccess.WriteTruffleStringNode writeTruffleStringNode) {
+                        @Cached FillNativeDataWithKind fillNativeDataWithKind) {
             assert stringObject.isNative();
             if (!stringObject.isMaterialized()) {
                 throw CompilerDirectives.shouldNotReachHere("unmaterialized PString should already have a native unicode stub");
             }
             int unicodeCharSize;
             boolean unicodeIsAscii = false;
-            long unicodeByteLength;
             TruffleString unicodeString = stringObject.getMaterialized();
             TruffleString.Encoding unicodeEncoding;
 
@@ -1405,24 +1480,7 @@ static Object doUnicode(PString stringObject,
                 unicodeCharSize = 4;
                 unicodeEncoding = TruffleString.Encoding.UTF_32;
             }
-            unicodeString = switchEncodingNode.execute(unicodeString, unicodeEncoding);
-            unicodeByteLength = unicodeString.byteLength(unicodeEncoding);
-            long dataSize = unicodeByteLength + unicodeCharSize;
-
-            long taggedPointer = stringObject.getNativePointer();
-            assert HandlePointerConverter.pointsToPyHandleSpace(taggedPointer);
-            long rawPointer = HandlePointerConverter.pointerToStub(taggedPointer);
-            long data = NativeMemory.malloc(dataSize);
-
-            // unicode object may have been interned already
-            int interned = GraalPyUnicodeObjectUtil.getInterned(rawPointer);
-            if (interned == GRAALPY_UNICODE_INTERN_STATE_UNDETERMINED) {
-                interned = GRAALPY_UNICODE_INTERN_STATE_NOT_INTERNED;
-            }
-
-            assert !GraalPyUnicodeObjectUtil.isCompact(rawPointer);
-            GraalPyUnicodeObjectUtil.initializeGraalPyUnicodeObject(rawPointer, data, unicodeByteLength / unicodeCharSize, unicodeByteLength, unicodeCharSize, unicodeIsAscii, interned, false);
-            writeTruffleStringNode.write(data, unicodeString, unicodeEncoding);
+            fillNativeDataWithKind.execute(inliningTarget, stringObject, unicodeCharSize, unicodeIsAscii, unicodeEncoding);
             return PNone.NO_VALUE;
         }
     }

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

@@ -824,19 +824,28 @@ private static void freeNativeStub(long pointer, boolean gc) {
         assert !HandlePointerConverter.pointsToPyFloatHandle(pointer);
         boolean isLoggable = LOGGER.isLoggable(Level.FINE);
         long rawPointer = HandlePointerConverter.pointerToStub(pointer);
-        if (GraalPyUnicodeObjectUtil.isNonCompactGraalPyUnicodeObject(rawPointer)) {
-            if (isLoggable) {
-                LOGGER.fine(PythonUtils.formatJString("releasing native non-compact GraalPyUnicodeObject stub 0x%x", rawPointer));
+        if (GraalPyUnicodeObjectUtil.isGraalPyUnicodeObject(rawPointer)) {
+            if (!GraalPyUnicodeObjectUtil.isCompact(rawPointer)) {
+                if (isLoggable) {
+                    LOGGER.fine(PythonUtils.formatJString("releasing native non-compact GraalPyUnicodeObject stub 0x%x", rawPointer));
+                }
+                long nonCompactDataPointer = GraalPyUnicodeObjectUtil.getNonCompactDataPointer(rawPointer);
+                /*
+                 * If 'rawPointer + sizeof(GraalPyUnicodeObject) == nonCompactDataPointer', this indicates that the compact flag is incorrect. However, this is
+                 * not guaranteed because the allocated memory of the GraalPyUnicodeObject and the data could, by accident, be consecutive. So, we cannot assert this.
+                 */
+                if (isLoggable && rawPointer + CStructs.GraalPyUnicodeObject.size() == nonCompactDataPointer) {
+                    LOGGER.fine(PythonUtils.formatJString("Non-compact data pointer looks like compact data (GraalPyUnicodeObject=0x%x, data=0x%x)", rawPointer, nonCompactDataPointer));
+                }
+                free(nonCompactDataPointer);
             }
-            long nonCompactDataPointer = GraalPyUnicodeObjectUtil.getNonCompactDataPointer(rawPointer);
-            /*
-             * If 'rawPointer + sizeof(GraalPyUnicodeObject) == nonCompactDataPointer', this indicates that the compact flag is incorrect. However, this is
-             * not guaranteed because the allocated memory of the GraalPyUnicodeObject and the data could, by accident, be consecutive. So, we cannot assert this.
-             */
-            if (isLoggable && rawPointer + CStructs.GraalPyUnicodeObject.size() == nonCompactDataPointer) {
-                LOGGER.fine(PythonUtils.formatJString("Non-compact data pointer looks like compact data (GraalPyUnicodeObject=0x%x, data=0x%x)", rawPointer, nonCompactDataPointer));
+            long utf8data = GraalPyUnicodeObjectUtil.getUtf8DataPointer(rawPointer);
+            if (utf8data != NULLPTR) {
+                if (isLoggable) {
+                    LOGGER.fine(PythonUtils.formatJString("releasing native utf8 data of GraalPyUnicodeObject stub 0x%x", utf8data));
+                }
+                free(utf8data);
             }
-            free(nonCompactDataPointer);
         }
         if (gc) {
             PyObjectGCDelNode.executeUncached(pointer);