vm: have to be extra careful when messing with return addresses

basis/cpu/x86/32/bootstrap.factor

@@ -63,6 +63,9 @@ IN: bootstrap.x86
     ds-reg ctx-reg context-datastack-offset [+] MOV
     rs-reg ctx-reg context-retainstack-offset [+] MOV ;
 
+: jit-scrub-return ( n -- )
+    ESP swap [+] 0 MOV ;
+
 [
     ! ctx-reg is preserved across the call because it is non-volatile
     ! in the C ABI
@@ -130,6 +133,7 @@ IN: bootstrap.x86
 
     ! Unwind stack frames
     ESP EDX MOV
+    0 jit-scrub-return
 
     jit-jump-quot
 ] \ unwind-native-frames define-sub-primitive
@@ -252,9 +256,7 @@ IN: bootstrap.x86
 
 ! Contexts
 : jit-switch-context ( reg -- )
-    ! Dummy return address -- it never gets returned to but it
-    ! must point to inside the current code block
-    ESP -4 [+] HEX: ffffffff MOV rc-absolute-cell rt-this jit-rel
+    -4 jit-scrub-return
 
     ! Save ds, rs registers
     jit-load-vm

basis/cpu/x86/64/bootstrap.factor

@@ -61,6 +61,9 @@ IN: bootstrap.x86
     ds-reg ctx-reg context-datastack-offset [+] MOV
     rs-reg ctx-reg context-retainstack-offset [+] MOV ;
 
+: jit-scrub-return ( n -- )
+    RSP swap [+] 0 MOV ;
+
 [
     ! ctx-reg is preserved across the call because it is non-volatile
     ! in the C ABI
@@ -111,6 +114,7 @@ IN: bootstrap.x86
 
     ! Unwind stack frames
     RSP arg2 MOV
+    0 jit-scrub-return
 
     ! Load VM pointer into vm-reg, since we're entering from
     ! C code
@@ -228,10 +232,7 @@ IN: bootstrap.x86
 
 ! Contexts
 : jit-switch-context ( reg -- )
-    ! Dummy return address -- it never gets returned to but it
-    ! must point to inside the current code block
-    R11 0 [RIP+] LEA
-    RSP -8 [+] R11 MOV
+    -8 jit-scrub-return
 
     ! Save ds, rs registers
     jit-save-context

vm/callstack.cpp

@@ -108,7 +108,25 @@ stack_frame *factor_vm::frame_successor(stack_frame *frame)
 	return (stack_frame *)((cell)frame - frame->size);
 }
 
-/* Allocates memory */
+cell factor_vm::frame_offset(stack_frame *frame)
+{
+	char *entry_point = (char *)frame_code(frame)->entry_point();
+	char *return_address = (char *)FRAME_RETURN_ADDRESS(frame,this);
+	if(return_address)
+		return return_address - entry_point;
+	else
+		return (cell)-1;
+}
+
+void factor_vm::set_frame_offset(stack_frame *frame, cell offset)
+{
+	char *entry_point = (char *)frame_code(frame)->entry_point();
+	if(offset == (cell)-1)
+		FRAME_RETURN_ADDRESS(frame,this) = NULL;
+	else
+		FRAME_RETURN_ADDRESS(frame,this) = entry_point + offset;
+}
+
 cell factor_vm::frame_scan(stack_frame *frame)
 {
 	switch(frame_type(frame))
@@ -120,13 +138,7 @@ cell factor_vm::frame_scan(stack_frame *frame)
 				obj = obj.as<word>()->def;
 
 			if(obj.type_p(QUOTATION_TYPE))
-			{
-				char *return_addr = (char *)FRAME_RETURN_ADDRESS(frame,this);
-				char *quot_entry_point = (char *)frame_code(frame)->entry_point();
-
-				return tag_fixnum(quot_code_offset_to_scan(
-					obj.value(),(cell)(return_addr - quot_entry_point)));
-			}    
+				return tag_fixnum(quot_code_offset_to_scan(obj.value(),frame_offset(frame)));
 			else
 				return false_object;
 		}
@@ -138,11 +150,6 @@ cell factor_vm::frame_scan(stack_frame *frame)
 	}
 }
 
