gh-111089: Add cache to PyUnicode_AsUTF8() for embedded NUL

Add PyASCIIObject.state.embed_null member to Python str objects. It
is used as a cache by PyUnicode_AsUTF8() to only check once if a
string contains a null character. Strings created by
PyUnicode_FromString() initializes *embed_null* since the string
cannot contain a null character.

Global static strings now also initialize the *embed_null* member.
The chr(0) singleton ("\0" string) is the only static string which
contains a null character.

Author: vstinner

Include/cpython/unicodeobject.h

@@ -142,9 +142,16 @@ typedef struct {
         unsigned int ascii:1;
         /* The object is statically allocated. */
         unsigned int statically_allocated:1;
+        // Does the string embed null characters? Possible values:
+        //   0: No
+        //   1: Yes
+        //   2: Unknown, the string must be scanned
+        //   3: Invalid state (must not be used)
+        // Cache used by PyUnicode_AsUTF8() to avoid calling strlen().
+        unsigned int embed_null:2;
         /* Padding to ensure that PyUnicode_DATA() is always aligned to
            4 bytes (see issue #19537 on m68k). */
-        unsigned int :24;
+        unsigned int :22;
     } state;
 } PyASCIIObject;
 

Include/internal/pycore_runtime_init.h

@@ -215,7 +215,7 @@ extern PyTypeObject _PyExc_MemoryError;
         _PyBytes_SIMPLE_INIT((CH), 1) \
     }
 
-#define _PyUnicode_ASCII_BASE_INIT(LITERAL, ASCII) \
+#define _PyUnicode_ASCII_BASE_INIT(LITERAL, ASCII, EMBED_NUL) \
     { \
         .ob_base = _PyObject_HEAD_INIT(&PyUnicode_Type), \
         .length = sizeof(LITERAL) - 1, \
@@ -225,11 +225,17 @@ extern PyTypeObject _PyExc_MemoryError;
             .compact = 1, \
             .ascii = (ASCII), \
             .statically_allocated = 1, \
+            .embed_null = (EMBED_NUL), \
         }, \
     }
 #define _PyASCIIObject_INIT(LITERAL) \
     { \
-        ._ascii = _PyUnicode_ASCII_BASE_INIT((LITERAL), 1), \
+        ._ascii = _PyUnicode_ASCII_BASE_INIT((LITERAL), 1, 0), \
+        ._data = (LITERAL) \
+    }
+#define _PyASCIIObject_INIT_embed_null(LITERAL) \
+    { \
+        ._ascii = _PyUnicode_ASCII_BASE_INIT((LITERAL), 1, 1), \
         ._data = (LITERAL) \
     }
 #define INIT_STR(NAME, LITERAL) \
@@ -239,7 +245,7 @@ extern PyTypeObject _PyExc_MemoryError;
 #define _PyUnicode_LATIN1_INIT(LITERAL, UTF8) \
     { \
         ._latin1 = { \
-            ._base = _PyUnicode_ASCII_BASE_INIT((LITERAL), 0), \
+            ._base = _PyUnicode_ASCII_BASE_INIT((LITERAL), 0, 0), \
             .utf8 = (UTF8), \
             .utf8_length = sizeof(UTF8) - 1, \
         }, \

Include/internal/pycore_runtime_init_generated.h

@@ -0,0 +1,5 @@
+Add ``PyASCIIObject.state.embed_null`` member to Python str objects. It is
+used as a cache by :c:func:`PyUnicode_AsUTF8` to only check once if a string
+contains a null character. Strings created by :c:func:`PyUnicode_FromString`
+initializes *embed_null* since the string cannot contain a null character.
+Patch by Victor Stinner.

Misc/NEWS.d/next/C API/2023-11-01-03-18-21.gh-issue-111089.GxXlz0.rst

@@ -301,7 +301,12 @@ unicode_fromstring(PyObject *self, PyObject *arg)
     if (!PyArg_Parse(arg, "z#", &s, &size)) {
         return NULL;
     }
-    return PyUnicode_FromString(s);
+    PyObject *unicode = PyUnicode_FromString(s);
+    if (unicode == NULL) {
+        return NULL;
+    }
+    assert(((PyASCIIObject*)unicode)->state.embed_null == 0);
+    return unicode;
 }
 
 /* Test PyUnicode_FromKindAndData() */

Modules/_testcapi/unicode.c

@@ -205,6 +205,10 @@ unicode_decode_utf8(const char *s, Py_ssize_t size,
 static inline int unicode_is_finalizing(void);
 static int unicode_is_singleton(PyObject *unicode);
 #endif
+static inline Py_ssize_t
+findchar(const void *s, int kind,
+         Py_ssize_t size, Py_UCS4 ch,
+         int direction);
 
 
 // Return a reference to the immortal empty string singleton.
@@ -623,6 +627,15 @@ _PyUnicode_CheckConsistency(PyObject *op, int check_content)
         }
         CHECK(PyUnicode_READ(kind, data, ascii->length) == 0);
     }
