future: fix completion of a future that ends with a delayed tail call

The completion needs to be set up as an lightweight contination so
that it can be captured.

Merge to v6.1
(cherry picked from commit 7a5746d9a7)

pkgs/racket-pkgs/racket-test/tests/future/tail-end.rkt

@@ -0,0 +1,19 @@
+#lang racket/base
+(require racket/future)
+
+;; This example ends up with a delayed tail call
+
+(let ()
+  (struct a (x))
+  (define an-a (a 10))
+
+  (define ((f arg) x)
+    (x arg))
+  (set! f f)
+
+  (void ((f an-a) a-x))
+
+  (define f1 (future (lambda ()
+                       ((f f) (f f)))))
+
+  (touch f1))

racket/src/racket/src/future.c

@@ -2373,9 +2373,8 @@ void *worker_thread_future_loop(void *arg)
           scheme_fill_lwc_start();
           jitcode = ((Scheme_Native_Closure *)rator)->code->start_code;
           v = scheme_call_as_lightweight_continuation(jitcode, rator, argc, argv);
-          if (SAME_OBJ(v, SCHEME_TAIL_CALL_WAITING)) {
-            v = scheme_ts_scheme_force_value_same_mark(v);
-          }
+          if (SAME_OBJ(v, SCHEME_TAIL_CALL_WAITING))
+            v = scheme_force_value_same_mark_as_lightweight_continuation(v);
         }
       }
 
@@ -2468,7 +2467,10 @@ static Scheme_Object *_apply_future_lw(future_t *ft)
                                             FUTURE_RUNSTACK_SIZE);
   
   if (SAME_OBJ(v, SCHEME_TAIL_CALL_WAITING)) {
-    v = scheme_ts_scheme_force_value_same_mark(v);
+    if (scheme_future_thread_state->is_runtime_thread)
+      v = scheme_force_value_same_mark(v);
+    else
+      v = scheme_force_value_same_mark_as_lightweight_continuation(v);
   }
 
   return v;

racket/src/racket/src/future.h

