Skip to content

JSON encoding speedups #647

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jcrist merged 4 commits into from
Feb 23, 2024
Merged

JSON encoding speedups #647

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9e16ac9
Simplify encode for Raw objects
jcrist Jan 23, 2024
71b868a
Accelerate JSON int & float encoding
jcrist Jan 23, 2024
85eb2f7
Accelerate datetime, date, and time JSON encoding
jcrist Jan 23, 2024
5c19008
A few simplifications
jcrist Jan 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
147 changes: 60 additions & 87 deletions msgspec/_core.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10393,16 +10393,6 @@ datetime_from_epoch(
);
}

static inline char *
ms_write_fixint(char *p, uint32_t x, int width) {
p += width;
for (int i = 0; i < width; i++) {
*--p = (x % 10) + '0';
x = x / 10;
}
return p + width;
}

static inline const char *
ms_read_fixint(const char *buf, int width, int *out) {
int x = 0;
Expand All @@ -10423,11 +10413,11 @@ ms_encode_date(PyObject *obj, char *out)
uint8_t month = PyDateTime_GET_MONTH(obj);
uint8_t day = PyDateTime_GET_DAY(obj);

out = ms_write_fixint(out, year, 4);
*out++ = '-';
out = ms_write_fixint(out, month, 2);
*out++ = '-';
out = ms_write_fixint(out, day, 2);
write_u32_4_digits(year, out);
*(out + 4) = '-';
write_u32_2_digits(month, out + 5);
*(out + 7) = '-';
write_u32_2_digits(day, out + 8);
}

