py: Fix up object equality test.

It regressed a bit after implementing float/complex equality.  Now it
should be improved, and support more equality tests.

2 files changed, 27 insertions, 10 deletions

diff --git a/py/obj.c b/py/obj.c
index e7957e21ff..0cd8117819 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -143,20 +143,35 @@ machine_int_t mp_obj_hash(mp_obj_t o_in) {
 bool mp_obj_equal(mp_obj_t o1, mp_obj_t o2) {
     if (o1 == o2) {
         return true;
-    } else if (o1 == mp_const_none || o2 == mp_const_none) {
+    }
+    if (o1 == mp_const_none || o2 == mp_const_none) {
         return false;
-    } else if (MP_OBJ_IS_SMALL_INT(o1) || MP_OBJ_IS_SMALL_INT(o2)) {
-        if (MP_OBJ_IS_SMALL_INT(o1) && MP_OBJ_IS_SMALL_INT(o2)) {
+    }
+
+    // fast path for small ints
+    if (MP_OBJ_IS_SMALL_INT(o1)) {
+        if (MP_OBJ_IS_SMALL_INT(o2)) {
+            // both SMALL_INT, and not equal if we get here
             return false;
         } else {
-            if (MP_OBJ_IS_SMALL_INT(o1)) {
-                mp_obj_t temp = o2; o2 = o1; o1 = temp;
-            }
-            // o2 is the SMALL_INT, o1 is not
+            mp_obj_t temp = o2; o2 = o1; o1 = temp;
+            // o2 is now the SMALL_INT, o1 is not
             // fall through to generic op
         }
-    } else if (MP_OBJ_IS_STR(o1) && MP_OBJ_IS_STR(o2)) {
-        return mp_obj_str_equal(o1, o2);
+    }
+
+    // fast path for strings
+    if (MP_OBJ_IS_STR(o1)) {
+        if (MP_OBJ_IS_STR(o2)) {
+            // both strings, use special function
+            return mp_obj_str_equal(o1, o2);
+        } else {
+            // a string is never equal to anything else
+            return false;
+        }
+    } else if (MP_OBJ_IS_STR(o2)) {
+        // o1 is not a string (else caught above), so the objects are not equal
+        return false;
     }
 
     // generic type, call binary_op(MP_BINARY_OP_EQUAL)
diff --git a/py/objdict.c b/py/objdict.c
index 963e188074..8bdd0026a5 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -99,7 +99,6 @@ STATIC mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
             return MP_BOOL(elem != NULL);
         }
         case MP_BINARY_OP_EQUAL: {
-            // TODO: Support equality to other object types
             if (MP_OBJ_IS_TYPE(rhs_in, &mp_type_dict)) {
                 mp_obj_dict_t *rhs = rhs_in;
                 if (o->map.used != rhs->map.used) {
@@ -118,6 +117,9 @@ STATIC mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
                     }
                 }
                 return mp_const_true;
+            } else {
+                // dict is not equal to instance of any other type
+                return mp_const_false;
             }
         }
         default: