Merge remote-tracking branch 'upstream/master' into containment

1 files changed, 92 insertions, 0 deletions

diff --git a/py/objlist.c b/py/objlist.c
index 7d0402d28..fa8ec67d0 100644
--- a/py/objlist.c
+++ b/py/objlist.c
@@ -21,6 +21,7 @@ typedef struct _mp_obj_list_t {
 
 static mp_obj_t mp_obj_new_list_iterator(mp_obj_list_t *list, int cur);
 static mp_obj_list_t *list_new(uint n);
+static mp_obj_t list_extend(mp_obj_t self_in, mp_obj_t arg_in);
 
 /******************************************************************************/
 /* list                                                                       */
@@ -61,6 +62,61 @@ static mp_obj_t list_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args
     return NULL;
 }
 
+// Don't pass RT_COMPARE_OP_NOT_EQUAL here
+static bool list_cmp_helper(int op, mp_obj_t self_in, mp_obj_t another_in) {
+    assert(MP_OBJ_IS_TYPE(self_in, &list_type));
+    if (!MP_OBJ_IS_TYPE(another_in, &list_type)) {
+        return false;
+    }
+    mp_obj_list_t *self = self_in;
+    mp_obj_list_t *another = another_in;
+    if (op == RT_COMPARE_OP_EQUAL && self->len != another->len) {
+        return false;
+    }
+
+    // Let's deal only with > & >=
+    if (op == RT_COMPARE_OP_LESS || op == RT_COMPARE_OP_LESS_EQUAL) {
+        mp_obj_t t = self;
+        self = another;
+        another = t;
+        if (op == RT_COMPARE_OP_LESS) {
+            op = RT_COMPARE_OP_MORE;
+        } else {
+            op = RT_COMPARE_OP_MORE_EQUAL;
+        }
+    }
+
+    int len = self->len < another->len ? self->len : another->len;
+    bool eq_status = true; // empty lists are equal
+    bool rel_status;
+    for (int i = 0; i < len; i++) {
+        eq_status = mp_obj_equal(self->items[i], another->items[i]);
+        if (op == RT_COMPARE_OP_EQUAL && !eq_status) {
+            return false;
+        }
+        rel_status = (rt_binary_op(op, self->items[i], another->items[i]) == mp_const_true);
+        if (!eq_status && !rel_status) {
+            return false;
+        }
+    }
+
+    // If we had tie in the last element...
+    if (eq_status) {
+        // ... and we have lists of different lengths...
+        if (self->len != another->len) {
+            if (self->len < another->len) {
+                // ... then longer list length wins (we deal only with >)
+                return false;
+            }
+        } else if (op == RT_COMPARE_OP_MORE) {
+            // Otherwise, if we have strict relation, equality means failure
+            return false;
+        }
+    }
+
+    return true;
+}
+
 static mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
     mp_obj_list_t *o = lhs;
     switch (op) {
@@ -81,6 +137,14 @@ static mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
             memcpy(s->items + o->len, p->items, sizeof(mp_obj_t) * p->len);
             return s;
         }
+        case RT_BINARY_OP_INPLACE_ADD:
+        {
+            if (!MP_OBJ_IS_TYPE(rhs, &list_type)) {
+                return NULL;
+            }
+            list_extend(lhs, rhs);
+            return o;
+        }
         case RT_BINARY_OP_MULTIPLY:
         {
             if (!MP_OBJ_IS_SMALL_INT(rhs)) {
@@ -96,6 +160,15 @@ static mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
             }
             return s;
         }
+        case RT_COMPARE_OP_EQUAL:
+        case RT_COMPARE_OP_LESS:
+        case RT_COMPARE_OP_LESS_EQUAL:
+        case RT_COMPARE_OP_MORE:
+        case RT_COMPARE_OP_MORE_EQUAL:
+            return MP_BOOL(list_cmp_helper(op, lhs, rhs));
+        case RT_COMPARE_OP_NOT_EQUAL:
+            return MP_BOOL(!list_cmp_helper(RT_COMPARE_OP_EQUAL, lhs, rhs));
+
         default:
             // op not supported
             return NULL;
@@ -117,6 +190,23 @@ mp_obj_t mp_obj_list_append(mp_obj_t self_in, mp_obj_t arg) {
     return mp_const_none; // return None, as per CPython
 }
 
+static mp_obj_t list_extend(mp_obj_t self_in, mp_obj_t arg_in) {
+    assert(MP_OBJ_IS_TYPE(self_in, &list_type));
+    assert(MP_OBJ_IS_TYPE(arg_in, &list_type));
+    mp_obj_list_t *self = self_in;
+    mp_obj_list_t *arg = arg_in;
+
+    if (self->len + arg->len > self->alloc) {
+        // TODO: use alloc policy for "4"
+        self->items = m_renew(mp_obj_t, self->items, self->alloc, self->len + arg->len + 4);
+        self->alloc = self->len + arg->len + 4;
+    }
+
+    memcpy(self->items + self->len, arg->items, sizeof(mp_obj_t) * arg->len);
+    self->len += arg->len;
+    return mp_const_none; // return None, as per CPython
+}
+
 static mp_obj_t list_pop(int n_args, const mp_obj_t *args) {
     assert(1 <= n_args && n_args <= 2);
     assert(MP_OBJ_IS_TYPE(args[0], &list_type));
@@ -281,6 +371,7 @@ static mp_obj_t list_reverse(mp_obj_t self_in) {
 }
 
 static MP_DEFINE_CONST_FUN_OBJ_2(list_append_obj, mp_obj_list_append);
+static MP_DEFINE_CONST_FUN_OBJ_2(list_extend_obj, list_extend);
 static MP_DEFINE_CONST_FUN_OBJ_1(list_clear_obj, list_clear);
 static MP_DEFINE_CONST_FUN_OBJ_1(list_copy_obj, list_copy);
 static MP_DEFINE_CONST_FUN_OBJ_2(list_count_obj, list_count);
@@ -296,6 +387,7 @@ static const mp_method_t list_type_methods[] = {
     { "clear", &list_clear_obj },
     { "copy", &list_copy_obj },
     { "count", &list_count_obj },
+    { "extend", &list_extend_obj },
     { "index", &list_index_obj },
     { "insert", &list_insert_obj },
     { "pop", &list_pop_obj },