Fix a bug in continuation capture across multiple stack chunks

2 files changed, 33 insertions, 0 deletions

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