weakref: Don't leak cyclic weakrefs.

docs/library/index.rst

@@ -81,6 +81,7 @@ library.
    struct.rst
    sys.rst
    time.rst
+   weakref.rst
    zlib.rst
    _thread.rst
 

docs/library/weakref.rst

@@ -0,0 +1,74 @@
+:mod:`weakref` -- Python object lifetime management
+===================================================
+
+.. module:: weakref
+   :synopsis: Create weak references to Python objects
+
+|see_cpython_module| :mod:`python:weakref`.
+
+This module allows creation of weak references to Python objects.  A weak reference
+is a non-traceable reference to a heap-allocated Python object, so the garbage
+collector can still reclaim the object even though the weak reference refers to it.
+
+Python callbacks can be registered to be called when an object is reclaimed by the
+garbage collector.  This provides a safe way to clean up when objects are no longer
+needed.
+
+ref objects
+-----------
+
+A ref object is the simplest way to make a weak reference.
+
+.. class:: ref(object [, callback], /)
+
+   Return a weak reference to the given *object*.
+
+   If *callback* is given and is not ``None`` then, when *object* is reclaimed
+   by the garbage collector and if the weak reference object is still alive, the
+   *callback* will be called.  The *callback* will be passed the weak reference
+   object as its single argument.
+
+.. method:: ref.__call__()
+
+   Calling the weak reference object will return its referenced object if that
+   object is still alive.  Otherwise ``None`` will be returned.
+
+finalize objects
+----------------
+
+A finalize object is an extended version of a ref object that is more convenient to
+use, and allows more control over the callback.
+
+.. class:: finalize(object, callback, /, *args, **kwargs)
+
+   Return a weak reference to the given *object*.  In contrast to *weakref.ref*
+   objects, finalize objects are held onto internally and will not be collected until
+   *object* is collected.
+
+   A finalize object starts off alive.  It transitions to the dead state when the
+   finalize object is called, either explicitly or when *object* is collected.  It also
+   transitions to dead if the `finalize.detach()` method is called.
+
+   When *object* is reclaimed by the garbage collector (or the finalize object is
+   explicitly called by user code) and the finalize object is still in the alive state,
+   the *callback* will be called.  The *callback* will be passed arguments as:
+   ``callback(*args, **kwargs)``.
+
+.. method:: finalize.__call__()
+
+   If the finalize object is alive then it transitions to the dead state and returns
+   the value of ``callback(*args, **kwargs)``.  Otherwise ``None`` will be returned.
+
+.. method:: finalize.alive
+
+   Read-only boolean attribute that indicates if the finalizer is in the alive state.
+
+.. method:: finalize.peek()
+
+   If the finalize object is alive then return ``(object, callback, args, kwargs)``.
+   Otherwise return ``None``.
+
+.. method:: finalize.detach()
+
+   If the finalize object is alive then it transitions to the dead state and returns
+   ``(object, callback, args, kwargs)``. Otherwise ``None`` will be returned.

ports/webassembly/Makefile

@@ -130,7 +130,7 @@ OBJ += $(addprefix $(BUILD)/, $(SRC_C:.c=.o))
 ################################################################################
 # Main targets.
 
-.PHONY: all repl min test test_min
+.PHONY: all repl min test test//% test_min
 
 all: $(BUILD)/micropython.mjs
 
@@ -150,6 +150,9 @@ min: $(BUILD)/micropython.min.mjs
 test: $(BUILD)/micropython.mjs $(TOP)/tests/run-tests.py
 	cd $(TOP)/tests && MICROPY_MICROPYTHON_MJS=../ports/webassembly/$< ./run-tests.py -t webassembly
 
+test//%: $(BUILD)/micropython.mjs $(TOP)/tests/run-tests.py
+	cd $(TOP)/tests && MICROPY_MICROPYTHON_MJS=../ports/webassembly/$< ./run-tests.py -t webassembly -i "$*"
+
 test_min: $(BUILD)/micropython.min.mjs $(TOP)/tests/run-tests.py
 	cd $(TOP)/tests && MICROPY_MICROPYTHON_MJS=../ports/webassembly/$< ./run-tests.py -t webassembly
 

ports/webassembly/variants/pyscript/mpconfigvariant.h

@@ -1,3 +1,4 @@
 #define MICROPY_CONFIG_ROM_LEVEL                (MICROPY_CONFIG_ROM_LEVEL_FULL_FEATURES)
 #define MICROPY_GC_SPLIT_HEAP                   (1)
 #define MICROPY_GC_SPLIT_HEAP_AUTO              (1)
+#define MICROPY_PY_WEAKREF                      (1)

py/gc.c

