Add catch semantics to STM

Add catch handler as branch to support catch

Author: Yogesh Sajanikar

io-sim/src/Control/Monad/IOSim/Internal.hs

@@ -896,10 +896,33 @@ execAtomically !time !tid !tlbl !nextVid0 action0 k0 =
           go ctl' read written' writtenSeq' createdOuterSeq nextVid (k x)
 
       ThrowStm e ->
+        {-# SCC "execAtomically.go.ThrowStm" #-}
+        case ctl of 
+          AtomicallyFrame -> do
+            -- Revert all the TVar writes
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            k0 $ StmTxAborted [] (toException e)
+
+          BranchFrame (CatchStmA h) k writtenOuter writtenOuterSeq createdOuterSeq ctl' ->
+            {-# SCC "execAtomically.go.branchFrame" #-} do
+            -- Revert all the TVar writes within this orElse
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            -- Execute the catch handler with an empty written set
+            let ctl'' = BranchFrame EmptyStmA k writtenOuter writtenOuterSeq createdOuterSeq ctl'
+            go ctl'' read Map.empty [] [] nextVid (h e)
+            -- 
+          BranchFrame _ _k writtenOuter writtenOuterSeq createdOuterSeq ctl' ->
+            {-# SCC "execAtomically.go.branchFrame" #-} do
+            -- Revert all the TVar writes within this orElse
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            go ctl' read writtenOuter writtenOuterSeq createdOuterSeq nextVid (ThrowStm e)
+
+      CatchStm a h k -> 
         {-# SCC "execAtomically.go.ThrowStm" #-} do
-        -- Revert all the TVar writes
-        !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
-        k0 $ StmTxAborted [] (toException e)
+        -- Execute the left side in a new frame with an empty written set
+        let ctl' = BranchFrame (CatchStmA h) k written writtenSeq createdSeq ctl
+        go ctl' read Map.empty [] [] nextVid a
+        
 
       Retry ->
         {-# SCC "execAtomically.go.Retry" #-}

io-sim/src/Control/Monad/IOSim/Types.hs

@@ -176,6 +176,7 @@ runSTM (STM k) = k ReturnStm
 data StmA s a where
   ReturnStm    :: a -> StmA s a
   ThrowStm     :: SomeException -> StmA s a
+  CatchStm     :: StmA s a -> (SomeException -> StmA s a) -> (a -> StmA s b) -> StmA s b
 
   NewTVar      :: Maybe String -> x -> (TVar s x -> StmA s b) -> StmA s b
   LabelTVar    :: String -> TVar s a -> StmA s b -> StmA s b
@@ -314,6 +315,25 @@ instance MonadThrow (STM s) where
 instance Exceptions.MonadThrow (STM s) where
   throwM = MonadThrow.throwIO
 
+instance MonadCatch (STM s) where
+
+  catch action handler = STM $ oneShot $ \k -> CatchStm (runSTM action) (runSTM . handler') k
+    where 
+      handler' e = case fromException e of
+        Nothing -> throwIO e  -- Rethrow the exception if handler does not handle it.
+        Just e' -> handler e'
+
+  generalBracket acquire release use = do
+    resource <- acquire
+    b <- use resource `catch` \e -> do
+      _ <- release resource (ExitCaseException e)
+      throwIO e
+    c <- release resource (ExitCaseSuccess b)
+    return (b, c)
+
+instance Exceptions.MonadCatch (STM s) where
+  catch = MonadThrow.catch
+
 instance MonadCatch (IOSim s) where
   catch action handler =
     IOSim $ oneShot $ \k -> Catch (runIOSim action) (runIOSim . handler) k
@@ -836,7 +856,9 @@ data StmTxResult s a =
 
 -- | OrElse/Catch give rise to an alternate branch. A branch of a branch is an
 -- empty one.
-data BranchStmA s a = OrElseStmA (StmA s a) | EmptyStmA
+data BranchStmA s a = OrElseStmA (StmA s a)
+                    | CatchStmA (SomeException -> StmA s a)
+                    | EmptyStmA
 
 data StmStack s b a where
   -- | Executing in the context of a top level 'atomically'.

io-sim/src/Control/Monad/IOSimPOR/Internal.hs

@@ -1107,10 +1107,32 @@ execAtomically time tid tlbl nextVid0 action0 k0 =
           go ctl' read written' writtenSeq' createdOuterSeq nextVid (k x)
 
       ThrowStm e ->
+        {-# SCC "execAtomically.go.ThrowStm" #-}
+        case ctl of 
+          AtomicallyFrame -> do
+            -- Revert all the TVar writes
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            k0 $ StmTxAborted [] (toException e)
+
+          BranchFrame (CatchStmA h) k writtenOuter writtenOuterSeq createdOuterSeq ctl' ->
+            {-# SCC "execAtomically.go.branchFrame" #-} do
+            -- Revert all the TVar writes within this orElse
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            -- Execute the catch handler with an empty written set
+            let ctl'' = BranchFrame EmptyStmA k writtenOuter writtenOuterSeq createdOuterSeq ctl'
+            go ctl'' read Map.empty [] [] nextVid (h e)
+            -- 
+          BranchFrame _ _k writtenOuter writtenOuterSeq createdOuterSeq ctl' ->
+            {-# SCC "execAtomically.go.branchFrame" #-} do
+            -- Revert all the TVar writes within this orElse
+            !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
+            go ctl' read writtenOuter writtenOuterSeq createdOuterSeq nextVid (ThrowStm e)
+
+      CatchStm a h k -> 
         {-# SCC "execAtomically.go.ThrowStm" #-} do
-        -- Revert all the TVar writes
-        !_ <- traverse_ (\(SomeTVar tvar) -> revertTVar tvar) written
-        k0 $ StmTxAborted (Map.elems read) (toException e)
+        -- Execute the left side in a new frame with an empty written set
+        let ctl' = BranchFrame (CatchStmA h) k written writtenSeq createdSeq ctl
+        go ctl' read Map.empty [] [] nextVid a
 
       Retry ->
         {-# SCC "execAtomically.go.Retry" #-}

io-sim/test/Test/STM.hs

@@ -67,6 +67,7 @@ data Term (t :: Type) where
 
     Return    :: Expr t -> Term t
     Throw     :: Expr a -> Term t
+    Catch     :: Term t -> Expr a -> Term t -> Term t
     Retry     :: Term t
 
     ReadTVar  :: Name (TyVar t) -> Term t