From 57fce3bdb203e9701dbd81ee108189898e19911b Mon Sep 17 00:00:00 2001
From: Jim Mussared <jim.mussared@gmail.com>
Date: Thu, 23 Apr 2020 01:10:30 +1000
Subject: [PATCH] py/objdict: Fix popitem for ordered dicts.

The popitem method wasn't implemented for ordered dicts and would result in
an invalid state.

Fixes issue #5956.
---
 py/map.c                     |  1 +
 py/obj.h                     |  3 +++
 py/objdict.c                 | 16 ++++++++++++----
 tests/basics/ordereddict1.py | 20 ++++++++++++++++++++
 4 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/py/map.c b/py/map.c
index 2570941bf1..676c364da7 100644
--- a/py/map.c
+++ b/py/map.c
@@ -177,6 +177,7 @@ mp_map_elem_t *mp_map_lookup(mp_map_t *map, mp_obj_t index, mp_map_lookup_kind_t
                     --map->used;
                     memmove(elem, elem + 1, (top - elem - 1) * sizeof(*elem));
                     // put the found element after the end so the caller can access it if needed
+                    // note: caller must NULL the value so the GC can clean up (e.g. see dict_get_helper).
                     elem = &map->table[map->used];
                     elem->key = MP_OBJ_NULL;
                     elem->value = value;
diff --git a/py/obj.h b/py/obj.h
index d3abdab591..125acf118f 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -26,6 +26,8 @@
 #ifndef MICROPY_INCLUDED_PY_OBJ_H
 #define MICROPY_INCLUDED_PY_OBJ_H
 
+#include <assert.h>
+
 #include "py/mpconfig.h"
 #include "py/misc.h"
 #include "py/qstr.h"
@@ -423,6 +425,7 @@ typedef enum _mp_map_lookup_kind_t {
 extern const mp_map_t mp_const_empty_map;
 
 static inline bool mp_map_slot_is_filled(const mp_map_t *map, size_t pos) {
+    assert(pos < map->alloc);
     return (map)->table[pos].key != MP_OBJ_NULL && (map)->table[pos].key != MP_OBJ_SENTINEL;
 }
 
diff --git a/py/objdict.c b/py/objdict.c
index 2a72c62524..7690eeab29 100644
--- a/py/objdict.c
+++ b/py/objdict.c
@@ -44,13 +44,15 @@ STATIC mp_map_elem_t *dict_iter_next(mp_obj_dict_t *dict, size_t *cur) {
     size_t max = dict->map.alloc;
     mp_map_t *map = &dict->map;
 
-    for (size_t i = *cur; i < max; i++) {
+    size_t i = *cur;
+    for (; i < max; i++) {
         if (mp_map_slot_is_filled(map, i)) {
             *cur = i + 1;
             return &(map->table[i]);
         }
     }
 
+    assert(map->used == 0 || i == max);
     return NULL;
 }
 
@@ -321,11 +323,17 @@ STATIC mp_obj_t dict_popitem(mp_obj_t self_in) {
     mp_check_self(mp_obj_is_dict_type(self_in));
     mp_obj_dict_t *self = MP_OBJ_TO_PTR(self_in);
     mp_ensure_not_fixed(self);
-    size_t cur = 0;
-    mp_map_elem_t *next = dict_iter_next(self, &cur);
-    if (next == NULL) {
+    if (self->map.used == 0) {
         mp_raise_msg(&mp_type_KeyError, MP_ERROR_TEXT("popitem(): dictionary is empty"));
     }
+    size_t cur = 0;
+    #if MICROPY_PY_COLLECTIONS_ORDEREDDICT
+    if (self->map.is_ordered) {
+        cur = self->map.used - 1;
+    }
+    #endif
+    mp_map_elem_t *next = dict_iter_next(self, &cur);
+    assert(next);
     self->map.used--;
     mp_obj_t items[] = {next->key, next->value};
     next->key = MP_OBJ_SENTINEL; // must mark key as sentinel to indicate that it was deleted
diff --git a/tests/basics/ordereddict1.py b/tests/basics/ordereddict1.py
index d1633f0bb0..270deab384 100644
--- a/tests/basics/ordereddict1.py
+++ b/tests/basics/ordereddict1.py
@@ -24,3 +24,23 @@ d["abc"] = 123
 print(len(d))
 print(list(d.keys()))
 print(list(d.values()))
+
+# pop an element
+print(d.popitem())
+print(len(d))
+print(list(d.keys()))
+print(list(d.values()))
+
+# add an element after popping
+d["xyz"] = 321
+print(len(d))
+print(list(d.keys()))
+print(list(d.values()))
+
+# pop until empty
+print(d.popitem())
+print(d.popitem())
+try:
+    d.popitem()
+except:
+    print('empty')