foundry_cheatcodes/
toml.rs

1//! Implementations of [`Toml`](spec::Group::Toml) cheatcodes.
2
3use crate::{
4    Cheatcode, Cheatcodes, Result,
5    Vm::*,
6    json::{
7        check_json_key_exists, parse_json, parse_json_coerce, parse_json_coerce_default,
8        parse_json_keys, resolve_type, split_value_key,
9    },
10};
11use alloy_dyn_abi::DynSolType;
12use alloy_sol_types::SolValue;
13use foundry_common::{fmt::StructDefinitions, fs};
14use foundry_config::fs_permissions::FsAccessKind;
15use foundry_evm_core::evm::FoundryEvmNetwork;
16use serde_json::Value as JsonValue;
17use toml::Value as TomlValue;
18use toml_edit::{DocumentMut, Item, Table, TableLike};
19
20impl Cheatcode for keyExistsTomlCall {
21    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
22        let Self { toml, key } = self;
23        check_json_key_exists(&toml_to_json_string(toml)?, key)
24    }
25}
26
27impl Cheatcode for parseToml_0Call {
28    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
29        let Self { toml } = self;
30        parse_toml(
31            toml,
32            "$",
33            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
34        )
35    }
36}
37
38impl Cheatcode for parseToml_1Call {
39    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
40        let Self { toml, key } = self;
41        parse_toml(
42            toml,
43            key,
44            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
45        )
46    }
47}
48
49macro_rules! impl_parse_toml {
50    ($call:ident, $call_with_default:ident, $ty:expr) => {
51        impl Cheatcode for $call {
52            fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
53                let Self { toml, key } = self;
54                parse_toml_coerce(toml, key, &$ty)
55            }
56        }
57
58        impl Cheatcode for $call_with_default {
59            fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
60                let Self { toml, key, defaultValue } = self;
61                parse_toml_coerce_default(toml, key, &$ty, defaultValue)
62            }
63        }
64    };
65}
66
67impl_parse_toml!(parseTomlUint_0Call, parseTomlUint_1Call, DynSolType::Uint(256));
68impl_parse_toml!(
69    parseTomlUintArray_0Call,
70    parseTomlUintArray_1Call,
71    DynSolType::Array(Box::new(DynSolType::Uint(256)))
72);
73impl_parse_toml!(parseTomlInt_0Call, parseTomlInt_1Call, DynSolType::Int(256));
74impl_parse_toml!(
75    parseTomlIntArray_0Call,
76    parseTomlIntArray_1Call,
77    DynSolType::Array(Box::new(DynSolType::Int(256)))
78);
79impl_parse_toml!(parseTomlBool_0Call, parseTomlBool_1Call, DynSolType::Bool);
80impl_parse_toml!(
81    parseTomlBoolArray_0Call,
82    parseTomlBoolArray_1Call,
83    DynSolType::Array(Box::new(DynSolType::Bool))
84);
85impl_parse_toml!(parseTomlAddress_0Call, parseTomlAddress_1Call, DynSolType::Address);
86impl_parse_toml!(
87    parseTomlAddressArray_0Call,
88    parseTomlAddressArray_1Call,
89    DynSolType::Array(Box::new(DynSolType::Address))
90);
91impl_parse_toml!(parseTomlString_0Call, parseTomlString_1Call, DynSolType::String);
92impl_parse_toml!(
93    parseTomlStringArray_0Call,
94    parseTomlStringArray_1Call,
95    DynSolType::Array(Box::new(DynSolType::String))
96);
97impl_parse_toml!(parseTomlBytes_0Call, parseTomlBytes_1Call, DynSolType::Bytes);
98impl_parse_toml!(
99    parseTomlBytesArray_0Call,
100    parseTomlBytesArray_1Call,
101    DynSolType::Array(Box::new(DynSolType::Bytes))
102);
103impl_parse_toml!(parseTomlBytes32_0Call, parseTomlBytes32_1Call, DynSolType::FixedBytes(32));
104impl_parse_toml!(
105    parseTomlBytes32Array_0Call,
106    parseTomlBytes32Array_1Call,
107    DynSolType::Array(Box::new(DynSolType::FixedBytes(32)))
108);
109
110impl Cheatcode for parseTomlType_0Call {
111    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
112        let Self { toml, typeDescription } = self;
113        parse_toml_coerce(
114            toml,
115            "$",
116            &resolve_type(
117                typeDescription,
118                state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
119            )?,
120        )
121        .map(|v| v.abi_encode())
122    }
123}
124
125impl Cheatcode for parseTomlType_1Call {
126    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
127        let Self { toml, key, typeDescription } = self;
128        parse_toml_coerce(
129            toml,
130            key,
131            &resolve_type(
132                typeDescription,
133                state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
134            )?,
135        )
136        .map(|v| v.abi_encode())
137    }
138}
139
140impl Cheatcode for parseTomlTypeArrayCall {
141    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
142        let Self { toml, key, typeDescription } = self;
143        let ty = resolve_type(
144            typeDescription,
145            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
146        )?;
147        parse_toml_coerce(toml, key, &DynSolType::Array(Box::new(ty))).map(|v| v.abi_encode())
148    }
149}
150
151impl Cheatcode for parseTomlKeysCall {
152    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
153        let Self { toml, key } = self;
154        parse_toml_keys(toml, key)
155    }
156}
157
158impl Cheatcode for writeToml_0Call {
159    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
160        let Self { json, path } = self;
161        let value =
162            serde_json::from_str(json).unwrap_or_else(|_| JsonValue::String(json.to_owned()));
163
164        let toml_string = format_json_to_toml(value)?;
165        super::fs::write_file(state, path.as_ref(), toml_string.as_bytes())
166    }
167}
168
169impl Cheatcode for writeToml_1Call {
170    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
171        let Self { json: value, path, valueKey } = self;
172
173        // Read and parse the TOML file, keeping its formatting and comments.
174        // If the file doesn't exist, start with an empty document so the file is created.
175        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
176        let mut document = if data_path.exists() {
177            parse_toml_document(&fs::locked_read_to_string(&data_path)?)?
178        } else {
179            DocumentMut::new()
180        };
181        upsert_toml_value(&mut document, value, valueKey)?;
182
183        super::fs::write_file(state, path.as_ref(), document.to_string().as_bytes())
184    }
185}
186
187/// Parse
188fn parse_toml_str(toml: &str) -> Result<TomlValue> {
189    toml::from_str(toml).map_err(|e| fmt_err!("failed parsing TOML: {e}"))
190}
191
192/// Parse a TOML string and return the value at the given path.
193fn parse_toml(toml: &str, key: &str, struct_defs: Option<&StructDefinitions>) -> Result {
194    parse_json(&toml_to_json_string(toml)?, key, struct_defs)
195}
196
197/// Parse a TOML string and return the value at the given path, coercing it to the given type.
198fn parse_toml_coerce(toml: &str, key: &str, ty: &DynSolType) -> Result {
199    parse_json_coerce(&toml_to_json_string(toml)?, key, ty)
200}
201
202/// Parse a TOML string and return the value at the given path, coercing it to the given type, or
203/// return the default if the path does not exist.
204fn parse_toml_coerce_default<T: SolValue>(
205    toml: &str,
206    key: &str,
207    ty: &DynSolType,
208    default: &T,
209) -> Result {
210    parse_json_coerce_default(&toml_to_json_string(toml)?, key, ty, default)
211}
212
213/// Parse a TOML string and return an array of all keys at the given path.
214fn parse_toml_keys(toml: &str, key: &str) -> Result {
215    parse_json_keys(&toml_to_json_string(toml)?, key)
216}
217
218/// Convert a TOML string to a JSON string.
219fn toml_to_json_string(toml: &str) -> Result<String> {
220    let toml = parse_toml_str(toml)?;
221    let json = toml_to_json_value(toml);
222    serde_json::to_string(&json).map_err(|e| fmt_err!("failed to serialize JSON: {e}"))
223}
224
225/// Format a JSON value to a TOML pretty string.
226fn format_json_to_toml(json: JsonValue) -> Result<String> {
227    let toml = json_to_toml_value(json);
228    toml::to_string_pretty(&toml).map_err(|e| fmt_err!("failed to serialize TOML: {e}"))
229}
230
231/// Convert a TOML value to a JSON value.
232pub(super) fn toml_to_json_value(toml: TomlValue) -> JsonValue {
233    match toml {
234        TomlValue::String(s) => match s.as_str() {
235            "null" => JsonValue::Null,
236            _ => JsonValue::String(s),
237        },
238        TomlValue::Integer(i) => JsonValue::Number(i.into()),
239        TomlValue::Float(f) => match serde_json::Number::from_f64(f) {
240            Some(n) => JsonValue::Number(n),
241            None => JsonValue::String(f.to_string()),
242        },
243        TomlValue::Boolean(b) => JsonValue::Bool(b),
244        TomlValue::Array(a) => JsonValue::Array(a.into_iter().map(toml_to_json_value).collect()),
245        TomlValue::Table(t) => {
246            JsonValue::Object(t.into_iter().map(|(k, v)| (k, toml_to_json_value(v))).collect())
247        }
248        TomlValue::Datetime(d) => JsonValue::String(d.to_string()),
249    }
250}
251
252/// Convert a JSON value to a TOML value.
253fn json_to_toml_value(json: JsonValue) -> TomlValue {
254    match json {
255        JsonValue::String(s) => TomlValue::String(s),
256        JsonValue::Number(n) => match n.as_i64() {
257            Some(i) => TomlValue::Integer(i),
258            None => match n.as_f64() {
259                Some(f) => TomlValue::Float(f),
260                None => TomlValue::String(n.to_string()),
261            },
262        },
263        JsonValue::Bool(b) => TomlValue::Boolean(b),
264        JsonValue::Array(a) => TomlValue::Array(a.into_iter().map(json_to_toml_value).collect()),
265        JsonValue::Object(o) => {
266            TomlValue::Table(o.into_iter().map(|(k, v)| (k, json_to_toml_value(v))).collect())
267        }
268        JsonValue::Null => TomlValue::String("null".to_string()),
269    }
270}
271
272/// Parses a TOML string into a document that keeps its formatting and comments.
273fn parse_toml_document(toml: &str) -> Result<DocumentMut> {
274    toml.parse().map_err(|e| fmt_err!("failed parsing TOML: {e}"))
275}
276
277/// Inserts or replaces the value at `key` in a TOML document, creating intermediate tables if
278/// necessary.
279///
280/// Only the item at `key` is rewritten, so comments and formatting elsewhere in the document are
281/// kept.
282fn upsert_toml_value(document: &mut DocumentMut, value: &str, key: &str) -> Result<()> {
283    let parts = split_value_key(key)?;
284
285    // Separate the final key from the path.
286    // Traverse the tables, creating implicit intermediary ones if necessary.
287    if let Some((key_to_insert, path_to_parent)) = parts.split_last() {
288        let mut current_level = document.as_item_mut();
289
290        for segment in path_to_parent {
291            let is_inline = current_level.is_inline_table();
292            let Some(table) = current_level.as_table_like_mut() else {
293                return Err(fmt_err!("path segment '{segment}' does not resolve to an object."));
294            };
295            if !table.contains_key(segment) {
296                let mut intermediary = Table::new();
297                intermediary.set_implicit(true);
298                insert_toml_item(table, is_inline, segment, Item::Table(intermediary));
299            }
300            current_level = table.get_mut(segment).unwrap();
301        }
302
303        let is_inline = current_level.is_inline_table();
304        let Some(parent) = current_level.as_table_like_mut() else {
305            return Err(fmt_err!("final destination is not an object, cannot insert key."));
306        };
307
308        let value =
309            serde_json::from_str(value).unwrap_or_else(|_| JsonValue::String(value.to_owned()));
310        let mut item = json_to_toml_item(value)?;
311
312        // Replace an existing item in place: `insert` would reset the key's formatting, which holds
313        // the comments above it.
314        match parent.get_mut(key_to_insert) {
315            Some(existing) if !existing.is_none() => {
316                if existing.is_value() {
317                    // Keep inline values inline instead of turning them into table sections.
318                    item = item.into_value().map_or_else(|item| item, Item::Value);
319                }
320                let shape_changed =
321                    std::mem::discriminant(existing) != std::mem::discriminant(&item);
322                // The comments above a table header are part of the table's decor.
323                let header_prefix =
324                    tables(existing).first().and_then(|table| table.decor().prefix()).cloned();
325                match (&*existing, &mut item) {
326                    (Item::Value(old), Item::Value(new)) => {
327                        *new.decor_mut() = old.decor().clone();
328                    }
329                    (old, new) => {
330                        for (old, new) in tables(old).into_iter().zip(tables_mut(new)) {
331                            *new.decor_mut() = old.decor().clone();
332                            new.set_position(old.position());
333                            if !old.is_implicit() {
334                                new.set_implicit(false);
335                            }
336                        }
337                    }
338                }
339                let new_is_value = item.is_value();
340                *existing = item;
341                if shape_changed {
342                    let mut key = parent.key_mut(key_to_insert).expect("replaced key must exist");
343                    let decor = key.leaf_decor_mut();
344                    decor.clear();
345                    if new_is_value {
346                        decor.set_suffix(" ");
347                        if let Some(prefix) = header_prefix {
348                            decor.set_prefix(prefix);
349                        }
350                    }
351                }
352            }
353            _ => insert_toml_item(parent, is_inline, key_to_insert, item),
354        }
355    }
356
357    Ok(())
358}
359
360/// Inserts a new `item` at `key` into `parent`.
361///
362/// Inline tables can only hold values. There, the new value takes over the whitespace or comment
363/// that followed the previous last value, so the inline table keeps its layout.
364fn insert_toml_item(parent: &mut dyn TableLike, is_inline: bool, key: &str, mut item: Item) {
365    if is_inline {
366        let trailing =
367            parent.iter_mut().filter_map(|(_, item)| item.as_value_mut()).last().and_then(|last| {
368                let suffix = last.decor().suffix().cloned();
369                last.decor_mut().set_suffix("");
370                suffix
371            });
372        item = item.into_value().map_or_else(|item| item, Item::Value);
373        if let Some(value) = item.as_value_mut() {
374            value.decor_mut().clear();
375            if let Some(trailing) = trailing {
376                value.decor_mut().set_suffix(trailing);
377            }
378        }
379    }
380    parent.insert(key, item);
381}
382
383/// Converts a JSON value to a TOML item, formatted the same way as [`format_json_to_toml`].
384fn json_to_toml_item(value: JsonValue) -> Result<Item> {
385    const KEY: &str = "value";
386
387    let wrapper = JsonValue::Object([(KEY.to_string(), value)].into_iter().collect());
388    let mut item = parse_toml_document(&format_json_to_toml(wrapper)?)?
389        .remove(KEY)
390        .ok_or_else(|| fmt_err!("failed to serialize TOML value"))?;
391    // Drop the layout of the temporary document so new tables are placed after their parent.
392    reset_table_layout(&mut item);
393    Ok(item)
394}
395
396/// Recursively clears the document position and header whitespace of all tables in `item`.
397fn reset_table_layout(item: &mut Item) {
398    for table in tables_mut(item) {
399        table.set_position(None);
400        table.decor_mut().clear();
401        for (_, item) in table.iter_mut() {
402            reset_table_layout(item);
403        }
404    }
405}
406
407/// Returns the tables of `item` that are written with a header.
408fn tables(item: &Item) -> Vec<&Table> {
409    match item {
410        Item::Table(table) => vec![table],
411        Item::ArrayOfTables(array) => array.iter().collect(),
412        _ => Vec::new(),
413    }
414}
415
416/// Returns the tables of `item` that are written with a header.
417fn tables_mut(item: &mut Item) -> Vec<&mut Table> {
418    match item {
419        Item::Table(table) => vec![table],
420        Item::ArrayOfTables(array) => array.iter_mut().collect(),
421        _ => Vec::new(),
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428
429    const CONFIG: &str = r#"# Deployment config.
430
431owner = "0x000000000000000000000000000000000000dEaD" # multisig
432deployed_at = 2024-04-27T11:57:21Z
433max_supply = 0xffff_ffff
434salt = 'literal-string'
435
436# Mainnet settings.
437[mainnet]
438token = "0x0000000000000000000000000000000000000000" # filled by script
439limits = { daily = 1_000, weekly = 5_000 }
440
441# Base settings.
442[base]
443chain_id = 8453
444"#;
445
446    fn upsert(toml: &str, value: &str, key: &str) -> Result<String> {
447        let mut document = parse_toml_document(toml)?;
448        upsert_toml_value(&mut document, value, key)?;
449        Ok(document.to_string())
450    }
451
452    #[test]
453    fn upsert_toml_keeps_comment_above_key() {
454        let toml = "key1 = \"1\"\n\n# this is key2\nkey2 = \"2\"\n";
455        assert_eq!(
456            upsert(toml, "abcd", ".key2").unwrap(),
457            "key1 = \"1\"\n\n# this is key2\nkey2 = \"abcd\"\n"
458        );
459    }
460
461    #[test]
462    fn upsert_toml_replaces_only_the_value() {
463        let address = "0x000000000000000000000000000000000000bEEF";
464        assert_eq!(
465            upsert(CONFIG, address, ".mainnet.token").unwrap(),
466            CONFIG.replace(
467                r#"token = "0x0000000000000000000000000000000000000000""#,
468                &format!(r#"token = "{address}""#)
469            )
470        );
471        assert_eq!(
472            upsert(CONFIG, "2000", ".mainnet.limits.daily").unwrap(),
473            CONFIG.replace("daily = 1_000", "daily = 2000")
474        );
475        assert_eq!(
476            upsert(CONFIG, "{\"chain_id\": 10}", "base").unwrap(),
477            CONFIG.replace("chain_id = 8453", "chain_id = 10")
478        );
479        assert_eq!(
480            upsert(CONFIG, "{\"daily\": 2000}", ".mainnet.limits").unwrap(),
481            CONFIG.replace("daily = 1_000, weekly = 5_000", "daily = 2000")
482        );
483    }
484
485    #[test]
486    fn upsert_toml_adds_keys() {
487        assert_eq!(
488            upsert(CONFIG, "30000000", ".base.gas_limit").unwrap(),
489            format!("{CONFIG}gas_limit = 30000000\n")
490        );
491        assert_eq!(
492            upsert(CONFIG, "{\"block\": 123}", ".optimism.contracts").unwrap(),
493            format!("{CONFIG}\n[optimism.contracts]\nblock = 123\n")
494        );
495        assert_eq!(
496            upsert(CONFIG, "{\"monthly\": 9}", ".mainnet.limits.extra").unwrap(),
497            CONFIG.replace("weekly = 5_000 }", "weekly = 5_000, extra = { monthly = 9 } }")
498        );
499    }
500
501    #[test]
502    fn upsert_toml_preserves_inline_table_formatting_when_adding_keys() {
503        let toml = "limits = { daily = 1_000, weekly = 5_000 }\n";
504        assert_eq!(
505            upsert(toml, "9000", ".limits.monthly").unwrap(),
506            "limits = { daily = 1_000, weekly = 5_000, monthly = 9000 }\n"
507        );
508        assert_eq!(
509            upsert(toml, "9000", ".limits.extra.monthly").unwrap(),
510            "limits = { daily = 1_000, weekly = 5_000, extra = { monthly = 9000 } }\n"
511        );
512
513        let toml = "limits={daily=1_000,weekly = 5_000}\n";
514        assert_eq!(
515            upsert(toml, "9000", ".limits.monthly").unwrap(),
516            "limits={daily=1_000,weekly = 5_000, monthly = 9000}\n"
517        );
518
519        let toml = "limits = {\n    # Keep this cap conservative.\n    daily  = 1_000,\n    weekly = 5_000,\n}\n";
520        assert_eq!(
521            upsert(toml, "9000", ".limits.monthly").unwrap(),
522            "limits = {\n    # Keep this cap conservative.\n    daily  = 1_000,\n    weekly = 5_000, monthly = 9000,\n}\n"
523        );
524    }
525
526    #[test]
527    fn upsert_toml_replaces_arrays_of_tables_in_place() {
528        let toml = "# Production RPC endpoints.\n[[rpc]]\nurl = \"old\"\n";
529        assert_eq!(
530            upsert(toml, r#"[{"url":"new"}]"#, ".rpc").unwrap(),
531            "# Production RPC endpoints.\n[[rpc]]\nurl = \"new\"\n"
532        );
533    }
534
535    #[test]
536    fn upsert_toml_keeps_comments_when_the_item_shape_changes() {
537        let toml = "# Deployment settings.\n[deployment] # maintained by the deploy script\nchain_id = 1\n";
538        assert_eq!(
539            upsert(toml, "disabled", ".deployment").unwrap(),
540            "# Deployment settings.\ndeployment = \"disabled\"\n"
541        );
542
543        let toml = "# Production RPC endpoint.\n[rpc]\nurl = \"a\"\n";
544        assert_eq!(
545            upsert(toml, r#"[{"url":"a"},{"url":"b"}]"#, ".rpc").unwrap(),
546            "# Production RPC endpoint.\n[[rpc]]\nurl = \"a\"\n\n[[rpc]]\nurl = \"b\"\n"
547        );
548
549        let toml = "# Production RPC endpoints.\nrpc = []\n";
550        assert_eq!(
551            upsert(toml, r#"[{"url":"a"},{"url":"b"}]"#, ".rpc").unwrap(),
552            "# Production RPC endpoints.\nrpc = [{ url = \"a\" }, { url = \"b\" }]\n"
553        );
554    }
555
556    #[test]
557    fn upsert_toml_formats_new_values_like_write_toml() {
558        let value = r#"{"list": ["0x01", "0x02"], "empty": {}, "nested": {"a": {"b": 1}}}"#;
559        let expected = format_json_to_toml(
560            serde_json::json!({ "new": serde_json::from_str::<JsonValue>(value).unwrap() }),
561        )
562        .unwrap();
563        assert_eq!(upsert("", value, ".new").unwrap(), expected);
564    }
565
566    #[test]
567    fn upsert_toml_errors() {
568        assert_eq!(
569            upsert(CONFIG, "1", ".owner.x").unwrap_err().to_string(),
570            "final destination is not an object, cannot insert key."
571        );
572        assert_eq!(
573            upsert(CONFIG, "1", ".owner.x.y").unwrap_err().to_string(),
574            "path segment 'x' does not resolve to an object."
575        );
576        assert_eq!(
577            upsert(CONFIG, "1", "$.").unwrap_err().to_string(),
578            "'valueKey' cannot be empty or just '$'"
579        );
580    }
581}
foundry_cheatcodes/toml.rs

foundry_cheatcodes/
toml.rs
