py: Fix bug in map lookup of interned string vs non-interned.

dpgeorge · dpgeorge · commit 186e463a9ea5 · 2014-04-28T12:11:57.000+01:00
Had choice of either interning or forcing full equality comparison, and chose latter. See comments in mp_map_lookup. Addresses issue adafruit#523.
diff --git a/py/map.c b/py/map.c
@@ -98,13 +98,33 @@ STATIC void mp_map_rehash(mp_map_t *map) {
 // MP_MAP_LOOKUP_REMOVE_IF_FOUND behaviour:
 //  - returns NULL if not found, else the slot if was found in with key null and value non-null
 mp_map_elem_t* mp_map_lookup(mp_map_t *map, mp_obj_t index, mp_map_lookup_kind_t lookup_kind) {
+
+    // Work out if we can compare just pointers
+    bool compare_only_ptrs = map->all_keys_are_qstrs;
+    if (compare_only_ptrs) {
+        if (MP_OBJ_IS_QSTR(index)) {
+            // Index is a qstr, so can just do ptr comparison.
+        } else if (MP_OBJ_IS_TYPE(index, &mp_type_str)) {
+            // Index is a non-interned string.
+            // We can either intern the string, or force a full equality comparison.
+            // We chose the latter, since interning costs time and potentially RAM,
+            // and it won't necessarily benefit subsequent calls because these calls
+            // most likely won't pass the newly-interned string.
+            compare_only_ptrs = false;
+        } else if (!(lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)) {
+            // If we are not adding, then we can return straight away a failed
+            // lookup because we know that the index will never be found.
+            return NULL;
+        }
+    }
+
     // if the map is a fixed array then we must do a brute force linear search
     if (map->table_is_fixed_array) {
         if (lookup_kind != MP_MAP_LOOKUP) {
             return NULL;
         }
         for (mp_map_elem_t *elem = &map->table[0], *top = &map->table[map->used]; elem < top; elem++) {
-            if (elem->key == index || (!map->all_keys_are_qstrs && mp_obj_equal(elem->key, index))) {
+            if (elem->key == index || (!compare_only_ptrs && mp_obj_equal(elem->key, index))) {
                 return elem;
             }
         }
@@ -148,7 +168,7 @@ mp_map_elem_t* mp_map_lookup(mp_map_t *map, mp_obj_t index, mp_map_lookup_kind_t
             if (avail_slot == NULL) {
                 avail_slot = slot;
             }
-        } else if (slot->key == index || (!map->all_keys_are_qstrs && mp_obj_equal(slot->key, index))) {
+        } else if (slot->key == index || (!compare_only_ptrs && mp_obj_equal(slot->key, index))) {
             // found index
             // Note: CPython does not replace the index; try x={True:'true'};x[1]='one';x
             if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
diff --git a/tests/basics/dict-intern.py b/tests/basics/dict-intern.py
@@ -0,0 +1,15 @@
+# check that interned strings are compared against non-interned strings
+
+di = {"key1": "value"}
+
+# lookup interned string
+k = "key1"
+print(k in di)
+
+# lookup non-interned string
+k2 = "key" + "1"
+print(k == k2)
+print(k2 in di)
+
+# lookup non-interned string
+print("".join(['k', 'e', 'y', '1']) in di)
