From cae18bba06c8d0474921721e098f5f0dbd620846 Mon Sep 17 00:00:00 2001
From: findolor <anakisci@gmail.com>
Date: Mon, 24 Feb 2025 10:31:26 +0300
Subject: [PATCH] add new method for parsing all the order token keys

---
 crates/settings/src/order.rs        | 37 ++++++++++++++++++++++++++++-
 crates/settings/src/yaml/dotrain.rs |  5 ++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/crates/settings/src/order.rs b/crates/settings/src/order.rs
index 3b05cbddf..3517bc543 100644
--- a/crates/settings/src/order.rs
+++ b/crates/settings/src/order.rs
@@ -2,7 +2,7 @@ use crate::{yaml::FieldErrorKind, *};
 use alloy::primitives::{private::rand, U256};
 use serde::{Deserialize, Serialize};
 use std::{
-    collections::HashMap,
+    collections::{BTreeSet, HashMap},
     str::FromStr,
     sync::{Arc, RwLock},
 };
@@ -434,6 +434,41 @@ impl OrderCfg {
 
         Ok(vault_ids)
     }
+
+    pub fn parse_io_token_keys(
+        documents: Vec<Arc<RwLock<StrictYaml>>>,
+        order_key: &str,
+    ) -> Result<Vec<String>, YamlError> {
+        let mut token_keys = BTreeSet::new();
+
+        for document in documents {
+            let document_read = document.read().map_err(|_| YamlError::ReadLockError)?;
+
+            if let Ok(orders_hash) = require_hash(&document_read, Some("orders"), None) {
+                if let Some(order_yaml) =
+                    orders_hash.get(&StrictYaml::String(order_key.to_string()))
+                {
+                    let location = format!("order '{}'", order_key);
+
+                    let inputs = require_vec(order_yaml, "inputs", Some(location.clone()))?;
+                    let outputs = require_vec(order_yaml, "outputs", Some(location.clone()))?;
+
+                    for input in inputs {
+                        let token_key =
+                            require_string(input, Some("token"), Some(location.clone()))?;
+                        token_keys.insert(token_key);
+                    }
+                    for output in outputs {
+                        let token_key =
+                            require_string(output, Some("token"), Some(location.clone()))?;
+                        token_keys.insert(token_key);
+                    }
+                }
+            }
+        }
+
+        Ok(token_keys.into_iter().collect())
+    }
 }
 
 impl YamlParsableHash for OrderCfg {
diff --git a/crates/settings/src/yaml/dotrain.rs b/crates/settings/src/yaml/dotrain.rs
index 4aac4842e..582e2c081 100644
--- a/crates/settings/src/yaml/dotrain.rs
+++ b/crates/settings/src/yaml/dotrain.rs
@@ -320,6 +320,11 @@ mod tests {
             OrderCfg::parse_vault_ids(dotrain_yaml.documents.clone(), &order.key, false).unwrap();
         assert_eq!(output_vault_ids.len(), 1);
         assert_eq!(output_vault_ids[0], Some("2".to_string()));
+        let io_token_keys =
+            OrderCfg::parse_io_token_keys(dotrain_yaml.documents.clone(), &order.key).unwrap();
+        assert_eq!(io_token_keys.len(), 2);
+        assert_eq!(io_token_keys[0], "token1");
+        assert_eq!(io_token_keys[1], "token2");
 
         let scenario_keys = dotrain_yaml.get_scenario_keys().unwrap();
         assert_eq!(scenario_keys.len(), 3);