/* Requires 21 bytes of scratch space */
Expand All @@ -10438,14 +10428,18 @@ ms_encode_time_parts(
PyObject *tzinfo, char *out, int out_offset
) {
char *p = out + out_offset;
p = ms_write_fixint(p, hour, 2);
write_u32_2_digits(hour, p);
p += 2;
*p++ = ':';
p = ms_write_fixint(p, minute, 2);
write_u32_2_digits(minute, p);
p += 2;
*p++ = ':';
p = ms_write_fixint(p, second, 2);
write_u32_2_digits(second, p);
p += 2;
if (microsecond) {
*p++ = '.';
p = ms_write_fixint(p, microsecond, 6);
write_u32_6_digits(microsecond, p);
p += 6;
}

if (tzinfo != Py_None) {
Expand Down Expand Up @@ -10506,9 +10500,11 @@ ms_encode_time_parts(
}
else {
*p++ = sign;
p = ms_write_fixint(p, offset_hour, 2);
write_u32_2_digits(offset_hour, p);
p += 2;
*p++ = ':';
p = ms_write_fixint(p, offset_min, 2);
write_u32_2_digits(offset_min, p);
p += 2;
}
}
}
Expand Down Expand Up @@ -10901,17 +10897,16 @@ ms_encode_timedelta(PyObject *obj, char *out) {
*out++ = 'P';

if (days != 0) {
int n = write_u64(days, out);
out += n;
out = write_u64(days, out);
*out++ = 'D';
}
if (secs != 0 || micros != 0) {
*out++ = 'T';
int n = write_u64(secs, out);
out += n;
out = write_u64(secs, out);
if (micros != 0) {
*out++ = '.';
out = ms_write_fixint(out, micros, 6);
write_u32_6_digits(micros, out);
out += 6;
while (*(out - 1) == '0') {
out--;
}
Expand Down Expand Up @@ -12337,10 +12332,7 @@ static int
mpack_encode_raw(EncoderState *self, PyObject *obj)
{
Raw *raw = (Raw *)obj;
if (ms_ensure_space(self, raw->len) < 0) return -1;
memcpy(self->output_buffer_raw + self->output_len, raw->buf, raw->len);
self->output_len += raw->len;
return 0;
return ms_write(self, raw->buf, raw->len);
}

static int
Expand Down Expand Up @@ -13292,27 +13284,6 @@ PyDoc_STRVAR(JSONEncoder__doc__,
static int json_encode_inline(EncoderState*, PyObject*);
static int json_encode(EncoderState*, PyObject*);

static MS_INLINE int
json_encode_none(EncoderState *self)
{
const char *buf = "null";
return ms_write(self, buf, 4);
}

static MS_INLINE int
json_encode_true(EncoderState *self)
{
const char *buf = "true";
return ms_write(self, buf, 4);
}

static MS_INLINE int
json_encode_false(EncoderState *self)
{
const char *buf = "false";
return ms_write(self, buf, 5);
}

static MS_NOINLINE int
json_encode_long_fallback(EncoderState *self, PyObject *obj) {
int out = -1;
Expand All @@ -13330,19 +13301,19 @@ json_encode_long_fallback(EncoderState *self, PyObject *obj) {

static MS_NOINLINE int
json_encode_long(EncoderState *self, PyObject *obj) {
char buf[20];
char *p = buf;
uint64_t x;
bool neg, overflow;
overflow = fast_long_extract_parts(obj, &neg, &x);
if (MS_UNLIKELY(overflow)) {
return json_encode_long_fallback(self, obj);
}
if (ms_ensure_space(self, 20) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
if (neg) {
*p++ = '-';
}
int n = write_u64(x, p);
return ms_write(self, buf, n + neg);
self->output_len = write_u64(x, p) - self->output_buffer_raw;
return 0;
}

static int
Expand All @@ -13354,23 +13325,22 @@ json_encode_long_as_str(EncoderState *self, PyObject *obj) {

static MS_NOINLINE int
json_encode_float(EncoderState *self, PyObject *obj) {
char buf[24];
double x = PyFloat_AS_DOUBLE(obj);
int n = write_f64(x, buf, false);
return ms_write(self, buf, n);
if (ms_ensure_space(self, 24) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
self->output_len += write_f64(x, p, false);
return 0;
}

static MS_NOINLINE int
json_encode_float_as_str(EncoderState *self, PyObject *obj) {
char buf[24];
double x = PyFloat_AS_DOUBLE(obj);
int size = write_f64(x, buf, true);
if (ms_ensure_space(self, size + 2) < 0) return -1;
if (ms_ensure_space(self, 26) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
*p++ = '"';
memcpy(p, buf, size);
*(p + size) = '"';
self->output_len += size + 2;
*p = '"';
int n = write_f64(x, p + 1, true);
*(p + 1 + n) = '"';
self->output_len += n + 2;
return 0;
}

Expand Down Expand Up @@ -13559,10 +13529,7 @@ static int
json_encode_raw(EncoderState *self, PyObject *obj)
{
Raw *raw = (Raw *)obj;
if (ms_ensure_space(self, raw->len) < 0) return -1;
memcpy(self->output_buffer_raw + self->output_len, raw->buf, raw->len);
self->output_len += raw->len;
return 0;
return ms_write(self, raw->buf, raw->len);
}

static int
Expand Down Expand Up @@ -13646,33 +13613,39 @@ json_encode_decimal(EncoderState *self, PyObject *obj)
static int
json_encode_date(EncoderState *self, PyObject *obj)
{
char buf[12];
buf[0] = '"';
buf[11] = '"';
ms_encode_date(obj, buf + 1);
return ms_write(self, buf, 12);
if (ms_ensure_space(self, 12) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
*p = '"';
ms_encode_date(obj, p + 1);
*(p + 11) = '"';
self->output_len += 12;
return 0;
}

static int
json_encode_time(EncoderState *self, PyObject *obj)
{
char buf[23];
buf[0] = '"';
int size = ms_encode_time(self->mod, obj, buf + 1);
if (ms_ensure_space(self, 23) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
*p = '"';
int size = ms_encode_time(self->mod, obj, p + 1);
if (size < 0) return -1;
buf[size + 1] = '"';
return ms_write(self, buf, size + 2);
*(p + size + 1) = '"';
self->output_len += (size + 2);
return 0;
}

static int
json_encode_datetime(EncoderState *self, PyObject *obj)
{
char buf[34];
buf[0] = '"';
int size = ms_encode_datetime(self->mod, obj, buf + 1);
if (ms_ensure_space(self, 34) < 0) return -1;
char *p = self->output_buffer_raw + self->output_len;
*p = '"';
int size = ms_encode_datetime(self->mod, obj, p + 1);
if (size < 0) return -1;
buf[size + 1] = '"';
return ms_write(self, buf, size + 2);
*(p + size + 1) = '"';
self->output_len += (size + 2);
return 0;
}

static int
Expand Down Expand Up @@ -14143,13 +14116,13 @@ json_encode_struct(EncoderState *self, PyObject *obj)
static MS_NOINLINE int
json_encode_uncommon(EncoderState *self, PyTypeObject *type, PyObject *obj) {
if (obj == Py_None) {
return json_encode_none(self);
return ms_write(self, "null", 4);
}
else if (obj == Py_True) {
return json_encode_true(self);
return ms_write(self, "true", 4);
}
else if (obj == Py_False) {
return json_encode_false(self);
return ms_write(self, "false", 5);
}
else if (Py_TYPE(type) == &StructMetaType) {
return json_encode_struct(self, obj);
Expand Down
25 changes: 21 additions & 4 deletions msgspec/itoa.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,18 @@ write_u32_8_digits(uint32_t x, char *buf) {
memcpy(buf + 6, DIGIT_TABLE + dd * 2, 2);
}

static MS_INLINE void
write_u32_6_digits(uint32_t x, char *buf) {
uint32_t aa, bbcc, bb, cc;
aa = (uint32_t)(((uint64_t)x * 109951163) >> 40); /* (x / 10000) */
bbcc = x - aa * 10000; /* (x % 10000) */
bb = (bbcc * 5243) >> 19; /* (bbcc / 100) */
cc = bbcc - bb * 100; /* (bbcc % 100) */
memcpy(buf + 0, DIGIT_TABLE + aa * 2, 2);
memcpy(buf + 2, DIGIT_TABLE + bb * 2, 2);
memcpy(buf + 4, DIGIT_TABLE + cc * 2, 2);
}

static MS_INLINE void
write_u32_4_digits(uint32_t x, char *buf) {
uint32_t aa, bb;
Expand All @@ -67,6 +79,11 @@ write_u32_4_digits(uint32_t x, char *buf) {
memcpy(buf + 2, DIGIT_TABLE + bb * 2, 2);
}

static MS_INLINE void
write_u32_2_digits(uint32_t x, char *buf) {
memcpy(buf, DIGIT_TABLE + x * 2, 2);
}

static MS_INLINE char *
write_u32_1_to_8_digits(uint32_t x, char *buf) {
uint32_t aa, bb, cc, dd, aabb, bbcc, ccdd, lz;
Expand Down Expand Up @@ -145,19 +162,19 @@ write_u64_5_to_8_digits(uint32_t x, char *buf) {
}

/* Write a uint64 to buf, requires 20 bytes of space */
static inline int
static inline char *
write_u64(uint64_t x, char *buf) {
uint64_t tmp, hgh;
uint32_t mid, low;

if (x < 100000000) { /* 1-8 digits */
return write_u32_1_to_8_digits((uint32_t)x, buf) - buf;
return write_u32_1_to_8_digits((uint32_t)x, buf);
} else if (x < (uint64_t)100000000 * 100000000) { /* 9-16 digits */
hgh = x / 100000000; /* (x / 100000000) */
low = (uint32_t)(x - hgh * 100000000); /* (x % 100000000) */
char *cur = write_u32_1_to_8_digits((uint32_t)hgh, buf);
write_u32_8_digits(low, cur);
return cur + 8 - buf;
return cur + 8;
} else { /* 17-20 digits */
tmp = x / 100000000; /* (x / 100000000) */
low = (uint32_t)(x - tmp * 100000000); /* (x % 100000000) */
Expand All @@ -166,7 +183,7 @@ write_u64(uint64_t x, char *buf) {
char *cur = write_u64_5_to_8_digits((uint32_t)hgh, buf);
write_u32_4_digits(mid, cur);
write_u32_8_digits(low, cur + 4);
return cur + 12 - buf;
return cur + 12;
}
}

Expand Down
Loading