unix: Allow to set heap size using "-X heapsize=N" option.

1 files changed, 55 insertions, 28 deletions

diff --git a/unix/main.c b/unix/main.c
index 9d07b842b1..32f3210c81 100644
--- a/unix/main.c
+++ b/unix/main.c
@@ -24,9 +24,10 @@
 #include <readline/history.h>
 #endif
 
+#if MICROPY_ENABLE_GC
 // Heap size of GC heap (if enabled)
-// TODO: allow to specify on command line
-#define HEAP_SIZE 128*1024
+long heap_size = 128*1024;
+#endif
 
 // Stack top at the start of program
 void *stack_top;
@@ -233,13 +234,34 @@ static mp_obj_t pyb_gc(void) {
 MP_DEFINE_CONST_FUN_OBJ_0(pyb_gc_obj, pyb_gc);
 #endif
 
+// Process options which set interpreter init options
+void pre_process_options(int argc, char **argv) {
+    for (int a = 1; a < argc; a++) {
+        if (argv[a][0] == '-') {
+            if (strcmp(argv[a], "-X") == 0) {
+                if (a + 1 >= argc) {
+                    exit(usage());
+                }
+#if MICROPY_ENABLE_GC
+                if (strncmp(argv[a + 1], "heapsize=", sizeof("heapsize=") - 1) == 0) {
+                    heap_size = strtol(argv[a + 1] + sizeof("heapsize=") - 1, NULL, 0);
+                }
+#endif
+                a++;
+            }
+        }
+    }
+}
+
 int main(int argc, char **argv) {
     volatile int stack_dummy;
     stack_top = (void*)&stack_dummy;
 
+    pre_process_options(argc, argv);
+
 #if MICROPY_ENABLE_GC
-    char *heap = malloc(HEAP_SIZE);
-    gc_init(heap, heap + HEAP_SIZE);
+    char *heap = malloc(heap_size);
+    gc_init(heap, heap + heap_size);
 #endif
 
     qstr_init();
@@ -319,37 +341,42 @@ int main(int argc, char **argv) {
     printf("    peak  %d\n", m_get_peak_bytes_allocated());
     */
 
-    if (argc == 1) {
-        do_repl();
-    } else {
-        for (int a = 1; a < argc; a++) {
-            if (argv[a][0] == '-') {
-                if (strcmp(argv[a], "-c") == 0) {
-                    if (a + 1 >= argc) {
-                        return usage();
-                    }
-                    do_str(argv[a + 1]);
-                    a += 1;
-                } else {
+    bool executed = false;
+    for (int a = 1; a < argc; a++) {
+        if (argv[a][0] == '-') {
+            if (strcmp(argv[a], "-c") == 0) {
+                if (a + 1 >= argc) {
                     return usage();
                 }
+                do_str(argv[a + 1]);
+                executed = true;
+                a += 1;
+            } else if (strcmp(argv[a], "-X") == 0) {
+                a += 1;
             } else {
-                // Set base dir of the script as first entry in sys.path
-                char *basedir = realpath(argv[a], NULL);
-                if (basedir != NULL) {
-                    char *p = strrchr(basedir, '/');
-                    path_items[0] = MP_OBJ_NEW_QSTR(qstr_from_strn(basedir, p - basedir));
-                    free(basedir);
-                }
-                for (int i = a; i < argc; i++) {
-                    rt_list_append(py_argv, MP_OBJ_NEW_QSTR(qstr_from_str(argv[i])));
-                }
-                do_file(argv[a]);
-                break;
+                return usage();
             }
+        } else {
+            // Set base dir of the script as first entry in sys.path
+            char *basedir = realpath(argv[a], NULL);
+            if (basedir != NULL) {
+                char *p = strrchr(basedir, '/');
+                path_items[0] = MP_OBJ_NEW_QSTR(qstr_from_strn(basedir, p - basedir));
+                free(basedir);
+            }
+            for (int i = a; i < argc; i++) {
+                rt_list_append(py_argv, MP_OBJ_NEW_QSTR(qstr_from_str(argv[i])));
+            }
+            do_file(argv[a]);
+            executed = true;
+            break;
         }
     }
 
+    if (!executed) {
+        do_repl();
+    }
+
     rt_deinit();
 
     //printf("total bytes = %d\n", m_get_total_bytes_allocated());