From 2a70707c91dd5943f40c320ee409c74ab036ab63 Mon Sep 17 00:00:00 2001
From: Alice Hau <110418948+ahau-square@users.noreply.github.com>
Date: Thu, 10 Apr 2025 12:16:31 -0400
Subject: [PATCH] fix: goosebench selector collection issue (#2129)

---
 .../goose-bench/src/runners/model_runner.rs   | 39 ++++++++-----------
 crates/goose-bench/src/utilities.rs           | 14 -------
 2 files changed, 16 insertions(+), 37 deletions(-)

diff --git a/crates/goose-bench/src/runners/model_runner.rs b/crates/goose-bench/src/runners/model_runner.rs
index 0ae789ed..2937bd50 100644
--- a/crates/goose-bench/src/runners/model_runner.rs
+++ b/crates/goose-bench/src/runners/model_runner.rs
@@ -2,7 +2,7 @@ use crate::bench_config::{BenchEval, BenchModel, BenchRunConfig};
 use crate::eval_suites::EvaluationSuite;
 use crate::reporting::{BenchmarkResults, SuiteResult};
 use crate::runners::eval_runner::EvalRunner;
-use crate::utilities::{await_process_exits, parallel_bench_cmd, union_hashmaps};
+use crate::utilities::{await_process_exits, parallel_bench_cmd};
 use std::collections::HashMap;
 use std::fs::read_to_string;
 use std::io::{self, BufRead};
@@ -189,30 +189,23 @@ impl ModelRunner {
 
         Ok(results)
     }
+
     fn collect_evals_for_run(&self) -> HashMap<String, Vec<BenchEval>> {
         // convert suites map {suite_name => [eval_selector_str] to map suite_name => [BenchEval]
-        let suites = self
-            .config
-            .evals
-            .iter()
-            .map(|eval| {
-                EvaluationSuite::select(vec![eval.clone().selector])
-                    .iter()
-                    .map(|(suite, evals)| {
-                        let bench_evals = evals
-                            .iter()
-                            .map(|suite_eval| {
-                                let mut updated_eval = eval.clone();
-                                updated_eval.selector = (*suite_eval).to_string();
-                                updated_eval
-                            })
-                            .collect::<Vec<_>>();
-                        (suite.clone(), bench_evals)
-                    })
-                    .collect()
-            })
-            .collect();
-        union_hashmaps(suites)
+        let mut result: HashMap<String, Vec<BenchEval>> = HashMap::new();
+        for eval in self.config.evals.iter() {
+            let selected_suites = EvaluationSuite::select(vec![eval.selector.clone()]);
+            for (suite, evals) in selected_suites {
+                let entry: &mut Vec<BenchEval> = result.entry(suite).or_default();
+                entry.reserve(evals.len());
+                for suite_eval in evals {
+                    let mut updated_eval = eval.clone();
+                    updated_eval.selector = suite_eval.to_string();
+                    entry.push(updated_eval);
+                }
+            }
+        }
+        result
     }
 
     fn toolshim_envs(&self) -> Vec<(String, String)> {
diff --git a/crates/goose-bench/src/utilities.rs b/crates/goose-bench/src/utilities.rs
index 2caeba7d..0eb04412 100644
--- a/crates/goose-bench/src/utilities.rs
+++ b/crates/goose-bench/src/utilities.rs
@@ -1,21 +1,7 @@
-use std::collections::HashMap;
 use std::env;
 use std::process::{Child, Command};
 use std::thread::JoinHandle;
 
-pub fn union_hashmaps<K, V>(maps: Vec<HashMap<K, V>>) -> HashMap<K, V>
-where
-    K: Eq + std::hash::Hash,
-    V: Clone,
-{
-    // We can use the fold method to accumulate all maps into one
-    maps.into_iter().fold(HashMap::new(), |mut result, map| {
-        // For each map in the vector, extend the result with its entries
-        result.extend(map);
-        result
-    })
-}
-
 pub fn await_process_exits(
     child_processes: &mut [Child],
     handles: Vec<JoinHandle<anyhow::Result<()>>>,