Make the extra stack re-entrant and "hop back" to check for irqs.

When on the extra stack and it's necessary to check for irqs, switch
back to the original caller's stack to check for irqs.  Make the extra
stack re-entrant, so that a new user of the extra stack wont collide
with an existing user.

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

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;