defer macro

3 files changed, 28 insertions, 16 deletions

diff --git a/src/defer.rs b/src/defer.rs
new file mode 100644
index 0000000..362c7d1
--- /dev/null
+++ b/src/defer.rs
@@ -0,0 +1,16 @@
+pub struct Defer<T: FnOnce()>(pub Option<T>);
+
+impl<T: FnOnce()> Drop for Defer<T> {
+    fn drop(&mut self) {
+        if let Some(f) = self.0.take() {
+            f();
+        }
+    }
+}
+
+#[macro_export]
+macro_rules! defer {
+    ($($x:tt)*) => {
+        let _defer = crate::defer::Defer(Some(move || {$($x)*}));
+    };
+}
diff --git a/src/export_fun.rs b/src/export_fun.rs
index ecbdf46..a98699c 100644
--- a/src/export_fun.rs
+++ b/src/export_fun.rs
@@ -5,6 +5,7 @@ use nix::poll::PollFlags;
 use nix::poll::poll;
 
 use crate::Session;
+use crate::defer;
 use crate::run::Input;
 use crate::run::Output;
 use crate::run::get_command_kind;
@@ -32,7 +33,6 @@ use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::mpsc;
 use std::sync::mpsc::Receiver;
-use std::sync::mpsc::Sender;
 use std::thread;
 use std::time::Duration;
 
@@ -204,7 +204,9 @@ impl Drop for SocketDropper {
 
         // wait 1s for background to exit
         if let Err(e) = sr.recv.recv_timeout(Duration::from_secs(1)) {
-            eprintln!("background thread is still running({e:?}, session might not be cleaned up\r");
+            eprintln!(
+                "background thread is still running({e:?}, session might not be cleaned up\r"
+            );
         }
     }
 }
@@ -231,19 +233,9 @@ pub fn listen(session: Arc<Mutex<Session>>) -> impl Drop {
 
     let se = session.clone();
     thread::spawn(move || {
-        struct SessionRemover {
-            send: Sender<()>,
-            path: PathBuf,
-        }
-        impl Drop for SessionRemover {
-            fn drop(&mut self) {
-                let _ = fs::remove_dir_all(&self.path);
-                let _ = self.send.send(());
-            }
-        }
-        let _session_remover = SessionRemover {
-            path: session_dir,
-            send,
+        defer! {
+            let _ = fs::remove_dir_all(session_dir);
+            let _ = send.send(());
         };
 
         let listener = UnixListener::bind(socket_path).unwrap();
diff --git a/src/main.rs b/src/main.rs
index 17c4ba0..d822a03 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,7 @@ pub mod basedir;
 pub mod completion;
 pub mod cursor;
 pub mod date;
+pub mod defer;
 pub mod export_fun;
 pub mod history;
 pub mod linebuf;
@@ -350,7 +351,10 @@ fn event_loop() {
             b'\t' => {
                 let cmd = se.line.into_bytes();
                 drop(se);
-                let comp = parse::completion_context(&cmd, &mut Executor::new_for_completion(session.clone()));
+                let comp = parse::completion_context(
+                    &cmd,
+                    &mut Executor::new_for_completion(session.clone()),
+                );
                 let mut se = session.lock().unwrap();
                 match comp.kind {
                     parse::CompletionKind::Command => todo!(),