unix/moduselect: If file object passed to .register(), return it in .poll().

This makes unix "uselect" compatible with baremetal "uselect". Previosuly,
unix version accepted file/socket objects, but internally converted that
to file descriptors, and that's what .poll() returned. To acheive new
behavior, file-like objects are stored internally in an array, in addition
to existing array of struct pollfd. This array is created only on first
case of file-like object being passed to .register(). If only raw fd's are
passed, there will be no additional memory used comparing to the original
implementation.

1 files changed, 22 insertions, 1 deletions

diff --git a/unix/moduselect.c b/unix/moduselect.c
index 3b5daa53d6..a17d48f8c1 100644
--- a/unix/moduselect.c
+++ b/unix/moduselect.c
@@ -54,6 +54,7 @@ typedef struct _mp_obj_poll_t {
     unsigned short alloc;
     unsigned short len;
     struct pollfd *entries;
+    mp_obj_t *obj_map;
 } mp_obj_poll_t;
 
 STATIC int get_fd(mp_obj_t fdlike) {
@@ -75,6 +76,7 @@ STATIC int get_fd(mp_obj_t fdlike) {
 /// \method register(obj[, eventmask])
 STATIC mp_obj_t poll_register(size_t n_args, const mp_obj_t *args) {
     mp_obj_poll_t *self = MP_OBJ_TO_PTR(args[0]);
+    bool is_fd = MP_OBJ_IS_INT(args[1]);
     int fd = get_fd(args[1]);
 
     mp_uint_t flags;
@@ -101,11 +103,21 @@ STATIC mp_obj_t poll_register(size_t n_args, const mp_obj_t *args) {
     if (free_slot == NULL) {
         if (self->len >= self->alloc) {
             self->entries = m_renew(struct pollfd, self->entries, self->alloc, self->alloc + 4);
+            if (self->obj_map) {
+                self->obj_map = m_renew(mp_obj_t, self->obj_map, self->alloc, self->alloc + 4);
+            }
             self->alloc += 4;
         }
         free_slot = &self->entries[self->len++];
     }
 
+    if (!is_fd) {
+        if (self->obj_map == NULL) {
+            self->obj_map = m_new0(mp_obj_t, self->alloc);
+        }
+        self->obj_map[free_slot - self->entries] = args[1];
+    }
+
     free_slot->fd = fd;
     free_slot->events = flags;
     free_slot->revents = 0;
@@ -121,6 +133,9 @@ STATIC mp_obj_t poll_unregister(mp_obj_t self_in, mp_obj_t obj_in) {
     for (int i = self->len - 1; i >= 0; i--) {
         if (entries->fd == fd) {
             entries->fd = -1;
+            if (self->obj_map) {
+                self->obj_map[entries - self->entries] = NULL;
+            }
             break;
         }
         entries++;
@@ -181,7 +196,12 @@ STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) {
     for (int i = 0; i < self->len; i++, entries++) {
         if (entries->revents != 0) {
             mp_obj_tuple_t *t = MP_OBJ_TO_PTR(mp_obj_new_tuple(2, NULL));
-            t->items[0] = MP_OBJ_NEW_SMALL_INT(entries->fd);
+            // If there's an object stored, return it, otherwise raw fd
+            if (self->obj_map && self->obj_map[i] != NULL) {
+                t->items[0] = self->obj_map[i];
+            } else {
+                t->items[0] = MP_OBJ_NEW_SMALL_INT(entries->fd);
+            }
             t->items[1] = MP_OBJ_NEW_SMALL_INT(entries->revents);
             ret_list->items[ret_i++] = MP_OBJ_FROM_PTR(t);
             if (flags & FLAG_ONESHOT) {
@@ -218,6 +238,7 @@ STATIC mp_obj_t select_poll(size_t n_args, const mp_obj_t *args) {
     poll->entries = m_new(struct pollfd, alloc);
     poll->alloc = alloc;
     poll->len = 0;
+    poll->obj_map = NULL;
     return MP_OBJ_FROM_PTR(poll);
 }
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_select_poll_obj, 0, 1, select_poll);