py/ringbuf: Add micropython.ringbuffer() interface for general use.

Signed-off-by: Andrew Leech <[email protected]>

docs/library/micropython.rst

@@ -112,6 +112,25 @@ Functions
    This function can be used to prevent the capturing of Ctrl-C on the
    incoming stream of characters that is usually used for the REPL, in case
    that stream is used for other purposes.
+
+.. function:: RingIO(size)
+.. function:: RingIO(buffer)
+   :noindex:
+
+   Provides a fixed-size ringbuffer for bytes with a stream interface. Can be
+   considered like a fifo queue variant of `io.BytesIO`.
+
+   Can be created with integer size provided and a suitable buffer will be allocated.
+   Alternatively a `bytearray`, `memoryview` or similar object can be provided at
+   creation for in-place use.
+
+   The classic ringbuffer algorithm is used which allows for any size buffer to be
+   used however one byte will be consumed for tracking. If initialised with an integer 
+   size this will be accounted for, for example ``RingIO(16)`` will allocate a 17 byte
+   buffer internally so it can hold 16 bytes of data. 
+   When passing in a pre-allocated buffer however one byte less than its original
+   length will be available for storage, eg. ``RingIO(bytearray(16))`` will only
+   hold 15 bytes of data.
 
 .. function:: schedule(func, arg)
 

py/modmicropython.c

@@ -31,6 +31,7 @@
 #include "py/runtime.h"
 #include "py/gc.h"
 #include "py/mphal.h"
+#include "py/ringbuf.h"
 
 #if MICROPY_PY_MICROPYTHON
 
@@ -200,6 +201,9 @@ static const mp_rom_map_elem_t mp_module_micropython_globals_table[] = {
     #if MICROPY_KBD_EXCEPTION
     { MP_ROM_QSTR(MP_QSTR_kbd_intr), MP_ROM_PTR(&mp_micropython_kbd_intr_obj) },
     #endif
+    #if MICROPY_PY_MICROPYTHON_RINGIO
+    { MP_ROM_QSTR(MP_QSTR_RingIO), MP_ROM_PTR(&mp_type_micropython_ringio) },
+    #endif
     #if MICROPY_ENABLE_SCHEDULER
     { MP_ROM_QSTR(MP_QSTR_schedule), MP_ROM_PTR(&mp_micropython_schedule_obj) },
     #endif

py/mpconfig.h

@@ -1299,6 +1299,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_MICROPYTHON_HEAP_LOCKED (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EVERYTHING)
 #endif
 
+// Support for micropython.RingIO()
+#ifndef MICROPY_PY_MICROPYTHON_RINGIO
+#define MICROPY_PY_MICROPYTHON_RINGIO (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
+#endif
+
 // Whether to provide "array" module. Note that large chunk of the
 // underlying code is shared with "bytearray" builtin type, so to
 // get real savings, it should be disabled too.

py/ringbuf.c

@@ -75,14 +75,8 @@ int ringbuf_put16(ringbuf_t *r, uint16_t v) {
     return 0;
 }
 
-// Returns:
-//    0: Success
-//   -1: Not enough data available to complete read (try again later)
-//   -2: Requested read is larger than buffer - will never succeed
-int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) {
-    if (ringbuf_avail(r) < data_len) {
-        return (r->size <= data_len) ? -2 : -1;
-    }
+static inline MP_ALWAYSINLINE void _ringbuf_memcpy_get(ringbuf_t *r, uint8_t *data, size_t data_len) {
+    // Ensure bounds / space checking is done before running this
     uint32_t iget = r->iget;
     uint32_t iget_a = (iget + data_len) % r->size;
     uint8_t *datap = data;
@@ -94,17 +88,10 @@ int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) {
     }
     memcpy(datap, r->buf + iget, iget_a - iget);
     r->iget = iget_a;
-    return 0;
 }
 
