improve `dict` and add `lru_cache`

docs/modules/functools.md

@@ -7,6 +7,10 @@ label: functools
 
 A decorator that caches a function's return value each time it is called. If called later with the same arguments, the cached value is returned, and not re-evaluated.
 
+### `functools.lru_cache(maxsize=128)`
+
+A decorator that wraps a function with a memoizing callable that saves up to the maxsize most recent calls.
+
 ### `functools.reduce(function, sequence, initial=...)`
 
 Apply a function of two arguments cumulatively to the items of a sequence, from left to right, so as to reduce the sequence to a single value. For example, `functools.reduce(lambda x, y: x+y, [1, 2, 3, 4, 5])` calculates `((((1+2)+3)+4)+5)`. The left argument, `x`, is the accumulated value and the right argument, `y`, is the update value from the sequence. If the optional `initial` is present, it is placed before the items of the sequence in the calculation, and serves as a default when the sequence is empty.

include/pocketpy/pocketpy.h

@@ -735,7 +735,7 @@ enum py_PredefinedType {
     tp_locals,
     tp_code,
     tp_dict,
-    tp_dict_items,    // 1 slot
+    tp_dict_iterator, // 1 slot
     tp_property,      // 2 slots (getter + setter)
     tp_star_wrapper,  // 1 slot + int level
     tp_staticmethod,  // 1 slot

python/functools.py

@@ -8,6 +8,26 @@ class cache:
             self.cache[args] = self.f(*args)
         return self.cache[args]
     
+class lru_cache:
+    def __init__(self, maxsize=128):
+        self.maxsize = maxsize
+        self.cache = {}
+
+    def __call__(self, f):
+        def wrapped(*args):
+            if args in self.cache:
+                res = self.cache.pop(args)
+                self.cache[args] = res
+                return res
+            
+            res = f(*args)
+            if len(self.cache) >= self.maxsize:
+                first_key = next(iter(self.cache))
+                self.cache.pop(first_key)
+            self.cache[args] = res
+            return res
+        return wrapped
+    
 def reduce(function, sequence, initial=...):
     it = iter(sequence)
     if initial is ...:

src/interpreter/vm.c

@@ -132,7 +132,7 @@ void VM__ctor(VM* self) {
     validate(tp_code, pk_code__register());
 
     validate(tp_dict, pk_dict__register());
-    validate(tp_dict_items, pk_dict_items__register());
+    validate(tp_dict_iterator, pk_dict_items__register());
 
     validate(tp_property, pk_property__register());
     validate(tp_star_wrapper, pk_newtype("star_wrapper", tp_object, NULL, NULL, false, true));

src/public/py_dict.c

@@ -56,6 +56,7 @@ static uint32_t Dict__next_cap(uint32_t cap) {
 typedef struct {
     DictEntry* curr;
     DictEntry* end;
+    int mode;  // 0: keys, 1: values, 2: items
 } DictIterator;
 
 static void Dict__ctor(Dict* self, uint32_t capacity, int entries_capacity) {
@@ -241,9 +242,10 @@ static int Dict__pop(Dict* self, py_Ref key) {
     return 0;
 }
 
-static void DictIterator__ctor(DictIterator* self, Dict* dict) {
+static void DictIterator__ctor(DictIterator* self, Dict* dict, int mode) {
     self->curr = dict->entries.data;
     self->end = self->curr + dict->entries.length;
+    self->mode = mode;
 }
 
 static DictEntry* DictIterator__next(DictIterator* self) {
@@ -377,7 +379,7 @@ static bool dict__eq__(int argc, py_Ref argv) {
         return true;
     }
     DictIterator iter;
-    DictIterator__ctor(&iter, self);
+    DictIterator__ctor(&iter, self, 2);
     // for each self key
     while(1) {
         DictEntry* entry = DictIterator__next(&iter);
@@ -463,48 +465,34 @@ static bool dict_pop(int argc, py_Ref argv) {
     py_Ref default_val = argc == 3 ? py_arg(2) : py_None();
     int res = Dict__pop(self, py_arg(1));
     if(res == -1) return false;
-    if(res == 0) { py_assign(py_retval(), default_val); }
+    if(res == 0) py_assign(py_retval(), default_val);
     return true;
 }
 
-static bool dict_items(int argc, py_Ref argv) {
+static bool dict_keys(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     Dict* self = py_touserdata(argv);
-    DictIterator* ud = py_newobject(py_retval(), tp_dict_items, 1, sizeof(DictIterator));
-    DictIterator__ctor(ud, self);
+    DictIterator* ud = py_newobject(py_retval(), tp_dict_iterator, 1, sizeof(DictIterator));
+    DictIterator__ctor(ud, self, 0);
     py_setslot(py_retval(), 0, argv);  // keep a reference to the dict
     return true;
 }
 
-static bool dict_keys(int argc, py_Ref argv) {
+static bool dict_values(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     Dict* self = py_touserdata(argv);
-    py_Ref p = py_newtuple(py_retval(), self->length);
-    DictIterator iter;
-    DictIterator__ctor(&iter, self);
-    int i = 0;
-    while(1) {
-        DictEntry* entry = DictIterator__next(&iter);
-        if(!entry) break;
-        p[i++] = entry->key;
-    }
-    assert(i == self->length);
+    DictIterator* ud = py_newobject(py_retval(), tp_dict_iterator, 1, sizeof(DictIterator));
+    DictIterator__ctor(ud, self, 1);
+    py_setslot(py_retval(), 0, argv);  // keep a reference to the dict
     return true;
 }
 
-static bool dict_values(int argc, py_Ref argv) {
+static bool dict_items(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     Dict* self = py_touserdata(argv);
-    py_Ref p = py_newtuple(py_retval(), self->length);
-    DictIterator iter;
-    DictIterator__ctor(&iter, self);
-    int i = 0;
-    while(1) {
-        DictEntry* entry = DictIterator__next(&iter);
-        if(!entry) break;
-        p[i++] = entry->val;
-    }
-    assert(i == self->length);
+    DictIterator* ud = py_newobject(py_retval(), tp_dict_iterator, 1, sizeof(DictIterator));
+    DictIterator__ctor(ud, self, 2);
+    py_setslot(py_retval(), 0, argv);  // keep a reference to the dict
     return true;
 }
 
@@ -521,15 +509,16 @@ py_Type pk_dict__register() {
     py_bindmagic(type, __repr__, dict__repr__);
     py_bindmagic(type, __eq__, dict__eq__);
     py_bindmagic(type, __ne__, dict__ne__);
+    py_bindmagic(type, __iter__, dict_keys);
 
     py_bindmethod(type, "clear", dict_clear);
     py_bindmethod(type, "copy", dict_copy);
     py_bindmethod(type, "update", dict_update);
     py_bindmethod(type, "get", dict_get);
     py_bindmethod(type, "pop", dict_pop);
-    py_bindmethod(type, "items", dict_items);
     py_bindmethod(type, "keys", dict_keys);
     py_bindmethod(type, "values", dict_values);
+    py_bindmethod(type, "items", dict_items);
 
     py_setdict(py_tpobject(type), __hash__, py_None());
     return type;
@@ -541,16 +530,28 @@ static bool dict_items__next__(int argc, py_Ref argv) {
     DictIterator* iter = py_touserdata(py_arg(0));
     DictEntry* entry = (DictIterator__next(iter));
     if(entry) {
-        py_Ref p = py_newtuple(py_retval(), 2);
-        p[0] = entry->key;
-        p[1] = entry->val;
-        return true;
+        switch(iter->mode) {
+            case 0:  // keys
+                py_assign(py_retval(), &entry->key);
+                return true;
+            case 1:  // values
+                py_assign(py_retval(), &entry->val);
+                return true;
+            case 2:  // items
+            {
+                py_Ref p = py_newtuple(py_retval(), 2);
+                p[0] = entry->key;
+                p[1] = entry->val;
+                return true;
+            }
+            default: c11__unreachable();
+        }
     }
     return StopIteration();
 }
 
 py_Type pk_dict_items__register() {
-    py_Type type = pk_newtype("dict_items", tp_object, NULL, NULL, false, true);
+    py_Type type = pk_newtype("dict_iterator", tp_object, NULL, NULL, false, true);
     py_bindmagic(type, __iter__, pk_wrapper__self);
     py_bindmagic(type, __next__, dict_items__next__);
     return type;

tests/08_dict.py

@@ -32,8 +32,8 @@ assert tinydict == updated_dict
 
 dishes = {'eggs': 2, 'sausage': 1, 'bacon': 1, 'spam': 500}
 # dict is now ordered
-assert dishes.keys() == ('eggs', 'sausage', 'bacon', 'spam')
-assert dishes.values() == (2, 1, 1, 500)
+assert list(dishes.keys()) == ['eggs', 'sausage', 'bacon', 'spam']
+assert list(dishes.values()) == [2, 1, 1, 500]
 
 d={1:"a",2:"b",3:"c"}
 result=[]
@@ -44,8 +44,8 @@ assert len(result) == 6
 
 del d[2]
 assert len(d) == 2
-assert d.keys() == (1, 3)
-assert d.values() == ('a', 'c')
+assert list(d.keys()) == [1, 3]
+assert list(d.values()) == ['a', 'c']
 del d[1]
 del d[3]
 assert len(d) == 0
@@ -77,11 +77,11 @@ a = {'g': 0}
 
 a['ball_3'] = 0
 a['ball_4'] = 0
-assert a.keys() == ('g', 'ball_3', 'ball_4')
+assert list(a.keys()) == ['g', 'ball_3', 'ball_4']
 del a['ball_3']
-assert a.keys() == ('g', 'ball_4')
+assert list(a.keys()) == ['g', 'ball_4']
 del a['ball_4']
-assert a.keys() == ('g',)
+assert list(a.keys()) == ['g',]
 del a['g']
 assert len(a) == 0
 
@@ -147,4 +147,11 @@ for i in range(-1000, 1000):
 e = {}
 for i in range(-10000, 10000, 3):
     e[i] = i
-    assert e[i] == i
+    assert e[i] == i
+
+# test iter
+d = {'1': 1, 222: 2, '333': 3}
+assert list(d) == ['1', 222, '333']
+assert list(d.keys()) == ['1', 222, '333']
+assert list(d.values()) == [1, 2, 3]
+assert list(d.items()) == [('1', 1), (222, 2), ('333', 3)]

tests/73_functools.py

@@ -23,3 +23,34 @@ sub_10 = partial(sub, b=10)
 assert sub_10(20) == 10
 assert sub_10(30) == 20
 
+# test lru_cache
+from functools import lru_cache
+miss_keys = []
+
+@lru_cache(maxsize=3)
+def test_f(x: int):
+    miss_keys.append(x)
+    return x
+
+assert test_f(1) == 1 and miss_keys == [1]
+# [1]
+assert test_f(2) == 2 and miss_keys == [1, 2]
+# [1, 2]
+assert test_f(3) == 3 and miss_keys == [1, 2, 3]
+# [1, 2, 3]
+assert test_f(1) == 1 and miss_keys == [1, 2, 3]
+# [2, 3, 1]
+assert test_f(2) == 2 and miss_keys == [1, 2, 3]
+# [3, 1, 2]
+assert test_f(4) == 4 and miss_keys == [1, 2, 3, 4]
+# [1, 2, 4]
+assert test_f(3) == 3 and miss_keys == [1, 2, 3, 4, 3]
+# [2, 4, 3]
+assert test_f(2) == 2 and miss_keys == [1, 2, 3, 4, 3]
+# [4, 3, 2]
+assert test_f(5) == 5 and miss_keys == [1, 2, 3, 4, 3, 5]
+# [3, 2, 5]
+assert test_f(3) == 3 and miss_keys == [1, 2, 3, 4, 3, 5]
+# [2, 5, 3]
+assert test_f(1) == 1 and miss_keys == [1, 2, 3, 4, 3, 5, 1]
+# [5, 3, 1]