@@ -104,14 +104,21 @@
 #if MICROPY_ENABLE_FINALISER
 // FTB = finaliser table byte
 // if set, then the corresponding block may have a finaliser
-
 #define BLOCKS_PER_FTB (8)
-
 #define FTB_GET(area, block) ((area->gc_finaliser_table_start[(block) / BLOCKS_PER_FTB] >> ((block) & 7)) & 1)
 #define FTB_SET(area, block) do { area->gc_finaliser_table_start[(block) / BLOCKS_PER_FTB] |= (1 << ((block) & 7)); } while (0)
 #define FTB_CLEAR(area, block) do { area->gc_finaliser_table_start[(block) / BLOCKS_PER_FTB] &= (~(1 << ((block) & 7))); } while (0)
 #endif
 
+#if MICROPY_PY_WEAKREF
+// WTB = weakref table byte
+// if set, then the corresponding block may have a weakref in MP_STATE_MEM(mp_weakref_map).
+#define BLOCKS_PER_WTB (8)
+#define WTB_GET(area, block) ((area->gc_weakref_table_start[(block) / BLOCKS_PER_WTB] >> ((block) & 7)) & 1)
+#define WTB_SET(area, block) do { area->gc_weakref_table_start[(block) / BLOCKS_PER_WTB] |= (1 << ((block) & 7)); } while (0)
+#define WTB_CLEAR(area, block) do { area->gc_weakref_table_start[(block) / BLOCKS_PER_WTB] &= (~(1 << ((block) & 7))); } while (0)
+#endif
+
 #if MICROPY_PY_THREAD && !MICROPY_PY_THREAD_GIL
 #define GC_MUTEX_INIT() mp_thread_recursive_mutex_init(&MP_STATE_MEM(gc_mutex))
 #define GC_ENTER() mp_thread_recursive_mutex_lock(&MP_STATE_MEM(gc_mutex), 1)
