Merge pull request numpy#30 from ngoldbaum/string-array

store string data in a struct along with length

Author: ngoldbaum

stringdtype/meson.build

@@ -26,6 +26,8 @@ srcs = [
   'stringdtype/src/casts.h',
   'stringdtype/src/dtype.c',
   'stringdtype/src/main.c',
+  'stringdtype/src/static_string.c',
+  'stringdtype/src/static_string.h',
   'stringdtype/src/umath.c',
   'stringdtype/src/umath.h',
 ]

stringdtype/stringdtype/src/casts.c

@@ -1,6 +1,16 @@
 #include "casts.h"
 
 #include "dtype.h"
+#include "static_string.h"
+
+void
+gil_error(PyObject *type, const char *msg)
+{
+    PyGILState_STATE gstate;
+    gstate = PyGILState_Ensure();
+    PyErr_SetString(type, msg);
+    PyGILState_Release(gstate);
+}
 
 static NPY_CASTING
 string_to_string_resolve_descriptors(PyObject *NPY_UNUSED(self),
@@ -35,17 +45,19 @@ string_to_string(PyArrayMethod_Context *context, char *const data[],
                  NpyAuxData *NPY_UNUSED(auxdata))
 {
     npy_intp N = dimensions[0];
-    char **in = (char **)data[0];
-    char **out = (char **)data[1];
+    ss **in = (ss **)data[0];
+    ss **out = (ss **)data[1];
     // strides are in bytes but pointer offsets are in pointer widths, so
     // divide by the element size (one pointer width) to get the pointer offset
     npy_intp in_stride = strides[0] / context->descriptors[0]->elsize;
     npy_intp out_stride = strides[1] / context->descriptors[1]->elsize;
 
     while (N--) {
-        size_t length = strlen(*in);
-        out[0] = (char *)malloc((sizeof(char) * length) + 1);
-        strncpy(*out, *in, length + 1);
+        out[0] = ssdup(in[0]);
+        if (out[0] == NULL) {
+            gil_error(PyExc_MemoryError, "ssdup failed");
+            return -1;
+        }
         in += in_stride;
         out += out_stride;
     }
@@ -189,7 +201,7 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
 
     npy_intp N = dimensions[0];
     Py_UCS4 *in = (Py_UCS4 *)data[0];
-    char **out = (char **)data[1];
+    ss **out = (ss **)data[1];
 
     // 4 bytes per UCS4 character
     npy_intp in_stride = strides[0] / 4;
@@ -202,16 +214,14 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
         size_t num_codepoints = 0;
         if (utf8_size(in, max_in_size, &num_codepoints, &out_num_bytes) ==
             -1) {
-            // invalid codepoint found so acquire GIL, set error, return
-            PyGILState_STATE gstate;
-            gstate = PyGILState_Ensure();
-            PyErr_SetString(PyExc_TypeError,
-                            "Invalid unicode code point found");
-            PyGILState_Release(gstate);
+            gil_error(PyExc_TypeError, "Invalid unicode code point found");
             return -1;
         }
-        // one extra byte for null terminator
-        char *out_buf = malloc((out_num_bytes + 1) * sizeof(char));
+        ss *out_ss = ssnewempty(out_num_bytes);
+        if (out_ss == NULL) {
+            gil_error(PyExc_MemoryError, "ssnewempty failed");
+        }
+        char *out_buf = out_ss->buf;
         for (int i = 0; i < num_codepoints; i++) {
             // get code point
             Py_UCS4 code = in[i];
@@ -237,7 +247,7 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
         out_buf[out_num_bytes] = '\0';
 
         // set out to the address of the beginning of the string
-        out[0] = out_buf;
+        out[0] = out_ss;
 
         in += in_stride;
         out += out_stride;
@@ -318,7 +328,7 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
                   NpyAuxData *NPY_UNUSED(auxdata))
 {
     npy_intp N = dimensions[0];
-    char **in = (char **)data[0];
+    ss **in = (ss **)data[0];
     Py_UCS4 *out = (Py_UCS4 *)data[1];
     // strides are in bytes but pointer offsets are in pointer widths, so
     // divide by the element size (one pointer width) to get the pointer offset
@@ -329,7 +339,9 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
     long max_out_size = (context->descriptors[1]->elsize) / 4;
 
     while (N--) {
-        unsigned char *this_string = (unsigned char *)*in;
+        unsigned char *this_string = (unsigned char *)((*in)->buf);
+        size_t n_bytes = (*in)->len;
+        size_t tot_n_bytes = 0;
 
         for (int i = 0; i < max_out_size; i++) {
             Py_UCS4 code;
@@ -340,16 +352,13 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
 
             // move to next character
             this_string += num_bytes;
+            tot_n_bytes += num_bytes;
 
             // set output codepoint
             out[i] = code;
 
-            // check if this is the null terminator
-            if (code == 0) {
-                // fill all remaining characters (if any) with zero
-                for (int j = i + 1; j < max_out_size; j++) {
-                    out[j] = 0;
-                }
+            // stop if we've exhausted the input string
+            if (tot_n_bytes >= n_bytes) {
                 break;
             }
         }

stringdtype/stringdtype/src/dtype.c

@@ -1,6 +1,7 @@
 #include "dtype.h"
 
 #include "casts.h"
+#include "static_string.h"
 
 PyTypeObject *StringScalar_Type = NULL;
 
@@ -15,8 +16,9 @@ new_stringdtype_instance(void)
     if (new == NULL) {
         return NULL;
     }
-    new->base.elsize = sizeof(char *);
-    new->base.alignment = _Alignof(char *);
+    new->base.elsize = sizeof(ss *);
+    new->base.alignment = _Alignof(ss *);
+    new->base.flags |= NPY_NEEDS_INIT;
 
     return new;
 }
@@ -113,16 +115,26 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
         return -1;
     }
 
-    *dataptr = malloc(sizeof(char) * length + 1);
-    strncpy(*dataptr, val, length + 1);
+    ss *str_val = ssnewlen(val, length);
+    if (str_val == NULL) {
+        PyErr_SetString(PyExc_MemoryError, "ssnewlen failed");
+        return -1;
+    }
+    // the dtype instance has the NPY_NEEDS_INIT flag set,
+    // so if *dataptr is NULL, that means we're initializing
+    // the array and don't need to free an existing string
+    if (*dataptr != NULL) {
+        free((ss *)*dataptr);
+    }
+    *dataptr = (char *)str_val;
     Py_DECREF(val_obj);
     return 0;
 }
 
 static PyObject *
 stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
 {
-    PyObject *val_obj = PyUnicode_FromString(*dataptr);
+    PyObject *val_obj = PyUnicode_FromString(((ss *)*dataptr)->buf);
 
     if (val_obj == NULL) {
         return NULL;

stringdtype/stringdtype/src/static_string.c

@@ -0,0 +1,41 @@
+#include "static_string.h"
+
+// allocates a new ss string of length len, filling with the contents of init
+ss *
+ssnewlen(const char *init, size_t len)
+{
+    // one extra byte for null terminator
+    ss *ret = (ss *)malloc(sizeof(ss) + sizeof(char) * (len + 1));
+
+    if (ret == NULL) {
+        return NULL;
+    }
+
+    ret->len = len;
+
+    if (len > 0) {
+        memcpy(ret->buf, init, len);
+    }
+
+    ret->buf[len] = '\0';
+
+    return ret;
+}
+
+// returns a new heap-allocated copy of input string *s*
+ss *
+ssdup(const ss *s)
+{
+    return ssnewlen(s->buf, s->len);
+}
+
+// returns a new, empty string of length len
+// does not do any initialization, the caller must
+// initialize and null-terminate the string
+ss *
+ssnewempty(size_t len)
+{
+    ss *ret = (ss *)malloc(sizeof(ss) + sizeof(char) * (len + 1));
+    ret->len = len;
+    return ret;
+}

stringdtype/stringdtype/src/static_string.h

@@ -0,0 +1,26 @@
+#ifndef _NPY_STATIC_STRING_H
+#define _NPY_STATIC_STRING_H
+
+#include "stdlib.h"
+#include "string.h"
+
+typedef struct ss {
+    size_t len;
+    char buf[];
+} ss;
+
+// allocates a new ss string of length len, filling with the contents of init
+ss *
+ssnewlen(const char *init, size_t len);
+
+// returns a new heap-allocated copy of input string *s*
+ss *
+ssdup(const ss *s);
+
+// returns a new, empty string of length len
+// does not do any initialization, the caller must
+// initialize and null-terminate the string
+ss *
+ssnewempty(size_t len);
+
+#endif /*_NPY_STATIC_STRING_H */

stringdtype/stringdtype/src/umath.c

@@ -9,6 +9,7 @@
 #include "numpy/ufuncobject.h"
 
 #include "dtype.h"
+#include "static_string.h"
 #include "string.h"
 #include "umath.h"
 
@@ -19,8 +20,8 @@ string_equal_strided_loop(PyArrayMethod_Context *context, char *const data[],
                           NpyAuxData *NPY_UNUSED(auxdata))
 {
     npy_intp N = dimensions[0];
-    char **in1 = (char **)data[0];
-    char **in2 = (char **)data[1];
+    ss **in1 = (ss **)data[0];
+    ss **in2 = (ss **)data[1];
     npy_bool *out = (npy_bool *)data[2];
     // strides are in bytes but pointer offsets are in pointer widths, so
     // divide by the element size (one pointer width) to get the pointer offset
@@ -29,7 +30,10 @@ string_equal_strided_loop(PyArrayMethod_Context *context, char *const data[],
     npy_intp out_stride = strides[2];
 
     while (N--) {
-        if (strcmp(*in1, *in2) == 0) {
+        ss *s1 = *in1;
+        ss *s2 = *in2;
+
+        if (s1->len == s2->len && strncmp(s1->buf, s2->buf, s1->len) == 0) {
             *out = (npy_bool)1;
         }
         else {