refactor to avoid get_loop use in string to unicode cast

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -314,52 +314,10 @@ utf8_char_to_ucs4_code(unsigned char *c, size_t len, Py_UCS4 *code)
     }
 }
 
-typedef struct s2u_auxdata {
-    NpyAuxData base;
-    char *na_name;
-    size_t len;
-} s2u_auxdata;
-
-static void
-s2u_auxdata_free(s2u_auxdata *s2u_auxdata)
-{
-    PyMem_Free(s2u_auxdata->na_name);
-    PyMem_Free(s2u_auxdata);
-}
-
-static s2u_auxdata *
-get_s2u_auxdata(PyArray_Descr *from_dt)
-{
-    s2u_auxdata *res = PyMem_Calloc(1, sizeof(s2u_auxdata));
-    res->base.free = (void *)s2u_auxdata_free;
-
-    PyObject *na_pystr =
-            PyObject_Str(((StringDTypeObject *)from_dt)->na_object);
-
-    if (na_pystr == NULL) {
-        NPY_AUXDATA_FREE((NpyAuxData *)res);
-        return NULL;
-    }
-
-    Py_ssize_t size = 0;
-
-    const char *utf8_ptr = PyUnicode_AsUTF8AndSize(na_pystr, &size);
-
-    res->na_name = PyMem_Malloc((size_t)size);
-
-    memcpy(res->na_name, utf8_ptr, (size_t)size);
-
-    Py_DECREF(na_pystr);
-
-    res->len = (size_t)size;
-
-    return res;
-}
-
 static int
 string_to_unicode(PyArrayMethod_Context *context, char *const data[],
                   npy_intp const dimensions[], npy_intp const strides[],
-                  NpyAuxData *auxdata)
+                  NpyAuxData *NPY_UNUSED(auxdata))
 {
     StringDTypeObject *descr = (StringDTypeObject *)context->descriptors[0];
     int has_null = descr->na_object != NULL;
@@ -383,9 +341,8 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
         if (ss_isnull(s)) {
             if (has_null && !has_string_na) {
                 // lossy but not much else we can do
-                this_string =
-                        (unsigned char *)((s2u_auxdata *)auxdata)->na_name;
-                n_bytes = ((s2u_auxdata *)auxdata)->len;
+                this_string = (unsigned char *)descr->na_name.buf;
+                n_bytes = descr->na_name.len;
             }
             else {
                 this_string = (unsigned char *)(default_string.buf);
@@ -426,31 +383,9 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
     return 0;
 }
 
-static int
-string_to_unicode_get_loop(PyArrayMethod_Context *context, int aligned,
-                           int NPY_UNUSED(move_references),
-                           const npy_intp *strides,
-                           PyArrayMethod_StridedLoop **out_loop,
-                           NpyAuxData **out_transferdata,
-                           NPY_ARRAYMETHOD_FLAGS *flags)
-{
-    s2u_auxdata *s2u_auxdata = get_s2u_auxdata(context->descriptors[0]);
-
-    if (s2u_auxdata == NULL) {
-        return -1;
-    }
-    *out_transferdata = (NpyAuxData *)s2u_auxdata;
-
-    *out_loop = (PyArrayMethod_StridedLoop *)string_to_unicode;
-
-    *flags = NPY_METH_NO_FLOATINGPOINT_ERRORS;
-
-    return 0;
-}
-
 static PyType_Slot s2u_slots[] = {
         {NPY_METH_resolve_descriptors, &string_to_unicode_resolve_descriptors},
-        {_NPY_METH_get_loop, &string_to_unicode_get_loop},
+        {NPY_METH_strided_loop, &string_to_unicode},
         {0, NULL}};
 
 static char *s2u_name = "cast_StringDType_to_Unicode";

stringdtype/stringdtype/src/dtype.c

@@ -20,6 +20,7 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
 
     Py_XINCREF(na_object);
     ((StringDTypeObject *)new)->na_object = na_object;
+    ss na_name = NULL_STRING;
     int hasnull = na_object != NULL;
     int has_nan_na = 0;
     int has_string_na = 0;
@@ -30,9 +31,19 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
             has_string_na = 1;
             Py_ssize_t size = 0;
             const char *buf = PyUnicode_AsUTF8AndSize(na_object, &size);
-            default_string.len = size;
-            // discards const, how to avoid?
-            default_string.buf = (char *)buf;
+            default_string = NULL_STRING;
+            int res = ssnewlen(buf, (size_t)size, &default_string);
+            if (res == -1) {
+                PyErr_NoMemory();
+                Py_DECREF(new);
+                return NULL;
+            }
+            else if (res == -2) {
+                // this should never happen
+                assert(0);
+                Py_DECREF(new);
+                return NULL;
+            }
         }
         else {
             // treat as nan-like if != comparison returns a object whose truth
@@ -53,10 +64,32 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
             }
             Py_DECREF(eq);
         }