@@ -249,7 +249,7 @@ typedef struct fsemaphore_t {
 
 # include "jit_ts_protos.h"
 
-extern Scheme_Object *scheme_ts_scheme_force_value_same_mark(Scheme_Object *v);
+extern Scheme_Object *scheme_force_value_same_mark_as_lightweight_continuation(Scheme_Object *v);
 
 extern Scheme_Object **scheme_rtcall_on_demand(Scheme_Object **argv);
 extern uintptr_t scheme_rtcall_alloc(void);

racket/src/racket/src/jit.c

@@ -129,10 +129,11 @@ static Scheme_Object *clear_rs_arguments(Scheme_Object *v, int size, int delta)
 
 THREAD_LOCAL_DECL(Scheme_Current_LWC *scheme_current_lwc);
 
-Scheme_Object *scheme_call_as_lightweight_continuation(Scheme_Native_Proc *code,
-                                                       void *data,
-                                                       int argc, 
-                                                       Scheme_Object **argv)
+static Scheme_Object *do_call_as_lwc(Scheme_Native_Proc *code,
+                                     void *data,
+                                     int argc,
+                                     Scheme_Object **argv,
+                                     MZ_MARK_STACK_TYPE cont_mark_stack_start)
 {
 #ifdef JIT_THREAD_LOCAL
 # define THDLOC &BOTTOM_VARIABLE
@@ -140,11 +141,29 @@ Scheme_Object *scheme_call_as_lightweight_continuation(Scheme_Native_Proc *code,
 # define THDLOC NULL
 #endif
   scheme_current_lwc->runstack_start = MZ_RUNSTACK;
-  scheme_current_lwc->cont_mark_stack_start = MZ_CONT_MARK_STACK;
+  scheme_current_lwc->cont_mark_stack_start = cont_mark_stack_start;
   return sjc.native_starter_code(data, argc, argv, THDLOC, code, (void **)&scheme_current_lwc->stack_start);
 #undef THDLOC
 }
 
+Scheme_Object *scheme_call_as_lightweight_continuation(Scheme_Native_Proc *code,
+                                                       void *data,
+                                                       int argc,
+                                                       Scheme_Object **argv)
+{
+  return do_call_as_lwc(code, data, argc, argv, MZ_CONT_MARK_STACK);
+}
+
+#ifdef MZ_USE_FUTURES
+Scheme_Object *scheme_force_value_same_mark_as_lightweight_continuation(Scheme_Object *v)
+{
+  /* Providing 0 as cont_mark_stack_start is the "same_mark" part:
+     it preserves any continuation marks that are in place as part
+     of the continuation. */
+  return do_call_as_lwc(sjc.force_value_same_mark_code, NULL, 0, NULL, 0);
+}
+#endif
+
 void scheme_fill_stack_lwc_end(void) XFORM_SKIP_PROC
 {
 #ifdef JIT_THREAD_LOCAL
@@ -3152,7 +3171,7 @@ int scheme_generate(Scheme_Object *obj, mz_jit_state *jitter, int is_tail, int w
 /*                          procedure codegen                             */
 /*========================================================================*/
 
-static void generate_function_prolog(mz_jit_state *jitter, void *code, int max_let_depth)
+void scheme_generate_function_prolog(mz_jit_state *jitter)
 {
   int in;
   START_JIT_DATA();
@@ -3289,10 +3308,7 @@ static int do_generate_closure(mz_jit_state *jitter, void *_data)
   argv = gdata->argv;
   argv_delta = gdata->argv_delta;
 
-  generate_function_prolog(jitter, start_code, 
-			   /* max_extra_pushed may be wrong the first time around,
-			      but it will be right the last time around */
-			   WORDS_TO_BYTES(data->max_let_depth + jitter->max_extra_pushed));
+  scheme_generate_function_prolog(jitter);
   CHECK_LIMIT();
 
   cnt = generate_function_getarg(jitter, 
@@ -4004,7 +4020,7 @@ static int do_generate_case_lambda_dispatch(mz_jit_state *jitter, void *_data)
 
   start_code = jit_get_ip();
   
-  generate_function_prolog(jitter, start_code, data->ndata->max_let_depth);
+  scheme_generate_function_prolog(jitter);
   CHECK_LIMIT();
   
   if (generate_case_lambda_dispatch(jitter, data->c, data->ndata, 1)) {

racket/src/racket/src/jit.h

@@ -368,6 +368,7 @@ struct scheme_jit_common_record {
 #endif
   void *make_rest_list_code, *make_rest_list_clear_code;
   void *call_check_not_defined_code, *call_check_assign_not_defined_code;
+  void *force_value_same_mark_code;
 
   Continuation_Apply_Indirect continuation_apply_indirect_code;
 #ifdef MZ_USE_LWC
@@ -1481,6 +1482,7 @@ void scheme_jit_register_helper_func(mz_jit_state *jitter, void *code, int gcabl
 Scheme_Object *scheme_noncm_prim_indirect(Scheme_Prim proc, int argc);
 Scheme_Object *scheme_prim_indirect(Scheme_Primitive_Closure_Proc proc, int argc, Scheme_Object *self);
 #endif
+int scheme_generate_force_value_same_mark(mz_jit_state *jitter);
 
 /**********************************************************************/
 /*                             jitstack                               */
@@ -1522,6 +1524,8 @@ int scheme_generate(Scheme_Object *obj, mz_jit_state *jitter, int tail_ok, int w
                     Branch_Info *for_branch);
 int scheme_generate_unboxed(Scheme_Object *obj, mz_jit_state *jitter, int inlined_ok, int unbox_anyway);
 
+void scheme_generate_function_prolog(mz_jit_state *jitter);
+
 #ifdef USE_FLONUM_UNBOXING
 int scheme_generate_flonum_local_unboxing(mz_jit_state *jitter, int push, int no_store, int extfl);
 int scheme_generate_flonum_local_boxing(mz_jit_state *jitter, int pos, int offset, int target, int extfl);

racket/src/racket/src/jitcall.c

@@ -147,17 +147,12 @@ Scheme_Object *scheme_prim_indirect(Scheme_Primitive_Closure_Proc proc, int argc
     return proc(argc, MZ_RUNSTACK, self);
 }
 
+#endif
+
 /* Various specific 'futurized' versions of primitives that may 
    be invoked directly from JIT code and are not considered thread-safe 
    (are not invoked via apply_multi_from_native, etc.) */
 
-Scheme_Object *scheme_ts_scheme_force_value_same_mark(Scheme_Object *v)
-{
-  return ts_scheme_force_value_same_mark(v);
-}
-
-#endif
-
 #ifdef MZ_USE_FUTURES
 static Scheme_Object *ts__scheme_tail_apply_from_native(Scheme_Object *rator, int argc, Scheme_Object **argv)
   XFORM_SKIP_PROC
@@ -547,6 +542,17 @@ int scheme_generate_tail_call(mz_jit_state *jitter, int num_rands, int direct_na
   return 1;
 }
 
+int scheme_generate_force_value_same_mark(mz_jit_state *jitter)
+{
+  GC_CAN_IGNORE jit_insn *refr USED_ONLY_FOR_FUTURES;
+  jit_movi_p(JIT_R0, SCHEME_TAIL_CALL_WAITING);
+  mz_prepare(1);
+  jit_pusharg_p(JIT_R0);
+  (void)mz_finish_lwe(ts_scheme_force_value_same_mark, refr);
+  jit_retval(JIT_R0);
+  return 1;
+}
+
 int scheme_generate_finish_apply(mz_jit_state *jitter)
 {
   GC_CAN_IGNORE jit_insn *refr USED_ONLY_FOR_FUTURES;

racket/src/racket/src/jitcommon.c

@@ -3274,6 +3274,25 @@ static int common12(mz_jit_state *jitter, void *_data)
   return 1;
 }
 
+static int common13(mz_jit_state *jitter, void *_data)
+{
+  /* *** force_value_same_mark_code *** */
+  /* Helper for futures: a synthetic functon that just forces values,
+     which will bounce back to the runtime thread (but with lightweight
+     continuation capture in place). */
+  sjc.force_value_same_mark_code = jit_get_ip();
+  scheme_generate_function_prolog(jitter);
+  CHECK_LIMIT();
+
+  scheme_generate_force_value_same_mark(jitter);
+  CHECK_LIMIT();
+
+  mz_pop_threadlocal();
+  mz_pop_locals();
+  jit_ret();
+  return 1;
+}
+
 int scheme_do_generate_common(mz_jit_state *jitter, void *_data)
 {
   if (!common0(jitter, _data)) return 0;
@@ -3293,6 +3312,7 @@ int scheme_do_generate_common(mz_jit_state *jitter, void *_data)
   if (!common10(jitter, _data)) return 0;
   if (!common11(jitter, _data)) return 0;
   if (!common12(jitter, _data)) return 0;
+  if (!common13(jitter, _data)) return 0;
   return 1;
 }