From 06cfe1d8a548175546b73f9af588de17257e9c32 Mon Sep 17 00:00:00 2001
From: Andrzej Rybczak <arybczak@users.noreply.github.com>
Date: Wed, 4 Sep 2024 20:27:00 +0200
Subject: [PATCH] Fix copyRefs in Effect.Dispatch.Dynamic (#240)

Fixes https://github.com/haskell-effectful/effectful/issues/239.
---
 .../src/Effectful/Dispatch/Dynamic.hs         |  2 +-
 effectful/tests/EnvTests.hs                   | 20 ++++++++++++-------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/effectful-core/src/Effectful/Dispatch/Dynamic.hs b/effectful-core/src/Effectful/Dispatch/Dynamic.hs
index 88f3a7b..06b8414 100644
--- a/effectful-core/src/Effectful/Dispatch/Dynamic.hs
+++ b/effectful-core/src/Effectful/Dispatch/Dynamic.hs
@@ -1040,7 +1040,7 @@ copyRefs (Env soffset srefs sstorage) (Env doffset drefs dstorage) = do
     error "storages do not match"
   let size = sizeofPrimArray drefs - doffset
       es = reifyIndices @es @srcEs
-      esSize = length es
+      esSize = 2 * length es
   mrefs <- newPrimArray (esSize + size)
   copyPrimArray mrefs esSize drefs doffset size
   let writeRefs i = \case
diff --git a/effectful/tests/EnvTests.hs b/effectful/tests/EnvTests.hs
index c805279..a1be74a 100644
--- a/effectful/tests/EnvTests.hs
+++ b/effectful/tests/EnvTests.hs
@@ -215,11 +215,11 @@ doubleB = interpose_ $ \case
 
 test_borrowLend :: Assertion
 test_borrowLend = runEff $ do
-  runX 1 2 . evalState @[Int] [3, 4] . runReader () . runReader @[Int] [5, 6] $ do
-    U.assertEqual "expected result" [1, 2, 3, 4, 5, 6] =<< send X
+  runX 1 2 . evalState @[Int] [3] . runReader () . runReader @[Int] [4] $ do
+    U.assertEqual "expected result" [1,2,3,4,1,2,3,4] =<< send X
 
 data X :: Effect where
-  X :: (State [Int] :> es, Reader [Int] :> es) => X (Eff es) [Int]
+  X :: (State [Int] :> es, Reader [Int] :> es, Reader () :> es) => X (Eff es) [Int]
 type instance DispatchOf X = Dynamic
 
 runX :: Int -> Int -> Eff (X : es) a -> Eff es a
@@ -227,12 +227,18 @@ runX s0 r0 = reinterpret (evalState s0 . evalState () . runReader r0) $ \env ->
   X -> localSeqUnlift env $ \unlift -> do
     as <- localSeqLend @[State Int, Reader Int] env $ \withHandlerEffs -> do
       unlift . withHandlerEffs $ do
+        () <- ask
         s <- get @Int
         r <- ask @Int
-        pure [s, r]
-    bs <- localSeqBorrow @[State [Int], Reader [Int]] env $ \withLocalEffs -> do
-      withLocalEffs $ do
         ss <- get @[Int]
         rs <- ask @[Int]
-        pure $ ss ++ rs
+        pure $ [s, r] ++ ss ++ rs
+    bs <- localSeqBorrow @[Reader [Int], State [Int], Reader ()] env $ \withEffs -> do
+      withEffs $ do
+        () <- ask
+        s <- get @Int
+        r <- ask @Int
+        ss <- get @[Int]
+        rs <- ask @[Int]
+        pure $ [s, r] ++ ss ++ rs
     pure $ as ++ bs