Skip to content

Commit 1865600

Browse files
ilevkivskyiilevkivskyi
committed
Add support for ints
1 parent bfafc83 commit 1865600

File tree

6 files changed

+174
-42
lines changed

6 files changed

+174
-42
lines changed

‎mypyc/build.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -656,7 +656,14 @@ def mypycify(
656656
if include_native_lib:
657657
extensions.append(
658658
get_extension()(
659-
"native_buffer", sources=[os.path.join(include_dir(), "native_buffer_internal.c")]
659+
"native_buffer",
660+
sources=[
661+
os.path.join(include_dir(), "native_buffer_internal.c"),
662+
os.path.join(include_dir(), "int_ops.c"),
663+
os.path.join(include_dir(), "exc_ops.c"),
664+
os.path.join(include_dir(), "init.c"),
665+
os.path.join(include_dir(), "pythonsupport.c"),
666+
],
660667
)
661668
)
662669

‎mypyc/lib-rt/native_buffer_internal.c‎

Lines changed: 103 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
#define PY_SSIZE_T_CLEAN
22
#include <Python.h>
3+
#include "CPy.h"
34

45
#define START_SIZE 512
56

67
#define STR_LEN_TYPE unsigned int
7-
#define MAX_STR_SIZE (1 << sizeof(STR_LEN_TYPE) * 8)
8-
9-
#define FLOAT_ERR -113.0
8+
#define MAX_STR_SIZE (1L << sizeof(STR_LEN_TYPE) * 8)
109

1110
typedef struct {
1211
PyObject_HEAD
@@ -122,7 +121,7 @@ static PyMethodDef Buffer_methods[] = {
122121
};
123122

124123
static PyTypeObject BufferType = {
125-
.ob_base = PyVarObject_HEAD_INIT(NULL, 0)
124+
.ob_base = PyVarObject_HEAD_INIT(&PyType_Type, 0)
126125
.tp_name = "Buffer",
127126
.tp_doc = PyDoc_STR("Mypy cache buffer objects"),
128127
.tp_basicsize = sizeof(BufferObject),
@@ -137,7 +136,9 @@ static PyTypeObject BufferType = {
137136
static inline
138137
char _check_buffer(PyObject *data) {
139138
if (Py_TYPE(data) != &BufferType) {
140-
PyErr_SetString(PyExc_TypeError, "data must be a Buffer object");
139+
PyErr_Format(
140+
PyExc_TypeError, "data must be a Buffer object, got %s", Py_TYPE(data)->tp_name
141+
);
141142
return 2;
142143
}
143144
return 1;
@@ -168,7 +169,7 @@ char _check_read(BufferObject *data, Py_ssize_t need) {
168169
return 1;
169170
}
170171

171-
char read_bool_internal(PyObject *data) {
172+
static char read_bool_internal(PyObject *data) {
172173
if (_check_buffer(data) == 2)
173174
return 2;
174175

@@ -180,7 +181,7 @@ char read_bool_internal(PyObject *data) {
180181
return res;
181182
}
182183

183-
PyObject *read_bool(PyObject *self, PyObject *args, PyObject *kwds) {
184+
static PyObject *read_bool(PyObject *self, PyObject *args, PyObject *kwds) {
184185
static char *kwlist[] = {"data", NULL};
185186
PyObject *data = NULL;
186187
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &data))
@@ -193,7 +194,7 @@ PyObject *read_bool(PyObject *self, PyObject *args, PyObject *kwds) {
193194
return retval;
194195
}
195196

196-
char write_bool_internal(PyObject *data, char value) {
197+
static char write_bool_internal(PyObject *data, char value) {
197198
if (_check_buffer(data) == 2)
198199
return 2;
199200

@@ -206,7 +207,7 @@ char write_bool_internal(PyObject *data, char value) {
206207
return 1;
207208
}
208209

209-
PyObject *write_bool(PyObject *self, PyObject *args, PyObject *kwds) {
210+
static PyObject *write_bool(PyObject *self, PyObject *args, PyObject *kwds) {
210211
static char *kwlist[] = {"data", "value", NULL};
211212
PyObject *data = NULL;
212213
PyObject *value = NULL;
@@ -223,17 +224,16 @@ PyObject *write_bool(PyObject *self, PyObject *args, PyObject *kwds) {
223224
return Py_None;
224225
}
225226

226-
PyObject *read_str_internal(PyObject *data) {
227+
static PyObject *read_str_internal(PyObject *data) {
227228
if (_check_buffer(data) == 2)
228229
return NULL;
229230

230-
STR_LEN_TYPE size;
231-
if (_check_read((BufferObject *)data, sizeof(size)) == 2)
231+
if (_check_read((BufferObject *)data, sizeof(STR_LEN_TYPE)) == 2)
232232
return NULL;
233233
char *buf = ((BufferObject *)data)->buf;
234234
// Read string length.
235-
size = *(STR_LEN_TYPE *)(buf + ((BufferObject *)data)->pos);
236-
((BufferObject *)data)->pos += sizeof(size);
235+
STR_LEN_TYPE size = *(STR_LEN_TYPE *)(buf + ((BufferObject *)data)->pos);
236+
((BufferObject *)data)->pos += sizeof(STR_LEN_TYPE);
237237
if (_check_read((BufferObject *)data, size) == 2)
238238
return NULL;
239239
// Read string content.
@@ -246,15 +246,15 @@ PyObject *read_str_internal(PyObject *data) {
246246
return res;
247247
}
248248

249-
PyObject *read_str(PyObject *self, PyObject *args, PyObject *kwds) {
249+
static PyObject *read_str(PyObject *self, PyObject *args, PyObject *kwds) {
250250
static char *kwlist[] = {"data", NULL};
251251
PyObject *data = NULL;
252252
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &data))
253253
return NULL;
254254
return read_str_internal(data);
255255
}
256256

257-
char write_str_internal(PyObject *data, PyObject *value) {
257+
static char write_str_internal(PyObject *data, PyObject *value) {
258258
if (_check_buffer(data) == 2)
259259
return 2;
260260

@@ -285,7 +285,7 @@ char write_str_internal(PyObject *data, PyObject *value) {
285285
return 1;
286286
}
287287

288-
PyObject *write_str(PyObject *self, PyObject *args, PyObject *kwds) {
288+
static PyObject *write_str(PyObject *self, PyObject *args, PyObject *kwds) {
289289
static char *kwlist[] = {"data", "value", NULL};
290290
PyObject *data = NULL;
291291
PyObject *value = NULL;
@@ -302,31 +302,31 @@ PyObject *write_str(PyObject *self, PyObject *args, PyObject *kwds) {
302302
return Py_None;
303303
}
304304

305-
double read_float_internal(PyObject *data) {
305+
static double read_float_internal(PyObject *data) {
306306
if (_check_buffer(data) == 2)
307-
return FLOAT_ERR;
307+
return CPY_FLOAT_ERROR;
308308

309309
if (_check_read((BufferObject *)data, sizeof(double)) == 2)
310-
return FLOAT_ERR;
310+
return CPY_FLOAT_ERROR;
311311
char *buf = ((BufferObject *)data)->buf;
312312
double res = *(double *)(buf + ((BufferObject *)data)->pos);
313313
((BufferObject *)data)->pos += sizeof(double);
314314
return res;
315315
}
316316

317-
PyObject *read_float(PyObject *self, PyObject *args, PyObject *kwds) {
317+
static PyObject *read_float(PyObject *self, PyObject *args, PyObject *kwds) {
318318
static char *kwlist[] = {"data", NULL};
319319
PyObject *data = NULL;
320320
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &data))
321321
return NULL;
322322
double retval = read_float_internal(data);
323-
if (retval == FLOAT_ERR && PyErr_Occurred()) {
323+
if (retval == CPY_FLOAT_ERROR && PyErr_Occurred()) {
324324
return NULL;
325325
}
326326
return PyFloat_FromDouble(retval);
327327
}
328328

329-
char write_float_internal(PyObject *data, double value) {
329+
static char write_float_internal(PyObject *data, double value) {
330330
if (_check_buffer(data) == 2)
331331
return 2;
332332

@@ -339,7 +339,7 @@ char write_float_internal(PyObject *data, double value) {
339339
return 1;
340340
}
341341

342-
PyObject *write_float(PyObject *self, PyObject *args, PyObject *kwds) {
342+
static PyObject *write_float(PyObject *self, PyObject *args, PyObject *kwds) {
343343
static char *kwlist[] = {"data", "value", NULL};
344344
PyObject *data = NULL;
345345
PyObject *value = NULL;
@@ -356,13 +356,92 @@ PyObject *write_float(PyObject *self, PyObject *args, PyObject *kwds) {
356356
return Py_None;
357357
}
358358

359+
static CPyTagged read_int_internal(PyObject *data) {
360+
if (_check_buffer(data) == 2)
361+
return CPY_INT_TAG;
362+
363+
if (_check_read((BufferObject *)data, sizeof(CPyTagged)) == 2)
364+
return CPY_INT_TAG;
365+
char *buf = ((BufferObject *)data)->buf;
366+
367+
CPyTagged ret = *(CPyTagged *)(buf + ((BufferObject *)data)->pos);
368+
((BufferObject *)data)->pos += sizeof(CPyTagged);
369+
if ((ret & CPY_INT_TAG) == 0)
370+
return ret;
371+
PyObject *str_ret = read_str_internal(data);
372+
if (str_ret == NULL)
373+
return CPY_INT_TAG;
374+
PyObject* ret_long = PyLong_FromUnicodeObject(str_ret, 10);
375+
Py_DECREF(str_ret);
376+
return ((CPyTagged)ret_long) | CPY_INT_TAG;
377+
}
378+
379+
static PyObject *read_int(PyObject *self, PyObject *args, PyObject *kwds) {
380+
static char *kwlist[] = {"data", NULL};
381+
PyObject *data = NULL;
382+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &data))
383+
return NULL;
384+
CPyTagged retval = read_int_internal(data);
385+
if (retval == CPY_INT_TAG) {
386+
return NULL;
387+
}
388+
return CPyTagged_StealAsObject(retval);
389+
}
390+
391+
char write_int_internal(PyObject *data, CPyTagged value) {
392+
if (_check_buffer(data) == 2)
393+
return 2;
394+
395+
if (_check_size((BufferObject *)data, sizeof(CPyTagged)) == 2)
396+
return 2;
397+
char *buf = ((BufferObject *)data)->buf;
398+
if ((value & CPY_INT_TAG) == 0) {
399+
*(CPyTagged *)(buf + ((BufferObject *)data)->pos) = value;
400+
} else {
401+
*(CPyTagged *)(buf + ((BufferObject *)data)->pos) = CPY_INT_TAG;
402+
}
403+
((BufferObject *)data)->pos += sizeof(CPyTagged);
404+
((BufferObject *)data)->end += sizeof(CPyTagged);
405+
if ((value & CPY_INT_TAG) != 0) {
406+
PyObject *str_value = PyObject_Str(CPyTagged_LongAsObject(value));
407+
if (str_value == NULL)
408+
return 2;
409+
char res = write_str_internal(data, str_value);
410+
Py_DECREF(str_value);
411+
if (res == 2)
412+
return 2;
413+
}
414+
return 1;
415+
}
416+
417+
static PyObject *write_int(PyObject *self, PyObject *args, PyObject *kwds) {
418+
static char *kwlist[] = {"data", "value", NULL};
419+
PyObject *data = NULL;
420+
PyObject *value = NULL;
421+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO", kwlist, &data, &value))
422+
return NULL;
423+
if (!PyLong_Check(value)) {
424+
PyErr_SetString(PyExc_TypeError, "value must be an int");
425+
return NULL;
426+
}
427+
CPyTagged tagged_value = CPyTagged_BorrowFromObject(value);
428+
if (write_int_internal(data, tagged_value) == 2) {
429+
return NULL;
430+
}
431+
Py_INCREF(Py_None);
432+
return Py_None;
433+
}
434+
359435
static PyMethodDef native_buffer_module_methods[] = {
436+
// TODO: switch public wrappers to METH_FASTCALL.
360437
{"write_bool", (PyCFunction)write_bool, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("write a bool")},
361438
{"read_bool", (PyCFunction)read_bool, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("read a bool")},
362439
{"write_str", (PyCFunction)write_str, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("write a string")},
363440
{"read_str", (PyCFunction)read_str, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("read a string")},
364441
{"write_float", (PyCFunction)write_float, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("write a float")},
365442
{"read_float", (PyCFunction)read_float, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("read a float")},
443+
{"write_int", (PyCFunction)write_int, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("write an int")},
444+
{"read_int", (PyCFunction)read_int, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("read an int")},
366445
{NULL, NULL, 0, NULL}
367446
};
368447

‎mypyc/primitives/misc_ops.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,3 +407,19 @@
407407
c_function_name="read_float_internal",
408408
error_kind=ERR_MAGIC,
409409
)
410+
411+
function_op(
412+
name="native_buffer.write_int",
413+
arg_types=[object_rprimitive, int_rprimitive],
414+
return_type=none_rprimitive,
415+
c_function_name="write_int_internal",
416+
error_kind=ERR_MAGIC,
417+
)
418+
419+
function_op(
420+
name="native_buffer.read_int",
421+
arg_types=[object_rprimitive],
422+
return_type=int_rprimitive,
423+
c_function_name="read_int_internal",
424+
error_kind=ERR_MAGIC,
425+
)

‎mypyc/test-data/irbuild-classes.test‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1410,42 +1410,50 @@ class TestOverload:
14101410
return x
14111411

14121412
[case testNativeBufferFastPath]
1413-
from native_buffer import Buffer, write_bool, read_bool, write_str, read_str, write_float, read_float
1413+
from native_buffer import (
1414+
Buffer, write_bool, read_bool, write_str, read_str, write_float, read_float, write_int, read_int
1415+
)
14141416

14151417
def foo() -> None:
14161418
b = Buffer()
14171419
write_str(b, "foo")
14181420
write_bool(b, True)
14191421
write_float(b, 0.1)
1422+
write_int(b, 1)
14201423

14211424
b = Buffer(b.getvalue())
14221425
x = read_str(b)
14231426
y = read_bool(b)
14241427
z = read_float(b)
1428+
t = read_int(b)
14251429
[out]
14261430
def foo():
14271431
r0, b :: native_buffer.Buffer
14281432
r1 :: str
1429-
r2, r3, r4 :: None
1430-
r5 :: bytes
1431-
r6 :: native_buffer.Buffer
1432-
r7, x :: str
1433-
r8, y :: bool
1434-
r9, z :: float
1433+
r2, r3, r4, r5 :: None
1434+
r6 :: bytes
1435+
r7 :: native_buffer.Buffer
1436+
r8, x :: str
1437+
r9, y :: bool
1438+
r10, z :: float
1439+
r11, t :: int
14351440
L0:
14361441
r0 = Buffer_internal_empty()
14371442
b = r0
14381443
r1 = 'foo'
14391444
r2 = write_str_internal(b, r1)
14401445
r3 = write_bool_internal(b, 1)
14411446
r4 = write_float_internal(b, 0.1)
1442-
r5 = Buffer_getvalue_internal(b)
1443-
r6 = Buffer_internal(r5)
1444-
b = r6
1445-
r7 = read_str_internal(b)
1446-
x = r7
1447-
r8 = read_bool_internal(b)
1448-
y = r8
1449-
r9 = read_float_internal(b)
1450-
z = r9
1447+
r5 = write_int_internal(b, 2)
1448+
r6 = Buffer_getvalue_internal(b)
1449+
r7 = Buffer_internal(r6)
1450+
b = r7
1451+
r8 = read_str_internal(b)
1452+
x = r8
1453+
r9 = read_bool_internal(b)
1454+
y = r9
1455+
r10 = read_float_internal(b)
1456+
z = r10
1457+
r11 = read_int_internal(b)
1458+
t = r11
14511459
return 1

0 commit comments

Comments
 (0)