From 109c1de015eeee385020233e2f8cc6f921149103 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Fri, 31 Oct 2014 21:30:46 +0000
Subject: [PATCH] py: Make gc.enable/disable just control auto-GC; alloc is
 still allowed.

gc.enable/disable are now the same as CPython: they just control whether
automatic garbage collection is enabled or not.  If disabled, you can
still allocate heap memory, and initiate a manual collection.
---
 docs/library/gc.rst            | 19 +++++++++----------
 py/gc.c                        |  9 +++++++--
 py/gc.h                        |  5 +++++
 py/malloc.c                    |  1 +
 py/modgc.c                     | 16 ++++++++++------
 py/objexcept.c                 |  1 +
 py/qstr.c                      |  1 +
 py/qstrdefs.h                  |  1 +
 tests/micropython/heapalloc.py |  7 ++++++-
 unix/gccollect.c               |  1 +
 10 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/docs/library/gc.rst b/docs/library/gc.rst
index 212fac1db8..3e9160f98d 100644
--- a/docs/library/gc.rst
+++ b/docs/library/gc.rst
@@ -4,23 +4,22 @@
 .. module:: gc
    :synopsis: control the garbage collector
 
-
-
 Functions
 ---------
 
+.. function:: enable()
+
+   Enable automatic garbage collection.
+
+.. function:: disable()
+
+   Disable automatic garbage collection.  Heap memory can still be allocated,
+   and garbage collection can still be initiated manually using :meth:`gc.collect`.
+
 .. function:: collect()
 
    Run a garbage collection.
 
-.. function:: disable()
-
-   Disable the garbage collector.
-
-.. function:: enable()
-
-   Enable the garbage collector.
-
 .. function:: mem_alloc()
 
    Return the number of bytes of heap RAM that are allocated.
diff --git a/py/gc.c b/py/gc.c
index ce2fa3a48d..fca5050e79 100644
--- a/py/gc.c
+++ b/py/gc.c
@@ -28,6 +28,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <stdbool.h>
+#include <stdint.h>
 
 #include "mpconfig.h"
 #include "misc.h"
@@ -68,7 +69,8 @@ STATIC mp_uint_t *gc_pool_end;
 STATIC int gc_stack_overflow;
 STATIC mp_uint_t gc_stack[STACK_SIZE];
 STATIC mp_uint_t *gc_sp;
-STATIC mp_uint_t gc_lock_depth;
+STATIC uint16_t gc_lock_depth;
+uint16_t gc_auto_collect_enabled;
 STATIC mp_uint_t gc_last_free_atb_index;
 
 // ATB = allocation table byte