-// Returns:
-//    0: Success
-//   -1: Not enough free space available to complete write (try again later)
-//   -2: Requested write is larger than buffer - will never succeed
-int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) {
-    if (ringbuf_free(r) < data_len) {
-        return (r->size <= data_len) ? -2 : -1;
-    }
+static inline MP_ALWAYSINLINE void _ringbuf_memcpy_put(ringbuf_t *r, const uint8_t *data, size_t data_len) {
+    // Ensure bounds / space checking is done before running this
     uint32_t iput = r->iput;
     uint32_t iput_a = (iput + data_len) % r->size;
     const uint8_t *datap = data;
@@ -116,5 +103,130 @@ int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) {
     }
     memcpy(r->buf + iput, datap, iput_a - iput);
     r->iput = iput_a;
+}
+
+// Returns:
+//    0: Success
+//   -1: Not enough data available to complete read (try again later)
+//   -2: Requested read is larger than buffer - will never succeed
+int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len) {
+    if (ringbuf_avail(r) < data_len) {
+        return (r->size <= data_len) ? -2 : -1;
+    }
+    _ringbuf_memcpy_get(r, data, data_len);
     return 0;
 }
+
+// Returns:
+//    0: Success
+//   -1: Not enough free space available to complete write (try again later)
+//   -2: Requested write is larger than buffer - will never succeed
+int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len) {
+    if (ringbuf_free(r) < data_len) {
+        return (r->size <= data_len) ? -2 : -1;
+    }
+    _ringbuf_memcpy_put(r, data, data_len);
+    return 0;
+}
+
+#if MICROPY_PY_MICROPYTHON_RINGIO
+#include "py/runtime.h"
+#include "py/stream.h"
+#include "py/mphal.h"
+
+typedef struct _micropython_ringio_obj_t {
+    mp_obj_base_t base;
+    ringbuf_t ringbuffer;
+} micropython_ringio_obj_t;
+
+static mp_obj_t micropython_ringio_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    mp_arg_check_num(n_args, n_kw, 1, 1, false);
+    mp_int_t buff_size = -1;
+    mp_buffer_info_t bufinfo = {NULL, 0, 0};
+
+    if (!mp_get_buffer(args[0], &bufinfo, MP_BUFFER_RW)) {
+        buff_size = mp_obj_get_int(args[0]);
+    }
+    micropython_ringio_obj_t *self = mp_obj_malloc(micropython_ringio_obj_t, type);
+    if (bufinfo.buf != NULL) {
+        // buffer passed in, use it directly for ringbuffer.
+        self->ringbuffer.buf = bufinfo.buf;
+        self->ringbuffer.size = bufinfo.len;
+        self->ringbuffer.iget = self->ringbuffer.iput = 0;
+    } else {
+        // Allocate new buffer, add one extra to buff_size as ringbuf consumes one byte for tracking.
+        ringbuf_alloc(&(self->ringbuffer), buff_size + 1);
+    }
+    return MP_OBJ_FROM_PTR(self);
+}
+
+static mp_uint_t micropython_ringio_read(mp_obj_t self_in, void *buf_in, mp_uint_t size, int *errcode) {
+    micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    size = MIN(size, ringbuf_avail(&self->ringbuffer));
+    _ringbuf_memcpy_get(&(self->ringbuffer), buf_in, size);
+    *errcode = 0;
+    return size;
+}
+
+static mp_uint_t micropython_ringio_write(mp_obj_t self_in, const void *buf_in, mp_uint_t size, int *errcode) {
+    micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    size = MIN(size, ringbuf_free(&self->ringbuffer));
+    _ringbuf_memcpy_put(&(self->ringbuffer), buf_in, size);
+    *errcode = 0;
+    return size;
+}
+
+static mp_uint_t micropython_ringio_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
+    micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    switch (request) {
+        case MP_STREAM_POLL: {
+            mp_uint_t ret = 0;
+            if ((arg & MP_STREAM_POLL_RD) && ringbuf_avail(&self->ringbuffer) > 0) {
+                ret |= MP_STREAM_POLL_RD;
+            }
+            if ((arg & MP_STREAM_POLL_WR) && ringbuf_free(&self->ringbuffer) > 0) {
+                ret |= MP_STREAM_POLL_WR;
+            }
+            return ret;
+        }
+        case MP_STREAM_CLOSE:
+            return 0;
+    }
+    *errcode = MP_EINVAL;
+    return MP_STREAM_ERROR;
+}
+
+static mp_obj_t micropython_ringio_any(mp_obj_t self_in) {
+    micropython_ringio_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    return MP_OBJ_NEW_SMALL_INT(ringbuf_avail(&self->ringbuffer));
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(micropython_ringio_any_obj, micropython_ringio_any);
+
+static const mp_rom_map_elem_t micropython_ringio_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_any), MP_ROM_PTR(&micropython_ringio_any_obj) },
+    { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
+    { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
+    { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
+    { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
+
+};
+static MP_DEFINE_CONST_DICT(micropython_ringio_locals_dict, micropython_ringio_locals_dict_table);
+
+static const mp_stream_p_t ringio_stream_p = {
+    .read = micropython_ringio_read,
+    .write = micropython_ringio_write,
+    .ioctl = micropython_ringio_ioctl,
+    .is_text = false,
+};
+
+MP_DEFINE_CONST_OBJ_TYPE(
+    mp_type_micropython_ringio,
+    MP_QSTR_RingIO,
+    MP_TYPE_FLAG_NONE,
+    make_new, micropython_ringio_make_new,
+    protocol, &ringio_stream_p,
+    locals_dict, &micropython_ringio_locals_dict
+    );
+
+#endif // MICROPY_PY_MICROPYTHON_RINGIO

py/ringbuf.h

@@ -28,9 +28,10 @@
 
 #include <stddef.h>
 #include <stdint.h>
+#include "py/mpconfig.h"
 
-#ifdef _MSC_VER
-#include "py/mpconfig.h" // For inline.
+#if MICROPY_PY_MICROPYTHON_RINGIO
+#include "py/obj.h"
 #endif
 
 typedef struct _ringbuf_t {
@@ -99,4 +100,8 @@ int ringbuf_put16(ringbuf_t *r, uint16_t v);
 int ringbuf_get_bytes(ringbuf_t *r, uint8_t *data, size_t data_len);
 int ringbuf_put_bytes(ringbuf_t *r, const uint8_t *data, size_t data_len);
 
+#if MICROPY_PY_MICROPYTHON_RINGIO
+extern const mp_obj_type_t mp_type_micropython_ringio;
+#endif
+
 #endif // MICROPY_INCLUDED_PY_RINGBUF_H

tests/micropython/ringbuffer.py

@@ -0,0 +1,48 @@
+# check that micropython.RingIO works correctly.
+
+import micropython
+
+try:
+    micropython.RingIO
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+rb = micropython.RingIO(16)
+print(rb)
+
+print(rb.any())
+
+rb.write(b"\x00")
+print(rb.any())
+
+rb.write(b"\x00")
+print(rb.any())
+
+print(rb.read(2))
+print(rb.any())
+
+
+rb.write(b"\x00\x01")
+print(rb.read())
+
+print(rb.read(1))
+
+# try to write more data than can fit at one go
+print(rb.write(b"\x00\x01" * 10))
+print(rb.write(b"\x00"))
+print(rb.read())
+
+
+ba = bytearray(17)
+rb = micropython.RingIO(ba)
+print(rb)
+print(rb.write(b"\x00\x01" * 10))
+print(rb.write(b"\x00"))
+print(rb.read())
+
+try:
+    # size must be int.
+    micropython.RingIO(None)
+except TypeError as ex:
+    print(ex)

tests/micropython/ringbuffer.py.exp

@@ -0,0 +1,16 @@
+<RingIO>
+0
+1
+2
+b'\x00\x00'
+0
+b'\x00\x01'
+b''
+16
+0
+b'\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01'
+<RingIO>
+16
+0
+b'\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01\x00\x01'
+can't convert NoneType to int

tests/micropython/ringbuffer_async.py

@@ -0,0 +1,37 @@
+# check that micropython.RingIO works correctly with asyncio.Stream.
+
+import micropython
+
+try:
+    import asyncio
+
+    asyncio.StreamWriter
+    micropython.RingIO
+except (AttributeError, ImportError):
+    print("SKIP")
+    raise SystemExit
+
+rb = micropython.RingIO(16)
+rba = asyncio.StreamWriter(rb)
+
+data = b"ABC123" * 20
+print("w", len(data))
+
+
+async def data_writer():
+    global data
+    rba.write(data)
+    await rba.drain()
+
+
+async def main():
+    task = asyncio.create_task(data_writer())
+    await asyncio.sleep_ms(10)
+    # buff = bytearray(len(data))
+    read = await rba.readexactly(len(data))
+    print(read)
+    print("r", len(read))
+    print(read == data)
+
+
+asyncio.run(main())

tests/micropython/ringbuffer_async.py.exp

@@ -0,0 +1,4 @@
+w 120
+b'ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123ABC123'
+r 120
+True