@@ -138,17 +145,23 @@ static void gc_sweep_free_blocks(void);
 // TODO waste less memory; currently requires that all entries in alloc_table have a corresponding block in pool
 static void gc_setup_area(mp_state_mem_area_t *area, void *start, void *end) {
     // calculate parameters for GC (T=total, A=alloc table, F=finaliser table, P=pool; all in bytes):
-    // T = A + F + P
+    // T = A + F + W + P
     //     F = A * BLOCKS_PER_ATB / BLOCKS_PER_FTB
+    //     W = A * BLOCKS_PER_ATB / BLOCKS_PER_WTB
     //     P = A * BLOCKS_PER_ATB * BYTES_PER_BLOCK
-    // => T = A * (1 + BLOCKS_PER_ATB / BLOCKS_PER_FTB + BLOCKS_PER_ATB * BYTES_PER_BLOCK)
+    // => T = A * (1 + BLOCKS_PER_ATB / BLOCKS_PER_FTB + BLOCKS_PER_ATB / BLOCKS_PER_WTB + BLOCKS_PER_ATB * BYTES_PER_BLOCK)
     size_t total_byte_len = (byte *)end - (byte *)start;
-    #if MICROPY_ENABLE_FINALISER
+    #if MICROPY_ENABLE_FINALISER || MICROPY_PY_WEAKREF
     area->gc_alloc_table_byte_len = (total_byte_len - ALLOC_TABLE_GAP_BYTE)
         * MP_BITS_PER_BYTE
         / (
             MP_BITS_PER_BYTE
+            #if MICROPY_ENABLE_FINALISER
             + MP_BITS_PER_BYTE * BLOCKS_PER_ATB / BLOCKS_PER_FTB
+            #endif
+            #if MICROPY_PY_WEAKREF
+            + MP_BITS_PER_BYTE * BLOCKS_PER_ATB / BLOCKS_PER_WTB
+            #endif
             + MP_BITS_PER_BYTE * BLOCKS_PER_ATB * BYTES_PER_BLOCK
             );
     #else
@@ -157,26 +170,36 @@ static void gc_setup_area(mp_state_mem_area_t *area, void *start, void *end) {
 
     area->gc_alloc_table_start = (byte *)start;
 
+    // Allocate FTB and WTB blocks if they are enabled.
+    byte *next_table = area->gc_alloc_table_start + area->gc_alloc_table_byte_len + ALLOC_TABLE_GAP_BYTE;
+    (void)next_table;
     #if MICROPY_ENABLE_FINALISER
     size_t gc_finaliser_table_byte_len = (area->gc_alloc_table_byte_len * BLOCKS_PER_ATB + BLOCKS_PER_FTB - 1) / BLOCKS_PER_FTB;
-    area->gc_finaliser_table_start = area->gc_alloc_table_start + area->gc_alloc_table_byte_len + ALLOC_TABLE_GAP_BYTE;
+    area->gc_finaliser_table_start = next_table;
+    next_table += gc_finaliser_table_byte_len;
+    #endif
+    #if MICROPY_PY_WEAKREF
+    size_t gc_weakref_table_byte_len = (area->gc_alloc_table_byte_len * BLOCKS_PER_ATB + BLOCKS_PER_WTB - 1) / BLOCKS_PER_WTB;
+    area->gc_weakref_table_start = next_table;
+    next_table += gc_weakref_table_byte_len;
     #endif
 
+    // Allocate the GC pool of heap blocks.
     size_t gc_pool_block_len = area->gc_alloc_table_byte_len * BLOCKS_PER_ATB;
     area->gc_pool_start = (byte *)end - gc_pool_block_len * BYTES_PER_BLOCK;
     area->gc_pool_end = end;
+    assert(area->gc_pool_start >= next_table);
 
-    #if MICROPY_ENABLE_FINALISER
-    assert(area->gc_pool_start >= area->gc_finaliser_table_start + gc_finaliser_table_byte_len);
-    #endif
-
-    #if MICROPY_ENABLE_FINALISER
-    // clear ATB's and FTB's
-    memset(area->gc_alloc_table_start, 0, gc_finaliser_table_byte_len + area->gc_alloc_table_byte_len + ALLOC_TABLE_GAP_BYTE);
-    #else
-    // clear ATB's
-    memset(area->gc_alloc_table_start, 0, area->gc_alloc_table_byte_len + ALLOC_TABLE_GAP_BYTE);
-    #endif
+    // Clear ATB's, and FTB's and WTB's if they are enabled.
+    memset(area->gc_alloc_table_start, 0,
+        area->gc_alloc_table_byte_len + ALLOC_TABLE_GAP_BYTE
+        #if MICROPY_ENABLE_FINALISER
+        + gc_finaliser_table_byte_len
+        #endif
+        #if MICROPY_PY_WEAKREF
+        + gc_weakref_table_byte_len
+        #endif
+        );
 
     area->gc_last_free_atb_index = 0;
     area->gc_last_used_block = 0;
@@ -196,6 +219,12 @@ static void gc_setup_area(mp_state_mem_area_t *area, void *start, void *end) {
         gc_finaliser_table_byte_len,
         gc_finaliser_table_byte_len * BLOCKS_PER_FTB);
     #endif
+    #if MICROPY_PY_WEAKREF
+    DEBUG_printf("  weakref table at %p, length " UINT_FMT " bytes, "
+        UINT_FMT " blocks\n", area->gc_weakref_table_start,
+        gc_weakref_table_byte_len,
+        gc_weakref_table_byte_len * BLOCKS_PER_WTB);
+    #endif
     DEBUG_printf("  pool at %p, length " UINT_FMT " bytes, "
         UINT_FMT " blocks\n", area->gc_pool_start,
         gc_pool_block_len * BYTES_PER_BLOCK, gc_pool_block_len);
@@ -226,6 +255,10 @@ void gc_init(void *start, void *end) {
     #endif
 
     GC_MUTEX_INIT();
+
+    #if MICROPY_PY_WEAKREF
+    mp_map_init(&MP_STATE_MEM(mp_weakref_map), 0);
+    #endif
 }
 
 #if MICROPY_GC_SPLIT_HEAP
@@ -310,6 +343,9 @@ static bool gc_try_add_heap(size_t failed_alloc) {
         #if MICROPY_ENABLE_FINALISER
         + total_blocks / BLOCKS_PER_FTB
         #endif
+        #if MICROPY_PY_WEAKREF
+        + total_blocks / BLOCKS_PER_WTB
+        #endif
         + total_blocks * BYTES_PER_BLOCK
         + ALLOC_TABLE_GAP_BYTE
         + sizeof(mp_state_mem_area_t);
@@ -547,6 +583,10 @@ void gc_sweep_all(void) {
 void gc_collect_end(void) {
     gc_deal_with_stack_overflow();
     gc_sweep_run_finalisers();
+    #if MICROPY_PY_WEAKREF
+    gc_collect_root((void **)&MP_STATE_MEM(mp_weakref_map).table, 1);
+    gc_deal_with_stack_overflow();
+    #endif
     gc_sweep_free_blocks();
     #if MICROPY_GC_SPLIT_HEAP
     MP_STATE_MEM(gc_last_free_area) = &MP_STATE_MEM(area);
@@ -556,6 +596,9 @@ void gc_collect_end(void) {
     }
     MP_STATE_THREAD(gc_lock_depth) &= ~GC_COLLECT_FLAG;
     GC_EXIT();
+    #if MICROPY_PY_WEAKREF
+    gc_weakref_sweep();
+    #endif
 }
 
 static void gc_deal_with_stack_overflow(void) {
@@ -581,12 +624,16 @@ static void gc_deal_with_stack_overflow(void) {
 
 // Run finalisers for all to-be-freed blocks
 static void gc_sweep_run_finalisers(void) {
-    #if MICROPY_ENABLE_FINALISER
+    #if MICROPY_ENABLE_FINALISER || MICROPY_PY_WEAKREF
+    #if MICROPY_ENABLE_FINALISER && MICROPY_PY_WEAKREF
+    MP_STATIC_ASSERT(BLOCKS_PER_FTB == BLOCKS_PER_WTB);
+    #endif
     for (const mp_state_mem_area_t *area = &MP_STATE_MEM(area); area != NULL; area = NEXT_AREA(area)) {
         assert(area->gc_last_used_block <= area->gc_alloc_table_byte_len * BLOCKS_PER_ATB);
         // Small speed optimisation: skip over empty FTB blocks
         size_t ftb_end = area->gc_last_used_block / BLOCKS_PER_FTB; // index is inclusive
         for (size_t ftb_idx = 0; ftb_idx <= ftb_end; ftb_idx++) {
+            #if MICROPY_ENABLE_FINALISER
             byte ftb = area->gc_finaliser_table_start[ftb_idx];
             size_t block = ftb_idx * BLOCKS_PER_FTB;
             while (ftb) {
@@ -616,9 +663,26 @@ static void gc_sweep_run_finalisers(void) {
                 ftb >>= 1;
                 block++;
             }
+            #endif
+            #if MICROPY_PY_WEAKREF
+            byte wtb = area->gc_weakref_table_start[ftb_idx];
+            block = ftb_idx * BLOCKS_PER_WTB;
+            while (wtb) {
+                MICROPY_GC_HOOK_LOOP(block);
+                if (wtb & 1) { // WTB_GET(area, block) shortcut
+                    if (ATB_GET_KIND(area, block) == AT_HEAD) {
+                        mp_obj_base_t *obj = (mp_obj_base_t *)PTR_FROM_BLOCK(area, block);
+                        gc_weakref_about_to_be_freed(obj);
+                        WTB_CLEAR(area, block);
+                    }
+                }
+                wtb >>= 1;
+                block++;
+            }
+            #endif
         }
     }
-    #endif // MICROPY_ENABLE_FINALISER
+    #endif // MICROPY_ENABLE_FINALISER || MICROPY_PY_WEAKREF
 }
 
 // Free unmarked heads and their tails
@@ -769,6 +833,25 @@ void gc_info(gc_info_t *info) {
     GC_EXIT();
 }
 
+#if MICROPY_PY_WEAKREF
+// Mark the GC heap pointer as having a weakref.
+void gc_weakref_mark(void *ptr) {
+    mp_state_mem_area_t *area;
+    #if MICROPY_GC_SPLIT_HEAP
+    area = gc_get_ptr_area(ptr);
+    assert(area);
+    #else
+    assert(VERIFY_PTR(ptr));
+    area = &MP_STATE_MEM(area);
+    #endif
+
+    size_t block = BLOCK_FROM_PTR(area, ptr);
+    assert(ATB_GET_KIND(area, block) == AT_HEAD);
+
+    WTB_SET(area, block);
+}
+#endif
+
 void *gc_alloc(size_t n_bytes, unsigned int alloc_flags) {
     bool has_finaliser = alloc_flags & GC_ALLOC_FLAG_HAS_FINALISER;
     size_t n_blocks = ((n_bytes + BYTES_PER_BLOCK - 1) & (~(BYTES_PER_BLOCK - 1))) / BYTES_PER_BLOCK;
@@ -967,6 +1050,11 @@ void gc_free(void *ptr) {
     FTB_CLEAR(area, block);
     #endif
 
+    #if MICROPY_PY_WEAKREF
+    // Objects that have a weak reference should not be explicitly freed.
+    assert(!WTB_GET(area, block));
+    #endif
+
     #if MICROPY_GC_SPLIT_HEAP
     if (MP_STATE_MEM(gc_last_free_area) != area) {
         // We freed something but it isn't the current area. Reset the

py/gc.h

@@ -58,6 +58,11 @@ void gc_collect_end(void);
 // Use this function to sweep the whole heap and run all finalisers
 void gc_sweep_all(void);
 
+// These functions are used to manage weakrefs.
+void gc_weakref_mark(void *ptr);
+void gc_weakref_about_to_be_freed(void *ptr);
+void gc_weakref_sweep(void);
+
 enum {
     GC_ALLOC_FLAG_HAS_FINALISER = 1,
 };