@@ -163,6 +165,9 @@ void gc_init(void *start, void *end) {
     // unlock the GC
     gc_lock_depth = 0;
 
+    // allow auto collection
+    gc_auto_collect_enabled = 1;
+
     DEBUG_printf("GC layout:\n");
     DEBUG_printf("  alloc table at %p, length " UINT_FMT " bytes, " UINT_FMT " blocks\n", gc_alloc_table_start, gc_alloc_table_byte_len, gc_alloc_table_byte_len * BLOCKS_PER_ATB);
 #if MICROPY_ENABLE_FINALISER
@@ -375,7 +380,7 @@ void *gc_alloc(mp_uint_t n_bytes, bool has_finaliser) {
     mp_uint_t end_block;
     mp_uint_t start_block;
     mp_uint_t n_free = 0;
-    int collected = 0;
+    int collected = !gc_auto_collect_enabled;
     for (;;) {
 
         // look for a run of n_blocks available blocks
diff --git a/py/gc.h b/py/gc.h
index dc276dd2b3..69fe83752b 100644
--- a/py/gc.h
+++ b/py/gc.h
@@ -32,6 +32,11 @@ void gc_lock(void);
 void gc_unlock(void);
 bool gc_is_locked(void);
 
+// This variable controls auto garbage collection.  If set to 0 then the
+// GC won't automatically run when gc_alloc can't find enough blocks.  But
+// you can still allocate/free memory and also explicitly call gc_collect.
+extern uint16_t gc_auto_collect_enabled;
+
 // A given port must implement gc_collect by using the other collect functions.
 void gc_collect(void);
 void gc_collect_start(void);
diff --git a/py/malloc.c b/py/malloc.c
index 1ee588fa0a..451e480e8e 100644
--- a/py/malloc.c
+++ b/py/malloc.c
@@ -27,6 +27,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <stdint.h>
 
 #include "mpconfig.h"
 #include "misc.h"
diff --git a/py/modgc.c b/py/modgc.c
index 8c5a7053ae..a27ef7eed9 100644
--- a/py/modgc.c
+++ b/py/modgc.c
@@ -24,15 +24,13 @@
  * THE SOFTWARE.
  */
 
+#include <stdint.h>
+
 #include "mpconfig.h"
 #include "misc.h"
 #include "qstr.h"
 #include "obj.h"
-#include "builtin.h"
 #include "runtime.h"
-#include "objlist.h"
-#include "objtuple.h"
-#include "objstr.h"
 #include "gc.h"
 
 #if MICROPY_PY_GC && MICROPY_ENABLE_GC
@@ -56,7 +54,7 @@ MP_DEFINE_CONST_FUN_OBJ_0(gc_collect_obj, py_gc_collect);
 /// \function disable()
 /// Disable the garbage collector.
 STATIC mp_obj_t gc_disable(void) {
-    gc_lock();
+    gc_auto_collect_enabled = 0;
     return mp_const_none;
 }
 MP_DEFINE_CONST_FUN_OBJ_0(gc_disable_obj, gc_disable);
@@ -64,11 +62,16 @@ MP_DEFINE_CONST_FUN_OBJ_0(gc_disable_obj, gc_disable);
 /// \function enable()
 /// Enable the garbage collector.
 STATIC mp_obj_t gc_enable(void) {
-    gc_unlock();
+    gc_auto_collect_enabled = 1;
     return mp_const_none;
 }
 MP_DEFINE_CONST_FUN_OBJ_0(gc_enable_obj, gc_enable);
 
+STATIC mp_obj_t gc_isenabled(void) {
+    return MP_BOOL(gc_auto_collect_enabled);
+}
+MP_DEFINE_CONST_FUN_OBJ_0(gc_isenabled_obj, gc_isenabled);
+
 /// \function mem_free()
 /// Return the number of bytes of available heap RAM.
 STATIC mp_obj_t gc_mem_free(void) {
@@ -92,6 +95,7 @@ STATIC const mp_map_elem_t mp_module_gc_globals_table[] = {
     { MP_OBJ_NEW_QSTR(MP_QSTR_collect), (mp_obj_t)&gc_collect_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_disable), (mp_obj_t)&gc_disable_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_enable), (mp_obj_t)&gc_enable_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR_isenabled), (mp_obj_t)&gc_isenabled_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_mem_free), (mp_obj_t)&gc_mem_free_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_mem_alloc), (mp_obj_t)&gc_mem_alloc_obj },
 };
diff --git a/py/objexcept.c b/py/objexcept.c
index 9b39788759..4b8d6ea08b 100644
--- a/py/objexcept.c
+++ b/py/objexcept.c
@@ -28,6 +28,7 @@
 #include <stdarg.h>
 #include <assert.h>
 #include <stdio.h>
+#include <stdint.h>
 
 #include "mpconfig.h"
 #include "nlr.h"
diff --git a/py/qstr.c b/py/qstr.c
index c2cfda8a37..e08de2ebab 100644
--- a/py/qstr.c
+++ b/py/qstr.c
@@ -26,6 +26,7 @@
 
 #include <assert.h>
 #include <string.h>
+#include <stdint.h>
 
 #include "mpconfig.h"
 #include "misc.h"
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 74a29d2ee8..ecb1b96236 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -470,6 +470,7 @@ Q(gc)
 Q(collect)
 Q(disable)
 Q(enable)
+Q(isenabled)
 Q(mem_free)
 Q(mem_alloc)
 #endif
diff --git a/tests/micropython/heapalloc.py b/tests/micropython/heapalloc.py
index c62428a084..b4b27d19b7 100644
--- a/tests/micropython/heapalloc.py
+++ b/tests/micropython/heapalloc.py
@@ -20,7 +20,12 @@ def h():
         g(i)            # default arg (second one)
         g(i, i)         # 2 args
 
-# call h with heap allocation disabled
+# call h with heap allocation disabled and all memory used up
 gc.disable()
+try:
+    while True:
+        'a'.lower # allocates 1 cell for boundmeth
+except MemoryError:
+    pass
 h()
 gc.enable()
diff --git a/unix/gccollect.c b/unix/gccollect.c
index 32b3d8bc61..117c13144f 100644
--- a/unix/gccollect.c
+++ b/unix/gccollect.c
@@ -25,6 +25,7 @@
  */
 
 #include <stdio.h>
+#include <stdint.h>
 
 #include "mpconfig.h"
 #include "misc.h"