+
+    if (_PyUnicode_STATE(ascii).embed_null != 2) {
+        Py_ssize_t pos = findchar(PyUnicode_DATA(ascii),
+                                  PyUnicode_KIND(ascii),
+                                  PyUnicode_GET_LENGTH(ascii),
+                                  0, 1);
+        assert(_PyUnicode_STATE(ascii).embed_null == (pos >= 0));
+    }
+
     return 1;
 
 #undef CHECK
@@ -1253,6 +1266,7 @@ PyUnicode_New(Py_ssize_t size, Py_UCS4 maxchar)
     _PyUnicode_STATE(unicode).compact = 1;
     _PyUnicode_STATE(unicode).ascii = is_ascii;
     _PyUnicode_STATE(unicode).statically_allocated = 0;
+    _PyUnicode_STATE(unicode).embed_null = 2;
     if (is_ascii) {
         ((char*)data)[size] = 0;
     }
@@ -1890,7 +1904,16 @@ PyUnicode_FromString(const char *u)
         PyErr_SetString(PyExc_OverflowError, "input too long");
         return NULL;
     }
-    return PyUnicode_DecodeUTF8Stateful(u, (Py_ssize_t)size, NULL, NULL);
+    PyObject *unicode;
+    unicode = PyUnicode_DecodeUTF8Stateful(u, (Py_ssize_t)size, NULL, NULL);
+    if (unicode != NULL) {
+        // PyUnicode_DecodeUTF8Stateful(u, strlen(u)) cannot create NUL
+        // characters: the UTF-8 decoder with the strict error handler only
+        // creates a NUL character if the input string contains a NUL byte
+        // which cannot be the case here.
+        _PyUnicode_STATE(unicode).embed_null = 0;
+    }
+    return unicode;
 }
 
 
@@ -1932,6 +1955,7 @@ _PyUnicode_FromId(_Py_Identifier *id)
     if (!obj) {
         return NULL;
     }
+    _PyUnicode_STATE(obj).embed_null = 0;
     PyUnicode_InternInPlace(&obj);
 
     if (index >= ids->size) {
@@ -3846,10 +3870,27 @@ PyUnicode_AsUTF8(PyObject *unicode)
 {
     Py_ssize_t size;
     const char *utf8 = PyUnicode_AsUTF8AndSize(unicode, &size);
-    if (utf8 != NULL && strlen(utf8) != (size_t)size) {
-        PyErr_SetString(PyExc_ValueError, "embedded null character");
+    if (utf8 == NULL) {
         return NULL;
     }
+
+    // Cache to avoid calling O(n) strlen() operation at every
+    // PyUnicode_AsUTF8() call on the same object.
+    if (_PyUnicode_STATE(unicode).embed_null == 2) {
+        if (strlen(utf8) != (size_t)size) {
+            _PyUnicode_STATE(unicode).embed_null = 1;
+        }
+        else {
+            _PyUnicode_STATE(unicode).embed_null = 0;
+        }
+    }
+
+    if (_PyUnicode_STATE(unicode).embed_null == 1) {
+        PyErr_SetString(PyExc_ValueError,
+                        "embedded null character");
+        return NULL;
+    }
+
     return utf8;
 }
 

Objects/unicodeobject.c

@@ -232,6 +232,14 @@ def open_for_changes(filename, orig):
 def generate_global_strings(identifiers, strings):
     filename = os.path.join(INTERNAL, 'pycore_global_strings.h')
 
+    # NUL characters are not supported; see _PyASCIIObject_INIT_embed_null().
+    for identifier in identifiers:
+        if "\0" in identifier:
+            raise Exception(f"an identifier contains a null character: {identifier!r}")
+    for string in strings:
+        if "\0" in string:
+            raise Exception(f"a string contains a null character: {string!r}")
+
     # Read the non-generated part of the file.
     with open(filename) as infile:
         orig = infile.read()
@@ -321,7 +329,10 @@ def generate_runtime_init(identifiers, strings):
         printer.write('')
         with printer.block('#define _Py_str_ascii_INIT', continuation=True):
             for i in range(128):
-                printer.write(f'_PyASCIIObject_INIT("\\x{i:02x}"),')
+                if i == 0:
+                    printer.write(f'_PyASCIIObject_INIT_embed_null("\\x{i:02x}"),')
+                else:
+                    printer.write(f'_PyASCIIObject_INIT("\\x{i:02x}"),')
                 immortal_objects.append(f'(PyObject *)&_Py_SINGLETON(strings).ascii[{i}]')
         printer.write('')
         with printer.block('#define _Py_str_latin1_INIT', continuation=True):