vm: always generate safepoints in jit

Even if there's no stack frame we still need to safepoint before leaving the function. Fixes #332.

basis/bootstrap/image/image.factor

@@ -532,6 +532,8 @@ M: quotation '
     \ set-fpu-state set-fpu-state-word set
     \ signal-handler signal-handler-word set
     \ leaf-signal-handler leaf-signal-handler-word set
+    \ ffi-signal-handler ffi-signal-handler-word set
+    \ ffi-leaf-signal-handler ffi-leaf-signal-handler-word set
     undefined-def undefined-quot set ;
 
 : emit-special-objects ( -- )

basis/cpu/x86/32/bootstrap.factor

@@ -30,6 +30,7 @@ IN: bootstrap.x86
 : link-reg ( -- reg ) EBX ;
 : fixnum>slot@ ( -- ) temp0 2 SAR ;
 : rex-length ( -- n ) 0 ;
+: red-zone-size ( -- n ) 0 ;
 
 : jit-call ( name -- )
     0 CALL f rc-relative rel-dlsym ;
@@ -390,8 +391,9 @@ IN: bootstrap.x86
     EAX EDX [] MOV
     jit-jump-quot ;
 
-: jit-safepoint ( -- )
+[
     0 EAX MOVABS rc-absolute rel-safepoint ;
+] \ jit-safepoint jit-define
 
 [
     jit-start-context-and-delete

basis/cpu/x86/64/bootstrap.factor

@@ -334,8 +334,9 @@ IN: bootstrap.x86
     jit-push-param
     jit-jump-quot ;
 
-: jit-safepoint ( -- )
-    0 [RIP+] EAX MOV rc-relative rel-safepoint ;
+[
+    0 [RIP+] EAX MOV rc-relative rel-safepoint
+] \ jit-safepoint jit-define
 
 [
     jit-start-context-and-delete

basis/cpu/x86/64/unix/bootstrap.factor

@@ -12,6 +12,7 @@ IN: bootstrap.x86
 : arg2 ( -- reg ) RSI ;
 : arg3 ( -- reg ) RDX ;
 : arg4 ( -- reg ) RCX ;
+: red-zone-size ( -- n ) 128 ;
 
 << "vocab:cpu/x86/unix/bootstrap.factor" parse-file suffix! >> call
 << "vocab:cpu/x86/64/bootstrap.factor" parse-file suffix! >> call

basis/cpu/x86/64/windows/bootstrap.factor

@@ -21,6 +21,8 @@ DEFER: stack-reg
 : jit-install-seh ( -- ) stack-reg bootstrap-cell ADD ;
 : jit-update-seh ( ctx-reg -- ) drop ;
 
+: red-zone-size ( -- n ) 0 ;
+
 << "vocab:cpu/x86/windows/bootstrap.factor" parse-file suffix! >> call
 << "vocab:cpu/x86/64/bootstrap.factor" parse-file suffix! >> call
 << "vocab:cpu/x86/bootstrap.factor" parse-file suffix! >> call

basis/cpu/x86/bootstrap.factor

@@ -102,8 +102,8 @@ big-endian off
     0 CALL f rc-relative rel-word-pic
 ] jit-word-call jit-define
 
-! The signal-handler and leaf-signal-handler subprimitives are special-cased
-! in vm/quotations.cpp not to trigger generation of a stack frame, so they can
+! The *-signal-handler subprimitives are special-cased in vm/quotations.cpp
+! not to trigger generation of a stack frame, so they can
 ! peform their own prolog/epilog preserving registers.
 
 [| |
@@ -126,6 +126,22 @@ big-endian off
     leaf-frame-size cell - RET
 ] \ leaf-signal-handler define-sub-primitive
 
+[| |
+    jit-signal-handler-prolog :> frame-size
+    temp0 vm-reg vm-signal-handler-addr-offset [+] MOV
+    temp0 CALL
+    frame-size jit-signal-handler-epilog
+    red-zone-size RET
+] \ ffi-signal-handler define-sub-primitive
+
+[| |
+    jit-signal-handler-prolog :> frame-size
+    temp0 vm-reg vm-signal-handler-addr-offset [+] MOV
+    temp0 CALL
+    frame-size jit-signal-handler-epilog
+    red-zone-size 16 bootstrap-cell - + RET
+] \ ffi-leaf-signal-handler define-sub-primitive
+
 [
     ! load boolean
     temp0 ds-reg [] MOV
@@ -222,7 +238,6 @@ big-endian off
 ] jit-execute jit-define
 
 [
-    jit-safepoint
     stack-reg stack-frame-size bootstrap-cell - ADD
 ] jit-epilog jit-define
 

vm/quotations.cpp

@@ -88,7 +88,9 @@ bool quotation_jit::word_stack_frame_p(cell obj)
 	// See #295.
 	return (to_boolean(untag<word>(obj)->subprimitive)
 			&& obj != parent->special_objects[SIGNAL_HANDLER_WORD]
-			&& obj != parent->special_objects[LEAF_SIGNAL_HANDLER_WORD])
+			&& obj != parent->special_objects[LEAF_SIGNAL_HANDLER_WORD]
+			&& obj != parent->special_objects[FFI_SIGNAL_HANDLER_WORD]
+			&& obj != parent->special_objects[FFI_LEAF_SIGNAL_HANDLER_WORD])
 		|| obj == parent->special_objects[JIT_PRIMITIVE_WORD];
 }
 
