Alternative STM implementation for Haskell

main.hs

{-
  This is a simple consistency test.
-}

import Control.Concurrent
--import Control.Concurrent.STM
import STM
import Control.Monad

main = do
  a <- newTVarIO ""
  b <- newTVarIO True
  forkIO $ forever $ atomically $ do
    a' <- readTVar a
    b' <- readTVar b
    if b' then do
      writeTVar a $ "%" ++ a'
      writeTVar b $ not b'
    else
      return ()
  forkIO $ forever $ atomically $ do
    a' <- readTVar a
    b' <- readTVar b
    if not b' then do
      writeTVar a $ "^" ++ a'
      writeTVar b $ not b'
    else
      return ()
  threadDelay 1000000
  v <- atomically $ do
    a' <- readTVar a
    b' <- readTVar b
    return (a', b')
  print v

STM.hs

{-# LANGUAGE ExistentialQuantification #-}

{-
  The idea is that each variable is an IORef that is either in a stable state 
  or is being transformed. If it is being transformed then the IORef contains a 
  MidChange value. The first argument of the constructor is the state before 
  the modification, the second argument is the state afterwards and the third 
  argument is an IORef that says whether the transaction has completed.
-}

module STM where

import Control.Concurrent (yield)
import Data.IORef

data STMState a =
  At a |
  MidChange a a (IORef Bool)

type TVar a = IORef (STMState a)

data STM a =
  Done a |
  forall x. NewTVar x (TVar x -> STM a) |
  forall x. ReadTVar (TVar x) (x -> STM a) |
  forall x. WriteTVar (TVar x) x (STM a)

instance Monad STM where
  return = Done
  m >>= f =
    case m of
      Done x -> f x
      NewTVar x c -> NewTVar x (\i -> c i >>= f)
      ReadTVar x c -> ReadTVar x (\i -> c i >>= f)
      WriteTVar v x c -> WriteTVar v x (c >>= f)

newTVar = flip NewTVar return
readTVar = flip ReadTVar return
writeTVar v x = WriteTVar v x $ return ()

newTVarIO :: a -> IO (TVar a)
newTVarIO x =
  newIORef (At x)

reduce :: IORef Bool -> TVar a -> IO ()
reduce t var = do
  atomicModifyIORef var (\v ->
    case v of
      MidChange a b t' ->
        if t == t' then
          (At b, ())
        else
          (v, ())
      _ -> (v, ()))

readTVarIO :: TVar a -> IO a
readTVarIO var = do
  var' <- readIORef var
  case var' of
    At x -> return x
    MidChange before after switch -> do
      x <- readIORef switch
      if x then do
        reduce switch var
        return after
      else
        return before

atomicModifyTVar :: TVar a -> (a -> (a, b)) -> IO b
atomicModifyTVar var op = do
  undefined

atomically :: STM a -> IO a
atomically op = do
  t <- newIORef False
  out <- retry $ trans t op
  writeIORef t True
  return out
 where
  retry :: IO (Maybe a) -> IO a
  retry op = do
    v <- op
    case v of
      Just v' -> return v'
      Nothing -> do
        yield
        retry op
  continue :: IORef Bool -> STM a -> TVar x -> IO (Maybe a)
  continue t c v = do
    out <- trans t c
    case out of
      Nothing ->
        -- transaction failed, rollback.
        atomicModifyIORef v (\v' ->
          case v' of
            At _ ->
              error "Clobbered transaction"
            MidChange x _ t' ->
              if t /= t' then
                error "Clobbered transaction"
              else
                (At x, ()))
      Just _ ->
        return ()
    return out
  transVar :: (IORef Bool) -> TVar a -> (a -> a) -> IO (Maybe a)
  transVar t x fn = do
    v <- atomicModifyIORef x (\v ->
      case v of
        At v' ->
          (MidChange v' (fn v') t, Left v')
        MidChange a b t' ->
          if t' == t then
            (MidChange a (fn b) t, Left b)
          else
            (v, Right t'))
    case v of
      Left v' -> return $ Just v'
      Right t' -> do
        s <- readIORef t'
        if s then do
          reduce t' x
          transVar t x fn
        else
          return Nothing
  trans _ (Done x) =
    return $ Just x
  trans t (NewTVar x c) = do
    var <- newTVarIO x
    trans t $ c var
  trans t (ReadTVar x c) = do
    v <- transVar t x id
    case v of
      Just v' -> continue t (c v') x
      Nothing -> return Nothing
  trans t (WriteTVar v x c) = do
    v' <- transVar t v (const x)
    case v' of
      Just _ -> continue t c v
      Nothing -> return Nothing

Some comments on the functions appreciated .. I am not able to understand whats going on after atomically

@satvikc, I completely agree, looking back at this, it's difficult to tell exactly what's going on. I'll add comments at some point when I get the time.