1use std::collections::BTreeSet;
2
3use crate::Config;
4use alloy_primitives::map::HashMap;
5use figment::{
6 Figment, Profile, Provider,
7 value::{Dict, Map, Value},
8};
9use foundry_compilers::ProjectCompileOutput;
10use foundry_evm_networks::NetworkVariant;
11use itertools::Itertools;
12
13mod natspec;
14pub use natspec::*;
15
16const INLINE_CONFIG_PREFIX: &str = "forge-config:";
17
18type DataMap = Map<Profile, Dict>;
19
20#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
22pub enum InlineConfigErrorKind {
23 #[error(transparent)]
25 Parse(#[from] toml::de::Error),
26 #[error("invalid profile `{0}`; valid profiles: {1}")]
28 InvalidProfile(String, String),
29 #[error("invalid @custom:halmos annotation: {0}")]
31 InvalidHalmosConfig(String),
32}
33
34#[derive(Debug, thiserror::Error)]
37#[error("Inline config error at {location}: {kind}")]
38pub struct InlineConfigError {
39 pub location: String,
42 pub kind: InlineConfigErrorKind,
44}
45
46#[derive(Clone, Debug, Default)]
50pub struct InlineConfig {
51 contract_level: HashMap<String, DataMap>,
53 fn_level: HashMap<(String, String), DataMap>,
55}
56
57impl InlineConfig {
58 pub fn new() -> Self {
60 Self::default()
61 }
62
63 pub fn new_parsed(output: &ProjectCompileOutput, config: &Config) -> eyre::Result<Self> {
66 let natspecs: Vec<NatSpec> = NatSpec::parse(output, &config.root);
67 let profiles = &config.profiles;
68 let mut inline = Self::new();
69 for natspec in &natspecs {
70 inline.insert(natspec)?;
71 natspec.validate_profiles(profiles)?;
73 }
74 Ok(inline)
75 }
76
77 pub fn insert(&mut self, natspec: &NatSpec) -> Result<(), InlineConfigError> {
79 let map = if let Some(function) = &natspec.function {
80 self.fn_level.entry((natspec.contract.clone(), function.clone())).or_default()
81 } else {
82 self.contract_level.entry(natspec.contract.clone()).or_default()
83 };
84 if let Some(data) = parse_config_values(natspec, natspec.halmos_config_values()?)? {
85 extend_data_map(map, &data);
86 }
87 if let Some(data) = parse_config_values(natspec, natspec.config_values())? {
88 extend_data_map(map, &data);
89 }
90 Ok(())
91 }
92
93 pub const fn provide<'a>(
96 &'a self,
97 contract: &'a str,
98 function: &'a str,
99 ) -> InlineConfigProvider<'a> {
100 InlineConfigProvider { inline: self, contract, function }
101 }
102
103 pub fn merge(&self, contract: &str, function: &str, base: &Config) -> Figment {
106 Figment::from(base).merge(self.provide(contract, function))
107 }
108
109 pub fn contains_contract(&self, contract: &str) -> bool {
111 self.get_contract(contract).is_some_and(|map| !map.is_empty())
112 }
113
114 pub fn contains_function(&self, contract: &str, function: &str) -> bool {
118 self.get_function(contract, function).is_some_and(|map| !map.is_empty())
119 }
120
121 pub fn network_for(
124 &self,
125 profile: &Profile,
126 contract: &str,
127 function: &str,
128 ) -> Option<NetworkVariant> {
129 let data = self.provide(contract, function).data().ok()?;
130 let dict = data.get(profile).or_else(|| data.get(&Profile::Default))?;
131 if let Some(Value::Dict(_, networks)) = dict.get("networks")
132 && let Some(Value::String(_, s)) = networks.get("network")
133 {
134 return s.parse().ok();
135 }
136 None
137 }
138
139 pub fn referenced_override_networks(&self, profile: &Profile) -> Vec<NetworkVariant> {
143 let mut seen = BTreeSet::new();
144 for (contract, function) in self.fn_level.keys() {
145 if let Some(v) = self.network_for(profile, contract, function) {
146 seen.insert(v);
147 }
148 }
149 for contract in self.contract_level.keys() {
150 if let Some(v) = self.network_for(profile, contract, "") {
151 seen.insert(v);
152 }
153 }
154 seen.into_iter().collect()
155 }
156
157 fn get_contract(&self, contract: &str) -> Option<&DataMap> {
158 self.contract_level.get(contract)
159 }
160
161 fn get_function(&self, contract: &str, function: &str) -> Option<&DataMap> {
162 let key = (contract.to_string(), function.to_string());
163 self.fn_level.get(&key)
164 }
165}
166
167fn parse_config_values<'a>(
168 natspec: &NatSpec,
169 values: impl IntoIterator<Item = impl std::borrow::Borrow<str> + 'a>,
170) -> Result<Option<DataMap>, InlineConfigError> {
171 let joined = values
172 .into_iter()
173 .map(|s| {
174 let s = s.borrow();
175 if let Some(idx) = s.find('=') {
177 s[..idx].replace('-', "_") + &s[idx..]
178 } else {
179 s.to_string()
180 }
181 })
182 .format("\n")
183 .to_string();
184 if joined.is_empty() {
185 return Ok(None);
186 }
187 let data = toml::from_str::<DataMap>(&joined).map_err(|e| InlineConfigError {
188 location: natspec.location_string(),
189 kind: InlineConfigErrorKind::Parse(e),
190 })?;
191 Ok(Some(data))
192}
193
194#[derive(Clone, Debug)]
198pub struct InlineConfigProvider<'a> {
199 inline: &'a InlineConfig,
200 contract: &'a str,
201 function: &'a str,
202}
203
204impl Provider for InlineConfigProvider<'_> {
205 fn metadata(&self) -> figment::Metadata {
206 figment::Metadata::named("inline config")
207 }
208
209 fn data(&self) -> figment::Result<DataMap> {
210 let mut map = DataMap::new();
211 if let Some(new) = self.inline.get_contract(self.contract) {
212 extend_data_map(&mut map, new);
213 }
214 if let Some(new) = self.inline.get_function(self.contract, self.function) {
215 extend_data_map(&mut map, new);
216 }
217 Ok(map)
218 }
219}
220
221fn extend_data_map(map: &mut DataMap, new: &DataMap) {
222 for (profile, data) in new {
223 extend_dict(map.entry(profile.clone()).or_default(), data);
224 }
225}
226
227fn extend_dict(dict: &mut Dict, new: &Dict) {
228 for (k, v) in new {
229 match dict.entry(k.clone()) {
230 std::collections::btree_map::Entry::Vacant(entry) => {
231 entry.insert(v.clone());
232 }
233 std::collections::btree_map::Entry::Occupied(entry) => {
234 extend_value(entry.into_mut(), v);
235 }
236 }
237 }
238}
239
240fn extend_value(value: &mut Value, new: &Value) {
241 match (value, new) {
242 (Value::Dict(tag, dict), Value::Dict(new_tag, new_dict)) => {
243 *tag = *new_tag;
244 extend_dict(dict, new_dict);
245 }
246 (value, new) => *value = new.clone(),
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253
254 fn natspec(docs: &str) -> NatSpec {
255 NatSpec {
256 contract: "test/Symbolic.t.sol:Symbolic".to_string(),
257 function: Some("check".to_string()),
258 line: "10:5".to_string(),
259 docs: docs.to_string(),
260 }
261 }
262
263 #[test]
264 fn legacy_halmos_array_lengths_feed_symbolic_inline_config() {
265 let mut inline = InlineConfig::new();
266 inline
267 .insert(&natspec(
268 "@custom:halmos --array-lengths 2,4 --invariant-depth 12 --width 8 --depth 99",
269 ))
270 .unwrap();
271
272 let config = Config::default()
273 .merge_inline_provider(inline.provide("test/Symbolic.t.sol:Symbolic", "check"))
274 .unwrap();
275
276 assert_eq!(config.symbolic.array_lengths, vec![2, 4]);
277 assert_eq!(config.symbolic.invariant_depth, 12);
278 assert_eq!(config.symbolic.width, Some(8));
279 assert_eq!(config.symbolic.depth, Some(99));
280 }
281
282 #[test]
283 fn legacy_halmos_named_and_default_lengths_feed_symbolic_inline_config() {
284 let mut inline = InlineConfig::new();
285 inline
286 .insert(&natspec(
287 "@custom:halmos --array-lengths values={2,4},data=8 --default-array-lengths 0,1 --default-bytes-lengths 0,65",
288 ))
289 .unwrap();
290
291 let config = Config::default()
292 .merge_inline_provider(inline.provide("test/Symbolic.t.sol:Symbolic", "check"))
293 .unwrap();
294
295 assert_eq!(
296 config.symbolic.dynamic_lengths,
297 std::collections::BTreeMap::from([
298 ("data".to_string(), vec![8]),
299 ("values".to_string(), vec![2, 4]),
300 ])
301 );
302 assert_eq!(config.symbolic.default_array_lengths, vec![0, 1]);
303 assert_eq!(config.symbolic.default_bytes_lengths, vec![0, 65]);
304 }
305
306 #[test]
307 fn native_symbolic_inline_config_overrides_legacy_halmos_translation() {
308 let mut inline = InlineConfig::new();
309 inline
310 .insert(&natspec(
311 r#"
312@custom:halmos --array-lengths 2
313forge-config: default.symbolic.array_lengths = [3]
314forge-config: default.symbolic.default_dynamic_length = 4
315"#,
316 ))
317 .unwrap();
318
319 let config = Config::default()
320 .merge_inline_provider(inline.provide("test/Symbolic.t.sol:Symbolic", "check"))
321 .unwrap();
322
323 assert_eq!(config.symbolic.array_lengths, vec![3]);
324 assert_eq!(config.symbolic.default_dynamic_length, 4);
325 }
326}