+        PyObject *na_pystr = PyObject_Str(na_object);
+        if (na_pystr == NULL) {
+            Py_DECREF(new);
+            return NULL;
+        }
+
+        Py_ssize_t size = 0;
+        const char *utf8_ptr = PyUnicode_AsUTF8AndSize(na_pystr, &size);
+        int res = ssnewlen(utf8_ptr, (size_t)size, &na_name);
+        if (res == -1) {
+            PyErr_NoMemory();
+            Py_DECREF(new);
+            return NULL;
+        }
+        else if (res == -2) {
+            // this should never happen
+            assert(0);
+            Py_DECREF(new);
+            return NULL;
+        }
+        Py_DECREF(na_pystr);
     }
     ((StringDTypeObject *)new)->has_nan_na = has_nan_na;
     ((StringDTypeObject *)new)->has_string_na = has_string_na;
     ((StringDTypeObject *)new)->default_string = default_string;
+    ((StringDTypeObject *)new)->na_name = na_name;
     ((StringDTypeObject *)new)->coerce = coerce;
 
     PyArray_Descr *base = (PyArray_Descr *)new;
@@ -538,6 +571,9 @@ stringdtype_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds)
 static void
 stringdtype_dealloc(StringDTypeObject *self)
 {
+    Py_XDECREF(self->na_object);
+    ssfree(&self->default_string);
+    ssfree(&self->na_name);
     PyArrayDescr_Type.tp_dealloc((PyObject *)self);
 }
 

stringdtype/stringdtype/src/dtype.h

@@ -26,6 +26,7 @@ typedef struct {
     int has_nan_na;
     int has_string_na;
     ss default_string;
+    ss na_name;
 } StringDTypeObject;
 
 typedef struct {

stringdtype/stringdtype/src/static_string.c

@@ -2,9 +2,11 @@
 
 #include "static_string.h"
 
-// defined this way so it has an in-memory representation that is distinct
-// from NULL, allowing us to use NULL to represent a sentinel value
+// defined this way so EMPTY_STRING has an in-memory representation that is
+// distinct from a zero-filled struct, allowing us to use a NULL_STRING
+// to represent a sentinel value
 const ss EMPTY_STRING = {0, "\0"};
+const ss NULL_STRING = {0, NULL};
 
 int
 ssnewlen(const char *init, size_t len, ss *to_init)

stringdtype/stringdtype/src/static_string.h

@@ -9,7 +9,11 @@ typedef struct ss {
     char *buf;
 } ss;
 
+// represents the empty string and can be passed safely to ss API functions
 extern const ss EMPTY_STRING;
+// represents a sentinel value, *CANNOT* be passed safely to ss API functions,
+// use ss_isnull to check if a value is null before working with it.
+extern const ss NULL_STRING;
 
 // Allocates a new buffer for *to_init*, filling with the copied contents of
 // *init* and sets *to_init->len* to *len*. Returns -1 if malloc fails and -2