3 files changed, 56 insertions, 7 deletions

diff --git a/src/post.c b/src/post.c
index aa29151..d2f40f4 100644
--- a/src/post.c
+++ b/src/post.c
@@ -95,6 +95,9 @@ init_bda(void)
 
     add_e820((u32)MAKE_FLATPTR(ebda_seg, 0), GET_EBDA(ebda_seg, size) * 1024
              , E820_RESERVED);
+
+    // Init extra stack
+    StackPos = (void*)(&ExtraStack[BUILD_EXTRA_STACK_SIZE] - _datalow_base);
 }
 
 static void
diff --git a/src/stacks.c b/src/stacks.c
index 044d9ea..febd8bc 100644
--- a/src/stacks.c
+++ b/src/stacks.c
@@ -15,14 +15,24 @@
 
 // Space for a stack for 16bit code.
 u8 ExtraStack[BUILD_EXTRA_STACK_SIZE+1] VARLOW __aligned(8);
+u8 *StackPos VARLOW;
+
+// Test if currently on the extra stack
+static inline int
+on_extra_stack(void)
+{
+    return MODE16 && GET_SEG(SS) == SEG_LOW && getesp() > (u32)ExtraStack;
+}
 
 // Switch to the extra stack and call a function.
 inline u32
 stack_hop(u32 eax, u32 edx, void *func)
 {
+    if (on_extra_stack())
+        return ((u32 (*)(u32, u32))func)(eax, edx);
     ASSERT16();
-    u16 stack_seg = SEG_LOW, bkup_ss;
-    u32 bkup_esp;
+    u16 stack_seg = SEG_LOW;
+    u32 bkup_ss, bkup_esp;
     asm volatile(
         // Backup current %ss/%esp values.
         "movw %%ss, %w3\n"
@@ -31,14 +41,51 @@ stack_hop(u32 eax, u32 edx, void *func)
         "movw %w6, %%ds\n"
         "movw %w6, %%ss\n"
         "movl %5, %%esp\n"
+        "pushl %3\n"
+        "pushl %4\n"
         // Call func
         "calll *%2\n"
+        "popl %4\n"
+        "popl %3\n"
         // Restore segments and stack
         "movw %w3, %%ds\n"
         "movw %w3, %%ss\n"
         "movl %4, %%esp"
         : "+a" (eax), "+d" (edx), "+c" (func), "=&r" (bkup_ss), "=&r" (bkup_esp)
-        : "i" (&ExtraStack[BUILD_EXTRA_STACK_SIZE]), "r" (stack_seg)
+        : "m" (StackPos), "r" (stack_seg)
+        : "cc", "memory");
+    return eax;
+}
+
+// Switch back to original caller's stack and call a function.
+static u32
+stack_hop_back(u32 eax, u32 edx, void *func)
+{
+    if (!on_extra_stack())
+        return ((u32 (*)(u32, u32))func)(eax, edx);
+    ASSERT16();
+    u16 bkup_ss;
+    u32 bkup_stack_pos, temp;
+    asm volatile(
+        // Backup stack_pos and current %ss/%esp
+        "movl %6, %4\n"
+        "movw %%ss, %w3\n"
+        "movl %%esp, %6\n"
+        // Restore original callers' %ss/%esp
+        "movl -4(%4), %5\n"
+        "movl %5, %%ss\n"
+        "movl %%ds:-8(%4), %%esp\n"
+        "movl %5, %%ds\n"
+        // Call func
+        "calll *%2\n"
+        // Restore %ss/%esp and stack_pos
+        "movw %w3, %%ds\n"
+        "movw %w3, %%ss\n"
+        "movl %6, %%esp\n"
+        "movl %4, %6"
+        : "+a" (eax), "+d" (edx), "+c" (func), "=&r" (bkup_ss)
+          , "=&r" (bkup_stack_pos), "=&r" (temp), "+m" (StackPos)
+        :
         : "cc", "memory");
     return eax;
 }
@@ -245,8 +292,7 @@ void
 yield(void)
 {
     if (MODESEGMENT) {
-        // Just directly check irqs.
-        check_irqs();
+        stack_hop_back(0, 0, check_irqs);
         return;
     }
     extern void _cfunc16_check_irqs(void);
@@ -274,7 +320,7 @@ void
 yield_toirq(void)
 {
     if (MODESEGMENT) {
-        wait_irq();
+        stack_hop_back(0, 0, wait_irq);
         return;
     }
     if (CONFIG_THREADS && MainThread.next != &MainThread) {
diff --git a/src/util.h b/src/util.h
index f4a5ac8..a4fabd5 100644
--- a/src/util.h
+++ b/src/util.h
@@ -221,7 +221,7 @@ void nullTrailingSpace(char *buf);
 int get_keystroke(int msec);
 
 // stacks.c
-extern u8 ExtraStack[];
+extern u8 ExtraStack[], *StackPos;
 inline u32 stack_hop(u32 eax, u32 edx, void *func);
 u32 call32(void *func, u32 eax, u32 errret);
 struct bregs;