[ion][python] Implement an architecture-dependant collect_registers.

setjmp is not guaranteed to collect all registers without modification
on all platforms.

This fixes the following bug: when the pointer of a newly allocated
object on the Python heap is stored in rpb registers on x86_64 arch, it
was not collected by the garbarge collector.

ion/include/ion.h

@@ -17,6 +17,7 @@
 #include <ion/unicode/utf8_helper.h>
 #include <stdint.h>
 #include <string.h>
+#include <setjmp.h>
 
 /* ION is not your regular library. It is a library you link against, but it
  * will take care of configuring the whole environment for you. In POSIX terms,
@@ -48,6 +49,9 @@ void decompress(const uint8_t * src, uint8_t * dst, int srcSize, int dstSize);
 // Tells whether the stack pointer is within acceptable bounds
 bool stackSafe();
 
+// Collect registers in a buffer and returns the stack pointer
+uintptr_t collectRegisters(jmp_buf regs);
+
 }
 
 #endif

ion/src/device/Makefile

@@ -10,6 +10,8 @@ ifeq ($(EPSILON_TELEMETRY),1)
 ion_src += ion/src/shared/telemetry_console.cpp
 endif
 
+ion_src += ion/src/shared/collect_registers.cpp
+
 ION_DEVICE_SFLAGS = -Iion/src/device/$(MODEL) -Iion/src/device/shared
 
 $(call object_for,$(ion_device_src) $(ion_device_flasher_src) $(ion_device_bench_src)): SFLAGS += $(ION_DEVICE_SFLAGS)

ion/src/shared/collect_registers.cpp

@@ -0,0 +1,15 @@
+#include <ion.h>
+#include <setjmp.h>
+
+namespace Ion {
+
+/* Forbid inlining to ensure dummy to be at the top of the stack. Otherwise,
+ * LTO inlining can make regs lower on the stack than some just-allocated
+ * pointers. */
+__attribute__((noinline))uintptr_t collectRegisters(jmp_buf buf) {
+  setjmp(buf);
+  int dummy;
+  return (uintptr_t)&dummy;
+}
+
+}

ion/src/simulator/android/Makefile

@@ -7,6 +7,8 @@ ion_src += $(addprefix ion/src/simulator/shared/, \
   dummy/language.cpp \
 )
 
+ion_src += ion/src/shared/collect_registers.cpp
+
 ifeq ($(EPSILON_TELEMETRY),1)
 ion_src += ion/src/simulator/android/src/cpp/telemetry.cpp
 endif

ion/src/simulator/ios/Makefile

@@ -7,6 +7,8 @@ ion_src += $(addprefix ion/src/simulator/shared/, \
   dummy/callback.cpp \
 )
 
+ion_src += ion/src/shared/collect_registers.cpp
+
 $(call object_for,ion/src/simulator/shared/main.cpp) : SFLAGS += -DEPSILON_SDL_FULLSCREEN=1
 
 ifeq ($(EPSILON_TELEMETRY),1)

ion/src/simulator/linux/Makefile

