moonbitlang · peter-jerry-ye · Jun 13, 2025 · Jun 12, 2025 · Jun 13, 2025 · Jun 13, 2025
diff --git a/ref/ref.mbt b/ref/ref.mbt
@@ -38,7 +38,7 @@ pub fn[T] new(x : T) -> Ref[T] {
 ///   assert_eq(@ref.new(1).map(fn(a){ a + 1 }).val, 2)
 /// }
 /// ```
-pub fn[T, R] map(self : Ref[T], f : (T) -> R) -> Ref[R] {
+pub fn[T, R] map(self : Ref[T], f : (T) -> R raise?) -> Ref[R] raise? {
   { val: f(self.val) }
 }
 
@@ -65,12 +65,20 @@ pub fn[T, R] map(self : Ref[T], f : (T) -> R) -> Ref[R] {
 ///   assert_eq(x.val, 1)
 /// }
 /// ```
-pub fn[T, R] protect(self : Ref[T], a : T, f : () -> R) -> R {
+pub fn[T, R] protect(self : Ref[T], a : T, f : () -> R raise?) -> R raise? {
   let old = self.val
   self.val = a
-  let r = f()
-  self.val = old
-  r
+  try f() catch {
+    err => {
+      self.val = old
+      raise err
+    }
+  } else {
+    r => {
+      self.val = old
+      r
+    }
+  }
 }
 
 ///|
@@ -106,7 +114,7 @@ test "swap" {
 }
 
 ///|
-pub fn[T] update(self : Ref[T], f : (T) -> T) -> Unit {
+pub fn[T] update(self : Ref[T], f : (T) -> T raise?) -> Unit raise? {
   self.val = f(self.val)
 }
 

diff --git a/ref/ref.mbti b/ref/ref.mbti
@@ -5,22 +5,22 @@ import(
 )
 
 // Values
-fn[T, R] map(Ref[T], (T) -> R) -> Ref[R]
+fn[T, R] map(Ref[T], (T) -> R raise?) -> Ref[R] raise?
 
 fn[T] new(T) -> Ref[T]
 
-fn[T, R] protect(Ref[T], T, () -> R) -> R
+fn[T, R] protect(Ref[T], T, () -> R raise?) -> R raise?
 
 fn[T] swap(Ref[T], Ref[T]) -> Unit
 
-fn[T] update(Ref[T], (T) -> T) -> Unit
+fn[T] update(Ref[T], (T) -> T raise?) -> Unit raise?
 
 // Types and methods
-fn[T, R] Ref::map(Self[T], (T) -> R) -> Self[R]
+fn[T, R] Ref::map(Self[T], (T) -> R raise?) -> Self[R] raise?
 fn[T] Ref::new(T) -> Self[T]
-fn[T, R] Ref::protect(Self[T], T, () -> R) -> R
+fn[T, R] Ref::protect(Self[T], T, () -> R raise?) -> R raise?
 fn[T] Ref::swap(Self[T], Self[T]) -> Unit
-fn[T] Ref::update(Self[T], (T) -> T) -> Unit
+fn[T] Ref::update(Self[T], (T) -> T raise?) -> Unit raise?
 impl[X : @quickcheck.Arbitrary] @quickcheck.Arbitrary for Ref[X]
 
 // Type aliases

diff --git a/ref/ref_test.mbt b/ref/ref_test.mbt
@@ -29,6 +29,21 @@ test "protect" {
   assert_eq(x.val, 1)
 }
 
+///|
+test "protect with error" {
+  let x = @ref.new(1)
+  assert_eq(x.val, 1)
+  let mut r = 0
+  let result = try? x.protect(2, fn() {
+      r = x.val
+      fail("")
+      r = -x.val
+    })
+  assert_true(result is Err(_))
+  assert_eq(r, 2)
+  assert_eq(x.val, 1)
+}
+
 ///|
 test "update with increment function" {
   let a = @ref.new(1)