-cell factor_vm::frame_offset(stack_frame *frame)
-{
-	return (cell)FRAME_RETURN_ADDRESS(frame,this) - (cell)frame_code(frame)->entry_point();
-}
-
 struct stack_frame_accumulator {
 	factor_vm *parent;
 	growable_array frames;
@@ -209,9 +216,9 @@ void factor_vm::primitive_set_innermost_stack_frame_quot()
 	jit_compile_quot(quot.value(),true);
 
 	stack_frame *inner = innermost_stack_frame(callstack.untagged());
-	cell offset = (char *)FRAME_RETURN_ADDRESS(inner,this) - (char *)inner->entry_point;
+	cell offset = frame_offset(inner);
 	inner->entry_point = quot->entry_point;
-	FRAME_RETURN_ADDRESS(inner,this) = (char *)quot->entry_point + offset;
+	set_frame_offset(inner,offset);
 }
 
 void factor_vm::primitive_callstack_bounds()

vm/code_block_visitor.hpp

@@ -42,13 +42,10 @@ struct call_frame_code_block_visitor {
 
 	void operator()(stack_frame *frame)
 	{
-		code_block *old_block = parent->frame_code(frame);
-		cell offset = (char *)FRAME_RETURN_ADDRESS(frame,parent) - (char *)old_block;
-
-		const code_block *new_block = fixup.fixup_code(old_block);
-		frame->entry_point = new_block->entry_point();
-
-		FRAME_RETURN_ADDRESS(frame,parent) = (char *)new_block + offset;
+		cell offset = parent->frame_offset(frame);
+		code_block *compiled = fixup.fixup_code(parent->frame_code(frame));
+		frame->entry_point = compiled->entry_point();
+		parent->set_frame_offset(frame,offset);
 	}
 };
 

vm/gc.cpp

@@ -207,13 +207,15 @@ struct call_frame_scrubber {
 
 	void operator()(stack_frame *frame)
 	{
-		const code_block *compiled = parent->frame_code(frame);
+		cell return_address = parent->frame_offset(frame);
+		if(return_address == (cell)-1)
+			return;
+
+		code_block *compiled = parent->frame_code(frame);
 		gc_info *info = compiled->block_gc_info();
 
-		cell return_address = parent->frame_offset(frame);
 		assert(return_address < compiled->size());
 		int index = info->return_address_index(return_address);
-
 		if(index != -1)
 			ctx->scrub_stacks(info,index);
 	}

vm/slot_visitor.hpp

@@ -284,23 +284,26 @@ struct call_frame_slot_visitor {
 	*/
 	void operator()(stack_frame *frame)
 	{
-		const code_block *compiled = visitor->fixup.translate_code(parent->frame_code(frame));
-		gc_info *info = compiled->block_gc_info();
 		cell return_address = parent->frame_offset(frame);
+		if(return_address == (cell)-1)
+			return;
+
+		code_block *compiled = visitor->fixup.translate_code(parent->frame_code(frame));
+		gc_info *info = compiled->block_gc_info();
+
 		assert(return_address < compiled->size());
 		int index = info->return_address_index(return_address);
+		if(index == -1)
+			return;
+
+		u8 *bitmap = info->gc_info_bitmap();
+		cell base = info->spill_slot_base(index);
+		cell *stack_pointer = (cell *)(parent->frame_successor(frame) + 1);
 
-		if(index != -1)
+		for(cell spill_slot = 0; spill_slot < info->gc_root_count; spill_slot++)
 		{
-			u8 *bitmap = info->gc_info_bitmap();
-			cell base = info->spill_slot_base(index);
-			cell *stack_pointer = (cell *)(parent->frame_successor(frame) + 1);
-
-			for(cell spill_slot = 0; spill_slot < info->gc_root_count; spill_slot++)
-			{
-				if(bitmap_p(bitmap,base + spill_slot))
-					visitor->visit_handle(&stack_pointer[spill_slot]);
-			}
+			if(bitmap_p(bitmap,base + spill_slot))
+				visitor->visit_handle(&stack_pointer[spill_slot]);
 		}
 	}
 };

vm/vm.hpp

@@ -597,6 +597,7 @@ struct factor_vm
 	stack_frame *frame_successor(stack_frame *frame);
 	cell frame_scan(stack_frame *frame);
 	cell frame_offset(stack_frame *frame);
+	void set_frame_offset(stack_frame *frame, cell offset);
 	void primitive_callstack_to_array();
 	stack_frame *innermost_stack_frame(callstack *stack);
 	void primitive_innermost_stack_frame_executing();