From c1b93fefad153550800e250d9d15745ee784eadc Mon Sep 17 00:00:00 2001
From: leudz <dylan.ancel@gmail.com>
Date: Fri, 6 Sep 2024 19:46:04 +0200
Subject: [PATCH] Fix workload execution panic

---
 src/world.rs          | 56 ++++++++++++++++++-------------------------
 tests/workload/mod.rs | 13 ++++++++++
 2 files changed, 36 insertions(+), 33 deletions(-)

diff --git a/src/world.rs b/src/world.rs
index f618fad6..662b7609 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -796,49 +796,39 @@ let i = world.run(sys1);
 
                 let mut skip_first = false;
                 let single_system = batch.0.filter(|_| run_if.0).or_else(|| {
-                    skip_first = true;
-                    batch.1.first().copied().filter(|_| run_if.1[0])
+                    let system = batch.1.first().copied().filter(|_| run_if.1[0]);
+
+                    if system.is_some() {
+                        skip_first = true;
+                    }
+
+                    system
                 });
 
                 rayon::in_place_scope(|scope| {
                     scope.spawn(|_| {
-                        if batch.1.len() == 1 && !skip_first {
-                            if !run_if.1[0] {
-                                return;
+                        use rayon::prelude::*;
+
+                        let start = if skip_first {
+                            1
+                        } else {
+                            0
+                        };
+
+                        result = batch.1[start..].par_iter().zip(&run_if.1[start..]).try_for_each(|(&index, should_run)| {
+                            if !should_run {
+                                return Ok(());
                             }
 
                             #[cfg(feature = "tracing")]
-                            let system_span = tracing::info_span!(parent: parent_span.clone(), "system", name = ?system_names[batch.1[0]]);
+                            let system_span = tracing::info_span!(parent: parent_span.clone(), "system", name = ?system_names[index]);
                             #[cfg(feature = "tracing")]
                             let _system_span = system_span.enter();
 
-                            result = systems[batch.1[0]](self).map_err(|err| {
-                                error::RunWorkload::Run((system_names[batch.1[0]].clone(), err))
-                            });
-                        } else {
-                            use rayon::prelude::*;
-
-                            let start = if skip_first {
-                                1
-                            } else {
-                                0
-                            };
-
-                            result = batch.1[start..].par_iter().zip(&run_if.1[start..]).try_for_each(|(&index, should_run)| {
-                                if !should_run {
-                                    return Ok(());
-                                }
-
-                                #[cfg(feature = "tracing")]
-                                let system_span = tracing::info_span!(parent: parent_span.clone(), "system", name = ?system_names[index]);
-                                #[cfg(feature = "tracing")]
-                                let _system_span = system_span.enter();
-
-                                (systems[index])(self).map_err(|err| {
-                                    error::RunWorkload::Run((system_names[index].clone(), err))
-                                })
-                            });
-                        }
+                            (systems[index])(self).map_err(|err| {
+                                error::RunWorkload::Run((system_names[index].clone(), err))
+                            })
+                        });
                     });
 
                     if let Some(index) = single_system {
diff --git a/tests/workload/mod.rs b/tests/workload/mod.rs
index 06a15a22..194829dc 100644
--- a/tests/workload/mod.rs
+++ b/tests/workload/mod.rs
@@ -236,3 +236,16 @@ fn single_system_run_on_same_thread() {
 
     world.run_default_workload().unwrap();
 }
+
+/// Make sure that we don't panic in this scenario:
+/// - one system running on main thread
+/// - the system has a run_if that evaluates to false
+/// - no other system
+#[test]
+fn skip_first() {
+    let world = World::new();
+
+    world.add_workload(|| (|_: AllStoragesViewMut| {}).run_if(|| false));
+
+    world.run_default_workload().unwrap();
+}