py/modweakref: Implement weakref module with ref and finalize classes.

This adds support for the standard `weakref` module, to make weak
references to Python objects and have callbacks for when an object is
reclaimed by the GC.

This feature was requested by PyScript, to allow control over the lifetime
of external proxy objects (distinct from JS<->Python proxies).

Addresses issue micropython#646 (that's nearly a 12 year old issue!).

Functionality added here:
- `weakref.ref(object [, callback])` create a simple weak reference with
  optional callback to be called when the object is reclaimed by the GC
- `weakref.finalize(object, callback, /, *args, **kwargs)` create a
  finalize object that holds a weak reference to an object and allows more
  convenient callback usage and state change

The new module is enabled at the "everything" level.

The implementation aims to be as efficient as possible, by adding another
bit-per-block to the garbage collector, the WTB (weak table).  Similar to
the finalizer bit (FTB), if a GC block has its corresponding WTB bit set
then a weak reference to that block is held.  The details of that weak
reference are stored in a global map, `mp_weakref_map`, which maps weak
reference to ref/finalize objects, allowing the callbacks to be efficiently
found when the object is reclaimed.

With this feature enabled the overhead is:
- 1/128th of the available memory is used for the new WTB table (eg a 128k
  heap now needs an extra 1k for the WTB).
- Code size is increased.
- At garbage collection time, there is a small overhead to check if the
  collected objects had weak references.  This check is the same as the
  existing FTB finaliser scan, so shouldn't add much overhead.  If there
  are weak reference objects alive (ref/finalize objects) then additional
  time is taken to call the callbacks and do some accounting to clean up
  the used weak reference.

Signed-off-by: Damien George <damien@micropython.org>

Author: dpgeorge

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_VM(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);
@@ -310,6 +339,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);
@@ -556,6 +588,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 +616,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 +655,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 +825,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 +1042,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,
 };