From 8420ffe94afd2683c869123183c2d88ed01c41e5 Mon Sep 17 00:00:00 2001
From: Marcin Szamotulski <coot@coot.me>
Date: Tue, 5 Mar 2024 18:13:24 +0100
Subject: [PATCH] io-sim: use oneShot in IOSim instances

---
 io-sim/src/Control/Monad/IOSim/Types.hs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/io-sim/src/Control/Monad/IOSim/Types.hs b/io-sim/src/Control/Monad/IOSim/Types.hs
index cb8eccca..9fa5412b 100644
--- a/io-sim/src/Control/Monad/IOSim/Types.hs
+++ b/io-sim/src/Control/Monad/IOSim/Types.hs
@@ -255,16 +255,16 @@ instance Applicative (IOSim s) where
 
     {-# INLINE (<*>) #-}
     (<*>) = \df dx -> IOSim $ oneShot $ \k ->
-                        unIOSim df (\f -> unIOSim dx (\x -> k (f x)))
+                        unIOSim df (\f -> unIOSim dx (oneShot $ \x -> k (f x)))
 
     {-# INLINE (*>) #-}
-    (*>) = \dm dn -> IOSim $ oneShot $ \k -> unIOSim dm (\_ -> unIOSim dn k)
+    (*>) = \dm dn -> IOSim $ oneShot $ \k -> unIOSim dm (oneShot $ \_ -> unIOSim dn k)
 
 instance Monad (IOSim s) where
     return = pure
 
     {-# INLINE (>>=) #-}
-    (>>=) = \dm f -> IOSim $ oneShot $ \k -> unIOSim dm (\m -> unIOSim (f m) k)
+    (>>=) = \dm f -> IOSim $ oneShot $ \k -> unIOSim dm (oneShot $ \m -> unIOSim (f m) k)
 
     {-# INLINE (>>) #-}
     (>>) = (*>)