@@ -15,6 +15,8 @@ ion_src += $(addprefix ion/src/simulator/linux/, \
 
 ion_src += $(addprefix ion/src/simulator/shared/, \
   dummy/callback.cpp \
+  collect_registers_x86_64.s \
+  collect_registers.cpp \
 )
 
 ifeq ($(EPSILON_TELEMETRY),1)

ion/src/simulator/macos/Makefile

@@ -5,6 +5,8 @@ ion_src += $(addprefix ion/src/simulator/macos/, \
 ion_src += $(addprefix ion/src/simulator/shared/, \
   apple/language.m \
   dummy/callback.cpp \
+  collect_registers_x86_64.s \
+  collect_registers.cpp \
 )
 
 ifeq ($(EPSILON_TELEMETRY),1)

ion/src/simulator/shared/collect_registers.cpp

@@ -0,0 +1,17 @@
+#include <ion.h>
+
+extern "C" {
+
+// define in assembly code
+// Force the name as archs (linux/macos) don't mangle C names the same way
+extern uintptr_t collect_registers(uintptr_t * regs) asm ("_collect_registers");
+}
+namespace Ion {
+
+// Wrapper to avoid
+uintptr_t collectRegisters(jmp_buf buf) {
+  uintptr_t * regs = (uintptr_t *)buf;
+  return collect_registers(regs);
+}
+
+}

ion/src/simulator/shared/collect_registers_x86_64.s

@@ -0,0 +1,25 @@
+.text
+
+.global _collect_registers
+
+_collect_registers:
+        pushq   %r15
+        pushq   %r14
+        pushq   %r13
+        pushq   %r12
+        pushq   %rbp
+        pushq   %rbx
+        movq    %rbx, (%rdi)
+        movq    %rbp, 8(%rdi)
+        movq    %r12, 16(%rdi)
+        popq    %rbx
+        popq    %rbp
+        popq    %r12
+        movq    %r13, 24(%rdi)
+        movq    %r14, 32(%rdi)
+        popq    %r13
+        movq    %r15, 40(%rdi)
+        popq    %r14
+        popq    %r15
+        movq    %rsp, %rax
+        ret

ion/src/simulator/web/Makefile

@@ -22,6 +22,8 @@ ion_src += $(addprefix ion/src/simulator/shared/, \
   dummy/language.cpp \
 )
 
+ion_src += ion/src/shared/collect_registers.cpp
+
 ifeq ($(EPSILON_TELEMETRY),1)
 ion_src += ion/src/simulator/shared/dummy/telemetry_init.cpp
 ion_src += ion/src/shared/telemetry_console.cpp

ion/src/simulator/windows/Makefile

@@ -8,6 +8,8 @@ ion_src += $(addprefix ion/src/simulator/shared/, \
   dummy/callback.cpp \
 )
 
+ion_src += ion/src/shared/collect_registers.cpp
+
 ifeq ($(EPSILON_TELEMETRY),1)
 ion_src += ion/src/simulator/shared/dummy/telemetry_init.cpp
 ion_src += ion/src/shared/telemetry_console.cpp

liba/include/setjmp.h

@@ -1,6 +1,7 @@
 #ifndef LIBA_SETJMP_H
 #define LIBA_SETJMP_H
 
+#include <stdint.h>
 #include "private/macros.h"
 
 /* We are preseving registers:
@@ -14,7 +15,7 @@
 
 LIBA_BEGIN_DECLS
 
-typedef int jmp_buf[31];
+typedef uintptr_t jmp_buf[31];
 void longjmp(jmp_buf env, int val);
 int setjmp(jmp_buf env);
 

python/port/port.cpp

@@ -1,6 +1,6 @@
 #include "port.h"
 
-#include <ion/keyboard.h>
+#include <ion.h>
 
 #include <math.h>
 #include <stdint.h>
@@ -231,63 +231,54 @@ KDColor MicroPython::ColorParser::ParseColor(mp_obj_t input, ColorMode ColorMode
   mp_raise_TypeError("Color couldn't be parsed");
 }
 
-/* Forbid inlining to ensure regs to be at the top of the stack. Otherwise,
- * LTO inlining can make regs lower on the stack than some just-allocated
- * pointers. */
-__attribute__((noinline)) void gc_collect(void) {
+void gc_collect_regs_and_stack(void) {
+  // get the registers and the sp
+  jmp_buf regs;
+  uintptr_t sp = Ion::collectRegisters(regs);
+
   void * python_stack_top = MP_STATE_THREAD(stack_top);
   assert(python_stack_top != NULL);
 
-  gc_collect_start();
-
-  modturtle_gc_collect();
-  modpyplot_gc_collect();
-
-  /* get the registers.
-   * regs is the also the last object on the stack so the stack is bound by
-   * &regs and python_stack_top. */
-  jmp_buf regs;
-  /* TODO: we use setjmp to get the registers values to look for python heap
-   * root. However, the 'setjmp' does not guarantee that it gets all registers
-   * values. We should check our setjmp implementation for the device and
-   * ensure that it also works for other platforms. */
-  setjmp(regs);
-
-  void **regs_ptr = (void**)&regs;
-
   /* On the device, the stack is stored in reverse order, but it might not be
    * the case on a computer. We thus have to take the absolute value of the
    * addresses difference. */
   size_t stackLengthInByte;
   void ** scanStart;
-  if ((uintptr_t)python_stack_top > (uintptr_t)regs_ptr) {
+  if ((uintptr_t)python_stack_top > sp) {
 
     /* To compute the stack length:
-     *                                  regs
+     *                               registers
      *                             <----------->
      * STACK <-  ...|  |  |  |  |  |--|--|--|--|  |  |  |  |  |  |
-     *                             ^&regs                        ^python_stack_top
+     *                             ^sp                           ^python_stack_top
      * */
 
-    stackLengthInByte = (uintptr_t)python_stack_top - (uintptr_t)regs_ptr;
-    scanStart = regs_ptr;
+    stackLengthInByte = (uintptr_t)python_stack_top - sp;
+    scanStart = (void **)sp;
 
   } else {
 
     /* When computing the stack length, take into account regs' size.
-     *                                                 regs
+     *                                              registers
      *                                            <----------->
      * STACK ->  |  |  |  |  |  |  |  |  |  |  |  |--|--|--|--|  |  |  |...
-     *           ^python_stack_top                ^&regs
+     *           ^python_stack_top                ^sp
      * */
 
-    stackLengthInByte = (uintptr_t)regs_ptr - (uintptr_t)python_stack_top + sizeof(regs);
+    stackLengthInByte = sp - (uintptr_t)python_stack_top + sizeof(regs);
     scanStart = (void **)python_stack_top;
 
   }
   /* Memory error detectors might find an error here as they might split regs
    * and stack memory zones. */
   MicroPython::collectRootsAtAddress((char *)scanStart, stackLengthInByte);
+}
+
+void gc_collect(void) {
+  gc_collect_start();
+  modturtle_gc_collect();
+  modpyplot_gc_collect();
+  gc_collect_regs_and_stack();
   gc_collect_end();
 }