From 09a0c64bce93f5ebcea82e81b4b07ddd7ff76cc7 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Thu, 30 Jan 2014 10:05:33 +0000
Subject: [PATCH] py: Improve __bool__ and __len__ dispatch; add slots for
 them.

---
 py/obj.c                       | 25 ++++---------------------
 py/objdict.c                   |  2 +-
 py/objtype.c                   | 24 ++++++++++++++++++++++++
 py/qstrdefs.h                  |  2 ++
 py/runtime.c                   | 16 ++++++++--------
 tests/basics/slots_bool_len.py | 17 +++++++++++++++++
 6 files changed, 56 insertions(+), 30 deletions(-)
 create mode 100644 tests/basics/slots_bool_len.py

diff --git a/py/obj.c b/py/obj.c
index 983718004b..33d64c894f 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -236,33 +236,16 @@ uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index)
 
 // may return MP_OBJ_NULL
 mp_obj_t mp_obj_len_maybe(mp_obj_t o_in) {
-    mp_small_int_t len = 0;
     if (MP_OBJ_IS_STR(o_in)) {
-        len = mp_obj_str_get_len(o_in);
-    } else if (MP_OBJ_IS_TYPE(o_in, &tuple_type)) {
-        uint seq_len;
-        mp_obj_t *seq_items;
-        mp_obj_tuple_get(o_in, &seq_len, &seq_items);
-        len = seq_len;
-    } else if (MP_OBJ_IS_TYPE(o_in, &list_type)) {
-        uint seq_len;
-        mp_obj_t *seq_items;
-        mp_obj_list_get(o_in, &seq_len, &seq_items);
-        len = seq_len;
-    } else if (MP_OBJ_IS_TYPE(o_in, &dict_type)) {
-        len = mp_obj_dict_len(o_in);
+        return MP_OBJ_NEW_SMALL_INT((machine_int_t)mp_obj_str_get_len(o_in));
     } else {
         mp_obj_type_t *type = mp_obj_get_type(o_in);
         if (type->unary_op != NULL) {
-            mp_obj_t result = type->unary_op(RT_UNARY_OP_LEN, o_in);
-            if (result != MP_OBJ_NULL) {
-                return result;
-            }
+            return type->unary_op(RT_UNARY_OP_LEN, o_in);
+        } else {
+            return MP_OBJ_NULL;
         }
-
-        return MP_OBJ_NULL;
     }
-    return MP_OBJ_NEW_SMALL_INT(len);
 }
 
 // Return input argument. Useful as .getiter for objects which are
diff --git a/py/objdict.c b/py/objdict.c
index f78a7358b2..934eb50650 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -47,7 +47,7 @@ static mp_obj_t dict_unary_op(int op, mp_obj_t self_in) {
     mp_obj_dict_t *self = self_in;
     switch (op) {
         case RT_UNARY_OP_BOOL: return MP_BOOL(self->map.used != 0);
-        case RT_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT(self->map.used);
+        case RT_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT((machine_int_t)self->map.used);
         default: return MP_OBJ_NULL; // op not supported for None
     }
 }
diff --git a/py/objtype.c b/py/objtype.c
index 9cb9b8654d..45992b23a3 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -116,6 +116,29 @@ static mp_obj_t class_make_new(mp_obj_t self_in, uint n_args, uint n_kw, const m
     return o;
 }
 
+static const qstr unary_op_method_name[] = {
+    [RT_UNARY_OP_BOOL] = MP_QSTR___bool__,
+    [RT_UNARY_OP_LEN] = MP_QSTR___len__,
+    //[RT_UNARY_OP_POSITIVE,
+    //[RT_UNARY_OP_NEGATIVE,
+    //[RT_UNARY_OP_INVERT,
+    [RT_UNARY_OP_NOT] = MP_QSTR_, // not implemented, used to make sure array has full size
+};
+
+static mp_obj_t class_unary_op(int op, mp_obj_t self_in) {
+    mp_obj_class_t *self = self_in;
+    qstr op_name = unary_op_method_name[op];
+    if (op_name == 0) {
+        return MP_OBJ_NULL;
+    }
+    mp_obj_t member = mp_obj_class_lookup(self->base.type, op_name);
+    if (member != MP_OBJ_NULL) {
+        return rt_call_function_1(member, self_in);
+    } else {
+        return MP_OBJ_NULL;
+    }
+}
+
 static const qstr binary_op_method_name[] = {
     [RT_BINARY_OP_SUBSCR] = MP_QSTR___getitem__,
     /*
@@ -330,6 +353,7 @@ mp_obj_t mp_obj_new_type(const char *name, mp_obj_t bases_tuple, mp_obj_t locals
     o->name = name;
     o->print = class_print;
     o->make_new = class_make_new;
+    o->unary_op = class_unary_op;
     o->binary_op = class_binary_op;
     o->load_attr = class_load_attr;
     o->store_attr = class_store_attr;
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 81706841cd..50b426f1a4 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -13,6 +13,8 @@ Q(__next__)
 Q(__qualname__)
 Q(__repl_print__)
 
+Q(__bool__)
+Q(__len__)
 Q(__getitem__)
 Q(__add__)
 Q(__sub__)
diff --git a/py/runtime.c b/py/runtime.c
index 6b3c8dc1d0..b1027de932 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -304,23 +304,23 @@ void rt_assign_inline_asm_code(uint unique_code_id, void *fun, uint len, int n_a
 
 int rt_is_true(mp_obj_t arg) {
     DEBUG_OP_printf("is true %p\n", arg);
-    if (MP_OBJ_IS_SMALL_INT(arg)) {
+    if (arg == mp_const_false) {
+        return 0;
+    } else if (arg == mp_const_true) {
+        return 1;
+    } else if (arg == mp_const_none) {
+        return 0;
+    } else if (MP_OBJ_IS_SMALL_INT(arg)) {
         if (MP_OBJ_SMALL_INT_VALUE(arg) == 0) {
             return 0;
         } else {
             return 1;
         }
-    } else if (arg == mp_const_none) {
-        return 0;
-    } else if (arg == mp_const_false) {
-        return 0;
-    } else if (arg == mp_const_true) {
-        return 1;
     } else {
         mp_obj_type_t *type = mp_obj_get_type(arg);
         if (type->unary_op != NULL) {
             mp_obj_t result = type->unary_op(RT_UNARY_OP_BOOL, arg);
-            if (result != NULL) {
+            if (result != MP_OBJ_NULL) {
                 return result == mp_const_true;
             }
         }
diff --git a/tests/basics/slots_bool_len.py b/tests/basics/slots_bool_len.py
new file mode 100644
index 0000000000..481fe9b4e6
--- /dev/null
+++ b/tests/basics/slots_bool_len.py
@@ -0,0 +1,17 @@
+class A:
+    def __bool__(self):
+        print('__bool__')
+        return True
+    def __len__(self):
+        print('__len__')
+        return 1
+
+class B:
+    def __len__(self):
+        print('__len__')
+        return 0
+
+print(bool(A()))
+print(len(A()))
+print(bool(B()))
+print(len(B()))