Fix a bug in continuation capture across multiple stack chunks

3 files changed, 37 insertions, 2 deletions

diff --git a/rts/Continuation.c b/rts/Continuation.c
index 09be4d368e..fbc279574f 100644
--- a/rts/Continuation.c
+++ b/rts/Continuation.c
@@ -472,12 +472,14 @@ StgClosure *captureContinuationAndAbort(Capability *cap, StgTSO *tso, StgPromptT
     stack = pop_stack_chunk(cap, tso);
 
     for (StgWord i = 0; i < full_chunks; i++) {
-      memcpy(cont_stack, stack->sp, stack->stack_size * sizeof(StgWord));
-      cont_stack += stack->stack_size;
+      const size_t chunk_words = stack->stack + stack->stack_size - stack->sp - sizeofW(StgUnderflowFrame);
+      memcpy(cont_stack, stack->sp, chunk_words * sizeof(StgWord));
+      cont_stack += chunk_words;
       stack = pop_stack_chunk(cap, tso);
     }
 
     memcpy(cont_stack, stack->sp, last_chunk_words * sizeof(StgWord));
+    cont_stack += last_chunk_words;
     stack->sp += last_chunk_words;
   }
 
diff --git a/testsuite/tests/rts/continuations/all.T b/testsuite/tests/rts/continuations/all.T
index fb6b6f2ce1..7b35e29c00 100644
--- a/testsuite/tests/rts/continuations/all.T
+++ b/testsuite/tests/rts/continuations/all.T
@@ -2,3 +2,4 @@ test('cont_simple_shift', [extra_files(['ContIO.hs'])], multimod_compile_and_run
 test('cont_exn_masking', [extra_files(['ContIO.hs'])], multimod_compile_and_run, ['cont_exn_masking', ''])
 test('cont_missing_prompt_err', [extra_files(['ContIO.hs']), exit_code(1)], multimod_compile_and_run, ['cont_missing_prompt_err', ''])
 test('cont_nondet_handler', [extra_files(['ContIO.hs'])], multimod_compile_and_run, ['cont_nondet_handler', ''])
+test('cont_stack_overflow', [extra_files(['ContIO.hs'])], multimod_compile_and_run, ['cont_stack_overflow', '-with-rtsopts "-ki1k -kc2k -kb256"'])
diff --git a/testsuite/tests/rts/continuations/cont_stack_overflow.hs b/testsuite/tests/rts/continuations/cont_stack_overflow.hs
new file mode 100644
index 0000000000..9832310d41
--- /dev/null
+++ b/testsuite/tests/rts/continuations/cont_stack_overflow.hs
@@ -0,0 +1,32 @@
+-- This test is run with RTS options that instruct GHC to use a small stack
+-- chunk size (2k), which ensures this test exercises multi-chunk continuation
+-- captures and restores.
+
+import Control.Monad (unless)
+import ContIO
+
+data Answer
+  = Done Int
+  | Yield (IO Int -> IO Answer)
+
+getAnswer :: Answer -> Int
+getAnswer (Done n)  = n
+getAnswer (Yield _) = error "getAnswer"
+
+main :: IO ()
+main = do
+  tag <- newPromptTag
+  Yield k <- prompt tag $
+    Done <$> buildBigCont tag 6000
+  n <- getAnswer <$> k (getAnswer <$> k (pure 0))
+  unless (n == 36006000) $
+    error $ "produced wrong value: " ++ show n
+
+buildBigCont :: PromptTag Answer
+             -> Int
+             -> IO Int
+buildBigCont tag size
+  | size <= 0 = control0 tag (\k -> pure (Yield k))
+  | otherwise = do
+      n <- buildBigCont tag (size - 1)
+      pure $! n + size