1 files changed, 79 insertions, 35 deletions

diff --git a/src/libstd/unstable/finally.rs b/src/libstd/unstable/finally.rs
index 6faf69d2bb9..433accecdbc 100644
--- a/src/libstd/unstable/finally.rs
+++ b/src/libstd/unstable/finally.rs
@@ -12,6 +12,11 @@
 The Finally trait provides a method, `finally` on
 stack closures that emulates Java-style try/finally blocks.
 
+Using the `finally` method is sometimes convenient, but the type rules
+prohibit any shared, mutable state between the "try" case and the
+"finally" case. For advanced cases, the `try_finally` function can
+also be used. See that function for more details.
+
 # Example
 
  ```
@@ -31,53 +36,89 @@ pub trait Finally<T> {
     fn finally(&self, dtor: ||) -> T;
 }
 
-macro_rules! finally_fn {
-    ($fnty:ty) => {
-        impl<T> Finally<T> for $fnty {
-            fn finally(&self, dtor: ||) -> T {
-                let _d = Finallyalizer {
-                    dtor: dtor
-                };
-                (*self)()
-            }
-        }
+impl<'a,T> Finally<T> for 'a || -> T {
+    fn finally(&self, dtor: ||) -> T {
+        try_finally(&mut (), (),
+                    |_, _| (*self)(),
+                    |_| dtor())
     }
 }
 
-impl<'a,T> Finally<T> for 'a || -> T {
+impl<T> Finally<T> for fn() -> T {
     fn finally(&self, dtor: ||) -> T {
-        let _d = Finallyalizer {
-            dtor: dtor
-        };
-
-        (*self)()
+        try_finally(&mut (), (),
+                    |_, _| (*self)(),
+                    |_| dtor())
     }
 }
 
-finally_fn!(extern "Rust" fn() -> T)
+/**
+ * The most general form of the `finally` functions. The function
+ * `try_fn` will be invoked first; whether or not it fails, the
+ * function `finally_fn` will be invoked next. The two parameters
+ * `mutate` and `drop` are used to thread state through the two
+ * closures. `mutate` is used for any shared, mutable state that both
+ * closures require access to; `drop` is used for any state that the
+ * `try_fn` requires ownership of.
+ *
+ * **WARNING:** While shared, mutable state between the try and finally
+ * function is often necessary, one must be very careful; the `try`
+ * function could have failed at any point, so the values of the shared
+ * state may be inconsistent.
+ *
+ * # Example
+ *
+ * ```
+ * struct State<'a> { buffer: &'a mut [u8], len: uint }
+ * let mut state = State { buffer: buf, len: 0 };
+ * try_finally(
+ *     &mut state, (),
+ *     |state, ()| {
+ *         // use state.buffer, state.len
+ *     }
+ *     |state| {
+ *         // use state.buffer, state.len to cleanup
+ *     })
+ * ```
+ */
+pub fn try_finally<T,U,R>(mutate: &mut T,
+                          drop: U,
+                          try_fn: |&mut T, U| -> R,
+                          finally_fn: |&mut T|)
+                          -> R {
+    let f = Finallyalizer {
+        mutate: mutate,
+        dtor: finally_fn,
+    };
+    try_fn(&mut *f.mutate, drop)
+}
 
-struct Finallyalizer<'a> {
-    dtor: 'a ||
+struct Finallyalizer<'a,A> {
+    mutate: &'a mut A,
+    dtor: 'a |&mut A|
 }
 
 #[unsafe_destructor]
-impl<'a> Drop for Finallyalizer<'a> {
+impl<'a,A> Drop for Finallyalizer<'a,A> {
     #[inline]
     fn drop(&mut self) {
-        (self.dtor)();
+        (self.dtor)(self.mutate);
     }
 }
 
 #[test]
 fn test_success() {
     let mut i = 0;
-    (|| {
-        i = 10;
-    }).finally(|| {
-        assert!(!failing());
-        assert_eq!(i, 10);
-        i = 20;
-    });
+    try_finally(
+        &mut i, (),
+        |i, ()| {
+            *i = 10;
+        },
+        |i| {
+            assert!(!failing());
+            assert_eq!(*i, 10);
+            *i = 20;
+        });
     assert_eq!(i, 20);
 }
 
@@ -85,13 +126,16 @@ fn test_success() {
 #[should_fail]
 fn test_fail() {
     let mut i = 0;
-    (|| {
-        i = 10;
-        fail!();
-    }).finally(|| {
-        assert!(failing());
-        assert_eq!(i, 10);
-    })
+    try_finally(
+        &mut i, (),
+        |i, ()| {
+            *i = 10;
+            fail!();
+        },
+        |i| {
+            assert!(failing());
+            assert_eq!(*i, 10);
+        })
 }
 
 #[test]