diff --git a/defs.h b/defs.h
index 661f761..a472b2e 100644
--- a/defs.h
+++ b/defs.h
@@ -107,7 +107,7 @@ void            wakeup(void*);
 void            yield(void);
 
 // swtch.S
-void            swtch(struct context*, struct context*);
+void            swtch(struct context**, struct context**);
 
 // spinlock.c
 void            acquire(struct spinlock*);
diff --git a/proc.c b/proc.c
index fdd633d..5ecf474 100644
--- a/proc.c
+++ b/proc.c
@@ -137,9 +137,9 @@ copyproc(struct proc *p)
   }
 
   // Set up new context to start executing at forkret (see below).
-  memset(&np->context, 0, sizeof(np->context));
-  np->context.eip = (uint)forkret;
-  np->context.esp = (uint)np->tf;
+  np->context = (struct context *)np->tf - 1;
+  memset(np->context, 0, sizeof(*np->context));
+  np->context->eip = (uint)forkret;
 
   // Clear %eax so that fork system call returns 0 in child.
   np->tf->eax = 0;
@@ -477,7 +477,7 @@ procdump(void)
       state = "???";
     cprintf("%d %s %s", p->pid, state, p->name);
     if(p->state == SLEEPING){
-      getcallerpcs((uint*)p->context.ebp+2, pc);
+      getcallerpcs((uint*)p->context->ebp+2, pc);
       for(j=0; j<10 && pc[j] != 0; j++)
         cprintf(" %p", pc[j]);
     }
diff --git a/proc.h b/proc.h
index 502361d..41860c6 100644
--- a/proc.h
+++ b/proc.h
@@ -7,21 +7,17 @@
 #define NSEGS     6
 
 // Saved registers for kernel context switches.
-// Don't need to save all the %fs etc. segment registers,
+// Don't need to save all the segment registers (%cs, etc),
 // because they are constant across kernel contexts.
-// Save all the regular registers so we don't need to care
-// which are caller save, but not the return register %eax.
-// (Not saving %eax just simplifies the switching code.)
+// Stack pointer is encoded in the address of context,
+// which must be placed at the bottom of the stack.
 // The layout of context must match code in swtch.S.
 struct context {
-  int eip;
-  int esp;
-  int ebx;
-  int ecx;
-  int edx;
-  int esi;
-  int edi;
-  int ebp;
+  uint edi;
+  uint esi;
+  uint ebx;
+  uint ebp;
+  uint eip;
 };
 
 enum proc_state { UNUSED, EMBRYO, SLEEPING, RUNNABLE, RUNNING, ZOMBIE };
@@ -38,8 +34,8 @@ struct proc {
   int killed;               // If non-zero, have been killed
   struct file *ofile[NOFILE];  // Open files
   struct inode *cwd;        // Current directory
-  struct context context;   // Switch here to run process
-  struct trapframe *tf;     // Trap frame for current interrupt
+  struct context *context;  // Switch here to run process
+  struct trapframe *tf;     // Trap frame for current syscall
   char name[16];            // Process name (debugging)
 };
 
@@ -53,7 +49,7 @@ struct proc {
 struct cpu {
   uchar apicid;               // Local APIC ID
   struct proc *curproc;       // Process currently running.
-  struct context context;     // Switch here to enter scheduler
+  struct context *context;    // Switch here to enter scheduler
   struct taskstate ts;        // Used by x86 to find stack for interrupt
   struct segdesc gdt[NSEGS];  // x86 global descriptor table
   volatile uint booted;        // Has the CPU started?
diff --git a/swtch.S b/swtch.S
index 786e9ac..ada98f3 100644
--- a/swtch.S
+++ b/swtch.S
@@ -1,32 +1,26 @@
-#   void swtch(struct context *old, struct context *new);
+#   void swtch(struct context **old, struct context **new);
 #  
 # Save current register context in old
 # and then load register context from new.
 
 .globl swtch
 swtch:
-  # Save old registers
   movl 4(%esp), %eax
+  movl 8(%esp), %edx
 
-  popl 0(%eax)  # %eip
-  movl %esp, 4(%eax)
-  movl %ebx, 8(%eax)
-  movl %ecx, 12(%eax)
-  movl %edx, 16(%eax)
-  movl %esi, 20(%eax)
-  movl %edi, 24(%eax)
-  movl %ebp, 28(%eax)
+  # Save old callee-save registers
+  pushl %ebp
+  pushl %ebx
+  pushl %esi
+  pushl %edi
 
-  # Load new registers
-  movl 4(%esp), %eax  # not 8(%esp) - popped return address above
-
-  movl 28(%eax), %ebp
-  movl 24(%eax), %edi
-  movl 20(%eax), %esi
-  movl 16(%eax), %edx
-  movl 12(%eax), %ecx
-  movl 8(%eax), %ebx
-  movl 4(%eax), %esp
-  pushl 0(%eax)  # %eip
+  # Switch stacks
+  movl %esp, (%eax)
+  movl (%edx), %esp
 
+  # Load new callee-save registers
+  popl %edi
+  popl %esi
+  popl %ebx
+  popl %ebp
   ret