@@ -122,6 +124,12 @@ bool quotation_jit::trivial_quotation_p(array *elements)
 	return array_capacity(elements) == 1 && tagged<object>(array_nth(elements,0)).type_p(WORD_TYPE);
 }
 
+void quotation_jit::emit_epilog(bool stack_frame)
+{
+	emit(parent->special_objects[JIT_SAFEPOINT]);
+	if(stack_frame) emit(parent->special_objects[JIT_EPILOG]);
+}
+
 void quotation_jit::emit_quot(cell quot_)
 {
 	data_root<quotation> quot(quot_,parent);
@@ -172,7 +180,7 @@ void quotation_jit::iterate_quotation()
 			/* Everything else */
 			else if(i == length - 1)
 			{
-				if(stack_frame) emit(parent->special_objects[JIT_EPILOG]);
+				emit_epilog(stack_frame);
 				tail_call = true;
 				word_jump(obj.value());
 			}
@@ -210,7 +218,7 @@ void quotation_jit::iterate_quotation()
 			   mutually recursive in the library, but both still work) */
 			if(fast_if_p(i,length))
 			{
-				if(stack_frame) emit(parent->special_objects[JIT_EPILOG]);
+				emit_epilog(stack_frame);
 				tail_call = true;
 
 				emit_quot(array_nth(elements.untagged(),i));
@@ -247,7 +255,7 @@ void quotation_jit::iterate_quotation()
 			/* Method dispatch */
 			if(mega_lookup_p(i,length))
 			{
-				if(stack_frame) emit(parent->special_objects[JIT_EPILOG]);
+				emit_epilog(stack_frame);
 				tail_call = true;
 				emit_mega_cache_lookup(
 					array_nth(elements.untagged(),i),
@@ -271,7 +279,7 @@ void quotation_jit::iterate_quotation()
 	{
 		set_position(length);
 
-		if(stack_frame) emit(parent->special_objects[JIT_EPILOG]);
+		emit_epilog(stack_frame);
 		emit(parent->special_objects[JIT_RETURN]);
 	}
 }

vm/quotations.hpp

@@ -16,6 +16,7 @@ struct quotation_jit : public jit {
 	bool primitive_call_p(cell i, cell length);
 	bool trivial_quotation_p(array *elements);
 	void emit_quot(cell quot);
+	void emit_epilog(bool stack_frame);
 	bool fast_if_p(cell i, cell length);
 	bool fast_dip_p(cell i, cell length);
 	bool fast_2dip_p(cell i, cell length);