From d98f72435575ddf7ff4654bc5037374f1c47731e Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Fri, 21 Jun 2024 18:51:42 -0400
Subject: [PATCH] Refactor AsyncBlock::run

---
 dust-lang/src/abstract_tree/async_block.rs | 35 ++++++++++------------
 1 file changed, 15 insertions(+), 20 deletions(-)

diff --git a/dust-lang/src/abstract_tree/async_block.rs b/dust-lang/src/abstract_tree/async_block.rs
index a0a9672..a2f3212 100644
--- a/dust-lang/src/abstract_tree/async_block.rs
+++ b/dust-lang/src/abstract_tree/async_block.rs
@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     context::Context,
-    error::{PoisonError, RuntimeError, ValidationError},
+    error::{RuntimeError, ValidationError},
 };
 
 use super::{Evaluation, ExpectedType, Run, Statement, Type, Validate};
@@ -43,28 +43,23 @@ impl Run for AsyncBlock {
         self.statements
             .into_par_iter()
             .enumerate()
-            .find_map_any(|(index, statement)| {
-                let result = statement.run(&mut _context.clone(), false);
+            .find_map_any(
+                |(index, statement)| -> Option<Result<Option<Evaluation>, RuntimeError>> {
+                    let result = statement.run(&mut _context.clone(), false);
 
-                if index == statement_count - 1 {
-                    let get_write_lock = final_result.lock();
-
-                    match get_write_lock {
-                        Ok(mut final_result) => {
-                            *final_result = result;
-                            None
-                        }
-                        Err(_error) => Some(Err(RuntimeError::RwLockPoison(PoisonError))),
+                    if result.is_err() {
+                        return Some(result);
                     }
-                } else {
+
+                    if index == statement_count - 1 {
+                        // It is safe to unwrap here because only one thread uses the Mutex
+                        *final_result.lock().unwrap() = result;
+                    }
+
                     None
-                }
-            })
-            .unwrap_or(
-                final_result
-                    .into_inner()
-                    .map_err(|_| RuntimeError::RwLockPoison(PoisonError)),
-            )?
+                },
+            )
+            .unwrap_or(final_result.into_inner()?)
     }
 }