3 files changed, 42 insertions, 22 deletions
diff --git a/funcs.go b/funcs.go
index f87efb8..a7ff6e0 100644
--- a/funcs.go
+++ b/funcs.go
@@ -1,16 +1,33 @@
 package stm
 
 // Atomically executes the atomic function fn.
-func Atomically(fn func(*Tx)) {
+func Atomically(fn func(*Tx)) interface{} {
 retry:
 	// run the transaction
 	tx := &Tx{
 		reads:  make(map[*Var]uint64),
 		writes: make(map[*Var]interface{}),
 	}
-	if catchRetry(fn, tx) {
-		// wait for one of the variables we read to change before retrying
-		tx.wait()
+	var ret interface{}
+	if func() (retry bool) {
+		defer func() {
+			r := recover()
+			if r == nil {
+				return
+			}
+			if _ret, ok := r.(_return); ok {
+				ret = _ret.value
+			} else if r == Retry {
+				// wait for one of the variables we read to change before retrying
+				tx.wait()
+				retry = true
+			} else {
+				panic(r)
+			}
+		}()
+		fn(tx)
+		return false
+	}() {
 		goto retry
 	}
 	// verify the read log
@@ -25,6 +42,7 @@ retry:
 		globalCond.Broadcast()
 	}
 	globalLock.Unlock()
+	return ret
 }
 
 // AtomicGet is a helper function that atomically reads a value.
diff --git a/stm_test.go b/stm_test.go
index f030a8a..527e480 100644
--- a/stm_test.go
+++ b/stm_test.go
@@ -223,25 +223,19 @@ func testPingPong(t testing.TB, n int, afterHit func(string)) {
 	hits := NewVar(0)
 	ready := NewVar(true) // The ball is ready for hitting.
 	bat := func(from, to interface{}, noise string) {
-		done := false
-		for {
-			Atomically(func(tx *Tx) {
-				if tx.Get(doneVar).(bool) {
-					done = true
-					return
-				}
-				tx.Assert(tx.Get(ready).(bool))
-				if tx.Get(ball) == from {
-					tx.Set(ball, to)
-					tx.Set(hits, tx.Get(hits).(int)+1)
-					tx.Set(ready, false)
-					return
-				}
-				tx.Retry()
-			})
-			if done {
-				break
+		for !Atomically(func(tx *Tx) {
+			if tx.Get(doneVar).(bool) {
+				tx.Return(true)
+			}
+			tx.Assert(tx.Get(ready).(bool))
+			if tx.Get(ball) == from {
+				tx.Set(ball, to)
+				tx.Set(hits, tx.Get(hits).(int)+1)
+				tx.Set(ready, false)
+				tx.Return(false)
 			}
+			tx.Retry()
+		}).(bool) {
 			afterHit(noise)
 			AtomicSet(ready, true)
 		}
diff --git a/tx.go b/tx.go
index c2dcbdc..123c490 100644
--- a/tx.go
+++ b/tx.go
@@ -73,3 +73,11 @@ func (tx *Tx) Assert(p bool) {
 		tx.Retry()
 	}
 }
+
+func (tx *Tx) Return(v interface{}) {
+	panic(_return{v})
+}
+
+type _return struct {
+	value interface{}
+}