Add Unicode checks

- `_Py*UnicodeObject_CAST` macros assert they have the right bits for the shape
- `_PyUnicode_IsModifiable` & `_PyUnicodeWriter_CanWrite` assert that
  UTF-8 storage hasn't been allocated yet

Author: encukou

Include/cpython/unicodeobject.h

@@ -186,10 +186,10 @@ typedef struct {
     (assert(PyUnicode_Check(op)), \
      _Py_CAST(PyASCIIObject*, (op)))
 #define _PyCompactUnicodeObject_CAST(op) \
-    (assert(PyUnicode_Check(op)), \
+    (assert(!_PyASCIIObject_CAST(op)->state.ascii || !_PyASCIIObject_CAST(op)->state.compact), \
      _Py_CAST(PyCompactUnicodeObject*, (op)))
 #define _PyUnicodeObject_CAST(op) \
-    (assert(PyUnicode_Check(op)), \
+    (assert(!_PyASCIIObject_CAST(op)->state.compact), \
      _Py_CAST(PyUnicodeObject*, (op)))
 
 

Include/internal/pycore_unicodeobject.h

@@ -11,6 +11,7 @@ extern "C" {
 #include "pycore_fileutils.h"     // _Py_error_handler
 #include "pycore_ucnhash.h"       // _PyUnicode_Name_CAPI
 #include "pycore_runtime.h"       // _Py_LATIN1_CHR()
+#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_PTR_ACQUIRE
 
 
 // Maximum code point of Unicode 6.0: 0x10ffff (1,114,111).
@@ -108,6 +109,11 @@ _PyUnicode_EnsureUnicode(PyObject *obj)
     return 0;
 }
 
+static inline char* _PyUnicode_UTF8(PyObject *op)
+{
+    return FT_ATOMIC_LOAD_PTR_ACQUIRE(_PyCompactUnicodeObject_CAST(op)->utf8);
+}
+
 #ifndef NDEBUG
 static inline int
 _PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer)
@@ -125,6 +131,9 @@ _PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer)
     assert(PyUnstable_Unicode_GET_CACHED_HASH(buffer) == -1);
     assert(!PyUnicode_CHECK_INTERNED(buffer));
     assert(!_Py_IsImmortal(buffer));
+    assert(PyUnicode_IS_COMPACT_ASCII(buffer)
+           || _PyUnicode_UTF8(buffer) == NULL
+           || _PyUnicode_UTF8(buffer) == PyUnicode_DATA(buffer));
     return 1;
 }
 #endif

Objects/unicodeobject.c

@@ -114,11 +114,6 @@ NOTE: In the interpreter's initialization phase, some globals are currently
 #  define _PyUnicode_CHECK(op) PyUnicode_Check(op)
 #endif
 
-static inline char* _PyUnicode_UTF8(PyObject *op)
-{
-    return FT_ATOMIC_LOAD_PTR_ACQUIRE(_PyCompactUnicodeObject_CAST(op)->utf8);
-}
-
 static inline char* PyUnicode_UTF8(PyObject *op)
 {
     assert(_PyUnicode_CHECK(op));
@@ -1767,6 +1762,9 @@ _PyUnicode_IsModifiable(PyObject *unicode)
         return 0;
     if (PyUnicode_CHECK_INTERNED(unicode))
         return 0;
+    if (_PyUnicode_HAS_UTF8_MEMORY(unicode)) {
+        return 0;
+    }
 #ifdef Py_DEBUG
     /* singleton refcount is greater than 1 */
     assert(!unicode_is_singleton(unicode));