1 files changed, 21 insertions, 29 deletions

diff --git a/funcs.go b/funcs.go
index 8694eb4..07d35ec 100644
--- a/funcs.go
+++ b/funcs.go
@@ -2,19 +2,18 @@ package stm
 
 import (
 	"math/rand"
-	"reflect"
 	"runtime/pprof"
 	"sync"
 	"time"
 )
 
 var (
-	txPool = sync.Pool{New: func() interface{} {
+	txPool = sync.Pool{New: func() any {
 		expvars.Add("new txs", 1)
 		tx := &Tx{
-			reads:    make(map[*Var]VarValue),
-			writes:   make(map[*Var]interface{}),
-			watching: make(map[*Var]struct{}),
+			reads:    make(map[txVar]VarValue),
+			writes:   make(map[txVar]any),
+			watching: make(map[txVar]struct{}),
 		}
 		tx.cond.L = &tx.mu
 		return tx
@@ -40,7 +39,7 @@ func newTx() *Tx {
 	return tx
 }
 
-func WouldBlock(fn Operation) (block bool) {
+func WouldBlock[R any](fn Operation[R]) (block bool) {
 	tx := newTx()
 	tx.reset()
 	_, block = catchRetry(fn, tx)
@@ -52,7 +51,7 @@ func WouldBlock(fn Operation) (block bool) {
 }
 
 // Atomically executes the atomic function fn.
-func Atomically(op Operation) interface{} {
+func Atomically[R any](op Operation[R]) R {
 	expvars.Add("atomically", 1)
 	// run the transaction
 	tx := newTx()
@@ -104,12 +103,12 @@ retry:
 }
 
 // AtomicGet is a helper function that atomically reads a value.
-func AtomicGet(v *Var) interface{} {
-	return v.value.Load().(VarValue).Get()
+func AtomicGet[T any](v *Var[T]) T {
+	return v.value.Load().Get().(T)
 }
 
 // AtomicSet is a helper function that atomically writes a value.
-func AtomicSet(v *Var, val interface{}) {
+func AtomicSet[T any](v *Var[T], val T) {
 	v.mu.Lock()
 	v.changeValue(val)
 	v.mu.Unlock()
@@ -117,20 +116,19 @@ func AtomicSet(v *Var, val interface{}) {
 
 // Compose is a helper function that composes multiple transactions into a
 // single transaction.
-func Compose(fns ...Operation) Operation {
-	return func(tx *Tx) interface{} {
+func Compose[R any](fns ...Operation[R]) Operation[struct{}] {
+	return VoidOperation(func(tx *Tx) {
 		for _, f := range fns {
 			f(tx)
 		}
-		return nil
-	}
+	})
 }
 
 // Select runs the supplied functions in order. Execution stops when a
 // function succeeds without calling Retry. If no functions succeed, the
 // entire selection will be retried.
-func Select(fns ...Operation) Operation {
-	return func(tx *Tx) interface{} {
+func Select[R any](fns ...Operation[R]) Operation[R] {
+	return func(tx *Tx) R {
 		switch len(fns) {
 		case 0:
 			// empty Select blocks forever
@@ -140,7 +138,7 @@ func Select(fns ...Operation) Operation {
 			return fns[0](tx)
 		default:
 			oldWrites := tx.writes
-			tx.writes = make(map[*Var]interface{}, len(oldWrites))
+			tx.writes = make(map[txVar]any, len(oldWrites))
 			for k, v := range oldWrites {
 				tx.writes[k] = v
 			}
@@ -155,23 +153,17 @@ func Select(fns ...Operation) Operation {
 	}
 }
 
-type Operation func(*Tx) interface{}
+type Operation[R any] func(*Tx) R
 
-func VoidOperation(f func(*Tx)) Operation {
-	return func(tx *Tx) interface{} {
+func VoidOperation(f func(*Tx)) Operation[struct{}] {
+	return func(tx *Tx) struct{} {
 		f(tx)
-		return nil
+		return struct{}{}
 	}
 }
 
-func AtomicModify(v *Var, f interface{}) {
-	r := reflect.ValueOf(f)
+func AtomicModify[T any](v *Var[T], f func(T) T) {
 	Atomically(VoidOperation(func(tx *Tx) {
-		cur := reflect.ValueOf(tx.Get(v))
-		out := r.Call([]reflect.Value{cur})
-		if lenOut := len(out); lenOut != 1 {
-			panic(lenOut)
-		}
-		tx.Set(v, out[0].Interface())
+		v.Set(tx, f(v.Get(tx)))
 	}))
 }