chisel/
dispatcher.rs

Help
1//! Dispatcher
2//!
3//! This module contains the `ChiselDispatcher` struct, which handles the dispatching
4//! of both builtin commands and Solidity snippets.
5
6use crate::{
7    prelude::{
8        ChiselCommand, ChiselResult, ChiselSession, CmdCategory, CmdDescriptor,
9        SessionSourceConfig, SolidityHelper,
10    },
11    session_source::SessionSource,
12};
13use alloy_json_abi::{InternalType, JsonAbi};
14use alloy_primitives::{hex, Address};
15use forge_fmt::FormatterConfig;
16use foundry_config::{Config, RpcEndpointUrl};
17use foundry_evm::{
18    decode::decode_console_logs,
19    traces::{
20        decode_trace_arena,
21        identifier::{SignaturesIdentifier, TraceIdentifiers},
22        render_trace_arena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind,
23    },
24};
25use regex::Regex;
26use reqwest::Url;
27use serde::{Deserialize, Serialize};
28use solang_parser::diagnostics::Diagnostic;
29use std::{
30    borrow::Cow,
31    error::Error,
32    io::Write,
33    path::{Path, PathBuf},
34    process::Command,
35    sync::LazyLock,
36};
37use strum::IntoEnumIterator;
38use tracing::debug;
39use yansi::Paint;
40
41/// Prompt arrow character.
42pub const PROMPT_ARROW: char = '➜';
43/// Prompt arrow string.
44pub const PROMPT_ARROW_STR: &str = "➜";
45const DEFAULT_PROMPT: &str = "➜ ";
46
47/// Command leader character
48pub const COMMAND_LEADER: char = '!';
49/// Chisel character
50pub const CHISEL_CHAR: &str = "⚒️";
51
52/// Matches Solidity comments
53static COMMENT_RE: LazyLock<Regex> =
54    LazyLock::new(|| Regex::new(r"^\s*(?://.*\s*$)|(/*[\s\S]*?\*/\s*$)").unwrap());
55
56/// Matches Ethereum addresses that are not strings
57static ADDRESS_RE: LazyLock<Regex> = LazyLock::new(|| {
58    Regex::new(r#"(?m)(([^"']\s*)|^)(?P<address>0x[a-fA-F0-9]{40})((\s*[^"'\w])|$)"#).unwrap()
59});
60
61/// Chisel input dispatcher
62#[derive(Debug)]
63pub struct ChiselDispatcher {
64    /// A Chisel Session
65    pub session: ChiselSession,
66}
67
68/// Chisel dispatch result variants
69#[derive(Debug)]
70pub enum DispatchResult {
71    /// A Generic Dispatch Success
72    Success(Option<String>),
73    /// A Generic Failure
74    Failure(Option<String>),
75    /// A successful ChiselCommand Execution
76    CommandSuccess(Option<String>),
77    /// A failure to parse a Chisel Command
78    UnrecognizedCommand(Box<dyn Error>),
79    /// The solang parser failed
80    SolangParserFailed(Vec<Diagnostic>),
81    /// A Command Failed with error message
82    CommandFailed(String),
83    /// File IO Error
84    FileIoError(Box<dyn Error>),
85}
86
87impl DispatchResult {
88    /// Returns `true` if the result is an error.
89    pub fn is_error(&self) -> bool {
90        matches!(
91            self,
92            Self::Failure(_) |
93                Self::CommandFailed(_) |
94                Self::UnrecognizedCommand(_) |
95                Self::SolangParserFailed(_) |
96                Self::FileIoError(_)
97        )
98    }
99}
100
101/// A response from the Etherscan API's `getabi` action
102#[derive(Debug, Serialize, Deserialize)]
103pub struct EtherscanABIResponse {
104    /// The status of the response
105    /// "1" = success | "0" = failure
106    pub status: String,
107    /// The message supplied by the API
108    pub message: String,
109    /// The result returned by the API. Will be `None` if the request failed.
110    pub result: Option<String>,
111}
112
113/// Used to format ABI parameters into valid solidity function / error / event param syntax
114/// TODO: Smarter resolution of storage location, defaults to "memory" for all types
115/// that cannot be stored on the stack.
116macro_rules! format_param {
117    ($param:expr) => {{
118        let param = $param;
119        format!("{}{}", param.ty, if param.is_complex_type() { " memory" } else { "" })
120    }};
121}
122
123/// Helper function that formats solidity source with the given [FormatterConfig]
124pub fn format_source(source: &str, config: FormatterConfig) -> eyre::Result<String> {
125    match forge_fmt::parse(source) {
126        Ok(parsed) => {
127            let mut formatted_source = String::default();
128
129            if forge_fmt::format_to(&mut formatted_source, parsed, config).is_err() {
130                eyre::bail!("Could not format source!");
131            }
132
133            Ok(formatted_source)
134        }
135        Err(_) => eyre::bail!("Formatter could not parse source!"),
136    }
137}
138
139impl ChiselDispatcher {
140    /// Associated public function to create a new Dispatcher instance
141    pub fn new(config: SessionSourceConfig) -> eyre::Result<Self> {
142        ChiselSession::new(config).map(|session| Self { session })
143    }
144
145    /// Returns the optional ID of the current session.
146    pub fn id(&self) -> Option<&str> {
147        self.session.id.as_deref()
148    }
149
150    /// Returns the [`SessionSource`].
151    pub fn source(&self) -> &SessionSource {
152        &self.session.session_source
153    }
154
155    /// Returns the [`SessionSource`].
156    pub fn source_mut(&mut self) -> &mut SessionSource {
157        &mut self.session.session_source
158    }
159
160    fn format_source(&self) -> eyre::Result<String> {
161        format_source(
162            &self.source().to_repl_source(),
163            self.source().config.foundry_config.fmt.clone(),
164        )
165    }
166
167    /// Returns the prompt based on the current status of the Dispatcher
168    pub fn get_prompt(&self) -> Cow<'static, str> {
169        match self.session.id.as_deref() {
170            // `(ID: {id}) ➜ `
171            Some(id) => {
172                let mut prompt = String::with_capacity(DEFAULT_PROMPT.len() + id.len() + 7);
173                prompt.push_str("(ID: ");
174                prompt.push_str(id);
175                prompt.push_str(") ");
176                prompt.push_str(DEFAULT_PROMPT);
177                Cow::Owned(prompt)
178            }
179            // `➜ `
180            None => Cow::Borrowed(DEFAULT_PROMPT),
181        }
182    }
183
184    /// Dispatches a [ChiselCommand]
185    ///
186    /// ### Takes
187    ///
188    /// - A [ChiselCommand]
189    /// - An array of arguments
190    ///
191    /// ### Returns
192    ///
193    /// A [DispatchResult] containing feedback on the dispatch's execution.
194    pub async fn dispatch_command(&mut self, cmd: ChiselCommand, args: &[&str]) -> DispatchResult {
195        match cmd {
196            ChiselCommand::Help => {
197                let all_descriptors =
198                    ChiselCommand::iter().map(CmdDescriptor::from).collect::<Vec<CmdDescriptor>>();
199                DispatchResult::CommandSuccess(Some(format!(
200                    "{}\n{}",
201                    format!("{CHISEL_CHAR} Chisel help\n=============").cyan(),
202                    CmdCategory::iter()
203                        .map(|cat| {
204                            // Get commands in the current category
205                            let cat_cmds = &all_descriptors
206                                .iter()
207                                .filter(|(_, _, c)| {
208                                    std::mem::discriminant(c) == std::mem::discriminant(&cat)
209                                })
210                                .collect::<Vec<&CmdDescriptor>>();
211
212                            // Format the help menu for the current category
213                            format!(
214                                "{}\n{}\n",
215                                cat.magenta(),
216                                cat_cmds
217                                    .iter()
218                                    .map(|(cmds, desc, _)| format!(
219                                        "\t{} - {}",
220                                        cmds.iter()
221                                            .map(|cmd| format!("!{}", cmd.green()))
222                                            .collect::<Vec<_>>()
223                                            .join(" | "),
224                                        desc
225                                    ))
226                                    .collect::<Vec<String>>()
227                                    .join("\n")
228                            )
229                        })
230                        .collect::<Vec<String>>()
231                        .join("\n")
232                )))
233            }
234            ChiselCommand::Quit => {
235                // Exit the process with status code `0` for success.
236                std::process::exit(0);
237            }
238            ChiselCommand::Clear => {
239                self.source_mut().drain_run();
240                self.source_mut().drain_global_code();
241                self.source_mut().drain_top_level_code();
242                DispatchResult::CommandSuccess(Some(String::from("Cleared session!")))
243            }
244            ChiselCommand::Save => {
245                if args.len() <= 1 {
246                    // If a new name was supplied, overwrite the ID of the current session.
247                    if args.len() == 1 {
248                        // TODO: Should we delete the old cache file if the id of the session
249                        // changes?
250                        self.session.id = Some(args[0].to_owned());
251                    }
252
253                    if let Err(e) = self.session.write() {
254                        return DispatchResult::FileIoError(e.into())
255                    }
256                    DispatchResult::CommandSuccess(Some(format!(
257                        "Saved session to cache with ID = {}",
258                        self.session.id.as_ref().unwrap()
259                    )))
260                } else {
261                    DispatchResult::CommandFailed(Self::make_error(format!(
262                        "Too many arguments supplied: [{}]. Please check command syntax.",
263                        args.join(", ")
264                    )))
265                }
266            }
267            ChiselCommand::Load => {
268                if args.len() != 1 {
269                    // Must supply a session ID as the argument.
270                    return DispatchResult::CommandFailed(Self::make_error(
271                        "Must supply a session ID as the argument.",
272                    ))
273                }
274
275                // Use args as the name
276                let name = args[0];
277                // Try to save the current session before loading another
278                // Don't save an empty session
279                if !self.source().run_code.is_empty() {
280                    if let Err(e) = self.session.write() {
281                        return DispatchResult::FileIoError(e.into())
282                    }
283                    let _ = sh_println!("{}", "Saved current session!".green());
284                }
285
286                // Parse the arguments
287                let new_session = match name {
288                    "latest" => ChiselSession::latest(),
289                    _ => ChiselSession::load(name),
290                };
291
292                // WARNING: Overwrites the current session
293                if let Ok(mut new_session) = new_session {
294                    // Regenerate [IntermediateOutput]; It cannot be serialized.
295                    //
296                    // SAFETY
297                    // Should never panic due to the checks performed when the session was created
298                    // in the first place.
299                    new_session.session_source.build().unwrap();
300
301                    self.session = new_session;
302                    DispatchResult::CommandSuccess(Some(format!(
303                        "Loaded Chisel session! (ID = {})",
304                        self.session.id.as_ref().unwrap()
305                    )))
306                } else {
307                    DispatchResult::CommandFailed(Self::make_error("Failed to load session!"))
308                }
309            }
310            ChiselCommand::ListSessions => match ChiselSession::list_sessions() {
311                Ok(sessions) => DispatchResult::CommandSuccess(Some(format!(
312                    "{}\n{}",
313                    format!("{CHISEL_CHAR} Chisel Sessions").cyan(),
314                    sessions
315                        .iter()
316                        .map(|(time, name)| {
317                            format!("{} - {}", format!("{time:?}").blue(), name)
318                        })
319                        .collect::<Vec<String>>()
320                        .join("\n")
321                ))),
322                Err(_) => DispatchResult::CommandFailed(Self::make_error(
323                    "No sessions found. Use the `!save` command to save a session.",
324                )),
325            },
326            ChiselCommand::Source => match self.format_source() {
327                Ok(formatted_source) => DispatchResult::CommandSuccess(Some(
328                    SolidityHelper::new().highlight(&formatted_source).into_owned(),
329                )),
330                Err(_) => {
331                    DispatchResult::CommandFailed(String::from("Failed to format session source"))
332                }
333            },
334            ChiselCommand::ClearCache => match ChiselSession::clear_cache() {
335                Ok(_) => {
336                    self.session.id = None;
337                    DispatchResult::CommandSuccess(Some(String::from("Cleared chisel cache!")))
338                }
339                Err(_) => DispatchResult::CommandFailed(Self::make_error(
340                    "Failed to clear cache! Check file permissions or disk space.",
341                )),
342            },
343            ChiselCommand::Fork => {
344                if args.is_empty() || args[0].trim().is_empty() {
345                    self.source_mut().config.evm_opts.fork_url = None;
346                    return DispatchResult::CommandSuccess(Some(
347                        "Now using local environment.".to_string(),
348                    ))
349                }
350                if args.len() != 1 {
351                    return DispatchResult::CommandFailed(Self::make_error(
352                        "Must supply a session ID as the argument.",
353                    ))
354                }
355                let arg = *args.first().unwrap();
356
357                // If the argument is an RPC alias designated in the
358                // `[rpc_endpoints]` section of the `foundry.toml` within
359                // the pwd, use the URL matched to the key.
360                let endpoint = if let Some(endpoint) =
361                    self.source_mut().config.foundry_config.rpc_endpoints.get(arg)
362                {
363                    endpoint.clone()
364                } else {
365                    RpcEndpointUrl::Env(arg.to_string()).into()
366                };
367                let fork_url = match endpoint.resolve().url() {
368                    Ok(fork_url) => fork_url,
369                    Err(e) => {
370                        return DispatchResult::CommandFailed(Self::make_error(format!(
371                            "\"{}\" ENV Variable not set!",
372                            e.var
373                        )))
374                    }
375                };
376
377                // Check validity of URL
378                if Url::parse(&fork_url).is_err() {
379                    return DispatchResult::CommandFailed(Self::make_error(
380                        "Invalid fork URL! Please provide a valid RPC endpoint URL.",
381                    ))
382                }
383
384                // Create success message before moving the fork_url
385                let success_msg = format!("Set fork URL to {}", &fork_url.yellow());
386
387                // Update the fork_url inside of the [SessionSourceConfig]'s [EvmOpts]
388                // field
389                self.source_mut().config.evm_opts.fork_url = Some(fork_url);
390
391                // Clear the backend so that it is re-instantiated with the new fork
392                // upon the next execution of the session source.
393                self.source_mut().config.backend = None;
394
395                DispatchResult::CommandSuccess(Some(success_msg))
396            }
397            ChiselCommand::Traces => {
398                self.source_mut().config.traces = !self.source_mut().config.traces;
399                DispatchResult::CommandSuccess(Some(format!(
400                    "{} traces!",
401                    if self.source_mut().config.traces { "Enabled" } else { "Disabled" }
402                )))
403            }
404            ChiselCommand::Calldata => {
405                // remove empty space, double quotes, and 0x prefix
406                let arg = args
407                    .first()
408                    .map(|s| s.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\''))
409                    .map(|s| s.strip_prefix("0x").unwrap_or(s))
410                    .unwrap_or("");
411
412                if arg.is_empty() {
413                    self.source_mut().config.calldata = None;
414                    return DispatchResult::CommandSuccess(Some("Calldata cleared.".to_string()))
415                }
416
417                let calldata = hex::decode(arg);
418                match calldata {
419                    Ok(calldata) => {
420                        self.source_mut().config.calldata = Some(calldata);
421                        DispatchResult::CommandSuccess(Some(format!(
422                            "Set calldata to '{}'",
423                            arg.yellow()
424                        )))
425                    }
426                    Err(e) => DispatchResult::CommandFailed(Self::make_error(format!(
427                        "Invalid calldata: {e}"
428                    ))),
429                }
430            }
431            ChiselCommand::MemDump | ChiselCommand::StackDump => {
432                match self.source_mut().execute().await {
433                    Ok((_, res)) => {
434                        if let Some((stack, mem, _)) = res.state.as_ref() {
435                            if matches!(cmd, ChiselCommand::MemDump) {
436                                // Print memory by word
437                                (0..mem.len()).step_by(32).for_each(|i| {
438                                    let _ = sh_println!(
439                                        "{}: {}",
440                                        format!("[0x{:02x}:0x{:02x}]", i, i + 32).yellow(),
441                                        hex::encode_prefixed(&mem[i..i + 32]).cyan()
442                                    );
443                                });
444                            } else {
445                                // Print all stack items
446                                (0..stack.len()).rev().for_each(|i| {
447                                    let _ = sh_println!(
448                                        "{}: {}",
449                                        format!("[{}]", stack.len() - i - 1).yellow(),
450                                        format!("0x{:02x}", stack[i]).cyan()
451                                    );
452                                });
453                            }
454                            DispatchResult::CommandSuccess(None)
455                        } else {
456                            DispatchResult::CommandFailed(Self::make_error(
457                                "Run function is empty.",
458                            ))
459                        }
460                    }
461                    Err(e) => DispatchResult::CommandFailed(Self::make_error(e.to_string())),
462                }
463            }
464            ChiselCommand::Export => {
465                // Check if the current session inherits `Script.sol` before exporting
466
467                // Check if the pwd is a foundry project
468                if !Path::new("foundry.toml").exists() {
469                    return DispatchResult::CommandFailed(Self::make_error(
470                        "Must be in a foundry project to export source to script.",
471                    ));
472                }
473
474                // Create "script" dir if it does not already exist.
475                if !Path::new("script").exists() {
476                    if let Err(e) = std::fs::create_dir_all("script") {
477                        return DispatchResult::CommandFailed(Self::make_error(e.to_string()))
478                    }
479                }
480
481                match self.format_source() {
482                    Ok(formatted_source) => {
483                        // Write session source to `script/REPL.s.sol`
484                        if let Err(e) =
485                            std::fs::write(PathBuf::from("script/REPL.s.sol"), formatted_source)
486                        {
487                            return DispatchResult::CommandFailed(Self::make_error(e.to_string()))
488                        }
489
490                        DispatchResult::CommandSuccess(Some(String::from(
491                            "Exported session source to script/REPL.s.sol!",
492                        )))
493                    }
494                    Err(_) => DispatchResult::CommandFailed(String::from(
495                        "Failed to format session source",
496                    )),
497                }
498            }
499            ChiselCommand::Fetch => {
500                if args.len() != 2 {
501                    return DispatchResult::CommandFailed(Self::make_error(
502                        "Incorrect number of arguments supplied. Expected: <address> <name>",
503                    ))
504                }
505
506                let request_url = format!(
507                    "https://api.etherscan.io/api?module=contract&action=getabi&address={}{}",
508                    args[0],
509                    if let Some(api_key) =
510                        self.source().config.foundry_config.etherscan_api_key.as_ref()
511                    {
512                        format!("&apikey={api_key}")
513                    } else {
514                        String::default()
515                    }
516                );
517
518                // TODO: Not the cleanest method of building a solidity interface from
519                // the ABI, but does the trick. Might want to pull this logic elsewhere
520                // and/or refactor at some point.
521                match reqwest::get(&request_url).await {
522                    Ok(response) => {
523                        let json = response.json::<EtherscanABIResponse>().await.unwrap();
524                        if json.status == "1" && json.result.is_some() {
525                            let abi = json.result.unwrap();
526                            let abi: serde_json::Result<JsonAbi> = serde_json::from_str(&abi);
527                            if let Ok(abi) = abi {
528                                let mut interface = format!(
529                                    "// Interface of {}\ninterface {} {{\n",
530                                    args[0], args[1]
531                                );
532
533                                // Add error definitions
534                                abi.errors().for_each(|err| {
535                                    interface.push_str(&format!(
536                                        "\terror {}({});\n",
537                                        err.name,
538                                        err.inputs
539                                            .iter()
540                                            .map(|input| {
541                                                let mut param_type = &input.ty;
542                                                // If complex type then add the name of custom type.
543                                                // see <https://github.com/foundry-rs/foundry/issues/6618>.
544                                                if input.is_complex_type() {
545                                                    if let Some(
546                                                        InternalType::Enum { contract: _, ty } |
547                                                        InternalType::Struct { contract: _, ty } |
548                                                        InternalType::Other { contract: _, ty },
549                                                    ) = &input.internal_type
550                                                    {
551                                                        param_type = ty;
552                                                    }
553                                                }
554                                                format!("{} {}", param_type, input.name)
555                                            })
556                                            .collect::<Vec<_>>()
557                                            .join(",")
558                                    ));
559                                });
560                                // Add event definitions
561                                abi.events().for_each(|event| {
562                                    interface.push_str(&format!(
563                                        "\tevent {}({});\n",
564                                        event.name,
565                                        event
566                                            .inputs
567                                            .iter()
568                                            .map(|input| {
569                                                let mut formatted = input.ty.to_string();
570                                                if input.indexed {
571                                                    formatted.push_str(" indexed");
572                                                }
573                                                formatted
574                                            })
575                                            .collect::<Vec<_>>()
576                                            .join(",")
577                                    ));
578                                });
579                                // Add function definitions
580                                abi.functions().for_each(|func| {
581                                    interface.push_str(&format!(
582                                        "\tfunction {}({}) external{}{};\n",
583                                        func.name,
584                                        func.inputs
585                                            .iter()
586                                            .map(|input| format_param!(input))
587                                            .collect::<Vec<_>>()
588                                            .join(","),
589                                        match func.state_mutability {
590                                            alloy_json_abi::StateMutability::Pure => " pure",
591                                            alloy_json_abi::StateMutability::View => " view",
592                                            alloy_json_abi::StateMutability::Payable => " payable",
593                                            _ => "",
594                                        },
595                                        if func.outputs.is_empty() {
596                                            String::default()
597                                        } else {
598                                            format!(
599                                                " returns ({})",
600                                                func.outputs
601                                                    .iter()
602                                                    .map(|output| format_param!(output))
603                                                    .collect::<Vec<_>>()
604                                                    .join(",")
605                                            )
606                                        }
607                                    ));
608                                });
609                                // Close interface definition
610                                interface.push('}');
611
612                                // Add the interface to the source outright - no need to verify
613                                // syntax via compilation and/or
614                                // parsing.
615                                self.source_mut().with_global_code(&interface);
616
617                                DispatchResult::CommandSuccess(Some(format!(
618                                    "Added {}'s interface to source as `{}`",
619                                    args[0], args[1]
620                                )))
621                            } else {
622                                DispatchResult::CommandFailed(Self::make_error(
623                                    "Contract is not verified!",
624                                ))
625                            }
626                        } else if let Some(error_msg) = json.result {
627                            DispatchResult::CommandFailed(Self::make_error(format!(
628                                "Could not fetch interface - \"{error_msg}\""
629                            )))
630                        } else {
631                            DispatchResult::CommandFailed(Self::make_error(format!(
632                                "Could not fetch interface - \"{}\"",
633                                json.message
634                            )))
635                        }
636                    }
637                    Err(e) => DispatchResult::CommandFailed(Self::make_error(format!(
638                        "Failed to communicate with Etherscan API: {e}"
639                    ))),
640                }
641            }
642            ChiselCommand::Exec => {
643                if args.is_empty() {
644                    return DispatchResult::CommandFailed(Self::make_error(
645                        "No command supplied! Please provide a valid command after '!'.",
646                    ))
647                }
648
649                let mut cmd = Command::new(args[0]);
650                if args.len() > 1 {
651                    cmd.args(args[1..].iter().copied());
652                }
653
654                match cmd.output() {
655                    Ok(output) => {
656                        std::io::stdout().write_all(&output.stdout).unwrap();
657                        std::io::stdout().write_all(&output.stderr).unwrap();
658                        DispatchResult::CommandSuccess(None)
659                    }
660                    Err(e) => DispatchResult::CommandFailed(e.to_string()),
661                }
662            }
663            ChiselCommand::Edit => {
664                // create a temp file with the content of the run code
665                let mut temp_file_path = std::env::temp_dir();
666                temp_file_path.push("chisel-tmp.sol");
667                let result = std::fs::File::create(&temp_file_path)
668                    .map(|mut file| file.write_all(self.source().run_code.as_bytes()));
669                if let Err(e) = result {
670                    return DispatchResult::CommandFailed(format!(
671                        "Could not write to a temporary file: {e}"
672                    ))
673                }
674
675                // open the temp file with the editor
676                let editor = std::env::var("EDITOR").unwrap_or_else(|_| "vim".to_string());
677                let mut cmd = Command::new(editor);
678                cmd.arg(&temp_file_path);
679
680                match cmd.status() {
681                    Ok(status) => {
682                        if !status.success() {
683                            if let Some(status_code) = status.code() {
684                                return DispatchResult::CommandFailed(format!(
685                                    "Editor exited with status {status_code}"
686                                ))
687                            } else {
688                                return DispatchResult::CommandFailed(
689                                    "Editor exited without a status code".to_string(),
690                                )
691                            }
692                        }
693                    }
694                    Err(_) => {
695                        return DispatchResult::CommandFailed(
696                            "Editor exited without a status code".to_string(),
697                        )
698                    }
699                }
700
701                let mut new_session_source = self.source().clone();
702                if let Ok(edited_code) = std::fs::read_to_string(temp_file_path) {
703                    new_session_source.drain_run();
704                    new_session_source.with_run_code(&edited_code);
705                } else {
706                    return DispatchResult::CommandFailed(
707                        "Could not read the edited file".to_string(),
708                    )
709                }
710
711                // if the editor exited successfully, try to compile the new code
712                match new_session_source.execute().await {
713                    Ok((_, mut res)) => {
714                        let failed = !res.success;
715                        if new_session_source.config.traces || failed {
716                            if let Ok(decoder) =
717                                Self::decode_traces(&new_session_source.config, &mut res).await
718                            {
719                                if let Err(e) = Self::show_traces(&decoder, &mut res).await {
720                                    return DispatchResult::CommandFailed(e.to_string())
721                                };
722
723                                // Show console logs, if there are any
724                                let decoded_logs = decode_console_logs(&res.logs);
725                                if !decoded_logs.is_empty() {
726                                    let _ = sh_println!("{}", "Logs:".green());
727                                    for log in decoded_logs {
728                                        let _ = sh_println!("  {log}");
729                                    }
730                                }
731                            }
732
733                            if failed {
734                                // If the contract execution failed, continue on without
735                                // updating the source.
736                                return DispatchResult::CommandFailed(Self::make_error(
737                                    "Failed to execute edited contract!",
738                                ));
739                            }
740                        }
741
742                        // the code could be compiled, save it
743                        *self.source_mut() = new_session_source;
744                        DispatchResult::CommandSuccess(Some(String::from(
745                            "Successfully edited `run()` function's body!",
746                        )))
747                    }
748                    Err(_) => {
749                        DispatchResult::CommandFailed("The code could not be compiled".to_string())
750                    }
751                }
752            }
753            ChiselCommand::RawStack => {
754                let len = args.len();
755                if len != 1 {
756                    let msg = match len {
757                        0 => "No variable supplied!",
758                        _ => "!rawstack only takes one argument.",
759                    };
760                    return DispatchResult::CommandFailed(Self::make_error(msg))
761                }
762
763                // Store the variable that we want to inspect
764                let to_inspect = args.first().unwrap();
765
766                // Get a mutable reference to the session source
767                let source = self.source_mut();
768
769                // Copy the variable's stack contents into a bytes32 variable without updating
770                // the current session source.
771                let line = format!("bytes32 __raw__; assembly {{ __raw__ := {to_inspect} }}");
772                if let Ok((new_source, _)) = source.clone_with_new_line(line) {
773                    match new_source.inspect("__raw__").await {
774                        Ok((_, Some(res))) => return DispatchResult::CommandSuccess(Some(res)),
775                        Ok((_, None)) => {}
776                        Err(e) => return DispatchResult::CommandFailed(Self::make_error(e)),
777                    }
778                }
779
780                DispatchResult::CommandFailed(
781                    "Variable must exist within `run()` function.".to_string(),
782                )
783            }
784        }
785    }
786
787    /// Dispatches an input as a command via [Self::dispatch_command] or as a Solidity snippet.
788    pub async fn dispatch(&mut self, mut input: &str) -> DispatchResult {
789        // Check if the input is a builtin command.
790        // Commands are denoted with a `!` leading character.
791        if input.starts_with(COMMAND_LEADER) {
792            let split: Vec<&str> = input.split_whitespace().collect();
793            let raw_cmd = &split[0][1..];
794
795            return match raw_cmd.parse::<ChiselCommand>() {
796                Ok(cmd) => self.dispatch_command(cmd, &split[1..]).await,
797                Err(e) => DispatchResult::UnrecognizedCommand(e),
798            }
799        }
800        if input.trim().is_empty() {
801            debug!("empty dispatch input");
802            return DispatchResult::Success(None)
803        }
804
805        // Get a mutable reference to the session source
806        let source = self.source_mut();
807
808        // If the input is a comment, add it to the run code so we avoid running with empty input
809        if COMMENT_RE.is_match(input) {
810            debug!(%input, "matched comment");
811            source.with_run_code(input);
812            return DispatchResult::Success(None)
813        }
814
815        // If there is an address (or multiple addresses) in the input, ensure that they are
816        // encoded with a valid checksum per EIP-55.
817        let mut heap_input = input.to_string();
818        ADDRESS_RE.captures_iter(input).for_each(|m| {
819            // Convert the match to a string slice
820            let match_str = m.name("address").expect("exists").as_str();
821
822            // We can always safely unwrap here due to the regex matching.
823            let addr: Address = match_str.parse().expect("Valid address regex");
824            // Replace all occurrences of the address with a checksummed version
825            heap_input = heap_input.replace(match_str, &addr.to_string());
826        });
827        // Replace the old input with the formatted input.
828        input = &heap_input;
829
830        // Create new source with exact input appended and parse
831        let (mut new_source, do_execute) = match source.clone_with_new_line(input.to_string()) {
832            Ok(new) => new,
833            Err(e) => {
834                return DispatchResult::CommandFailed(Self::make_error(format!(
835                    "Failed to parse input! {e}"
836                )))
837            }
838        };
839
840        // TODO: Cloning / parsing the session source twice on non-inspected inputs kinda sucks.
841        // Should change up how this works.
842        match source.inspect(input).await {
843            // Continue and print
844            Ok((true, Some(res))) => {
845                let _ = sh_println!("{res}");
846            }
847            Ok((true, None)) => {}
848            // Return successfully
849            Ok((false, res)) => {
850                debug!(%input, ?res, "inspect success");
851                return DispatchResult::Success(res)
852            }
853
854            // Return with the error
855            Err(e) => return DispatchResult::CommandFailed(Self::make_error(e)),
856        }
857
858        if do_execute {
859            match new_source.execute().await {
860                Ok((_, mut res)) => {
861                    let failed = !res.success;
862
863                    // If traces are enabled or there was an error in execution, show the execution
864                    // traces.
865                    if new_source.config.traces || failed {
866                        if let Ok(decoder) = Self::decode_traces(&new_source.config, &mut res).await
867                        {
868                            if let Err(e) = Self::show_traces(&decoder, &mut res).await {
869                                return DispatchResult::CommandFailed(e.to_string())
870                            };
871
872                            // Show console logs, if there are any
873                            let decoded_logs = decode_console_logs(&res.logs);
874                            if !decoded_logs.is_empty() {
875                                let _ = sh_println!("{}", "Logs:".green());
876                                for log in decoded_logs {
877                                    let _ = sh_println!("  {log}");
878                                }
879                            }
880
881                            // If the contract execution failed, continue on without adding the new
882                            // line to the source.
883                            if failed {
884                                return DispatchResult::Failure(Some(Self::make_error(
885                                    "Failed to execute REPL contract!",
886                                )))
887                            }
888                        }
889                    }
890
891                    // Replace the old session source with the new version
892                    *self.source_mut() = new_source;
893
894                    DispatchResult::Success(None)
895                }
896                Err(e) => DispatchResult::Failure(Some(e.to_string())),
897            }
898        } else {
899            match new_source.build() {
900                Ok(out) => {
901                    debug!(%input, ?out, "skipped execute and rebuild source");
902                    *self.source_mut() = new_source;
903                    DispatchResult::Success(None)
904                }
905                Err(e) => DispatchResult::Failure(Some(e.to_string())),
906            }
907        }
908    }
909
910    /// Decodes traces in the [ChiselResult]
911    /// TODO: Add `known_contracts` back in.
912    ///
913    /// ### Takes
914    ///
915    /// - A reference to a [SessionSourceConfig]
916    /// - A mutable reference to a [ChiselResult]
917    ///
918    /// ### Returns
919    ///
920    /// Optionally, a [CallTraceDecoder]
921    pub async fn decode_traces(
922        session_config: &SessionSourceConfig,
923        result: &mut ChiselResult,
924        // known_contracts: &ContractsByArtifact,
925    ) -> eyre::Result<CallTraceDecoder> {
926        let mut decoder = CallTraceDecoderBuilder::new()
927            .with_labels(result.labeled_addresses.clone())
928            .with_signature_identifier(SignaturesIdentifier::new(
929                Config::foundry_cache_dir(),
930                session_config.foundry_config.offline,
931            )?)
932            .build();
933
934        let mut identifier = TraceIdentifiers::new().with_etherscan(
935            &session_config.foundry_config,
936            session_config.evm_opts.get_remote_chain_id().await,
937        )?;
938        if !identifier.is_empty() {
939            for (_, trace) in &mut result.traces {
940                decoder.identify(trace, &mut identifier);
941            }
942        }
943        Ok(decoder)
944    }
945
946    /// Display the gathered traces of a REPL execution.
947    ///
948    /// ### Takes
949    ///
950    /// - A reference to a [CallTraceDecoder]
951    /// - A mutable reference to a [ChiselResult]
952    ///
953    /// ### Returns
954    ///
955    /// Optionally, a unit type signifying a successful result.
956    pub async fn show_traces(
957        decoder: &CallTraceDecoder,
958        result: &mut ChiselResult,
959    ) -> eyre::Result<()> {
960        if result.traces.is_empty() {
961            eyre::bail!("Unexpected error: No traces gathered. Please report this as a bug: https://github.com/foundry-rs/foundry/issues/new?assignees=&labels=T-bug&template=BUG-FORM.yml");
962        }
963
964        sh_println!("{}", "Traces:".green())?;
965        for (kind, trace) in &mut result.traces {
966            // Display all Setup + Execution traces.
967            if matches!(kind, TraceKind::Setup | TraceKind::Execution) {
968                decode_trace_arena(trace, decoder).await?;
969                sh_println!("{}", render_trace_arena(trace))?;
970            }
971        }
972
973        Ok(())
974    }
975
976    /// Format a type that implements [std::fmt::Display] as a chisel error string.
977    ///
978    /// ### Takes
979    ///
980    /// A generic type implementing the [std::fmt::Display] trait.
981    ///
982    /// ### Returns
983    ///
984    /// A formatted error [String].
985    pub fn make_error<T: std::fmt::Display>(msg: T) -> String {
986        format!("{}", msg.red())
987    }
988}
989
990#[cfg(test)]
991mod tests {
992    use super::*;
993
994    #[test]
995    fn test_comment_regex() {
996        assert!(COMMENT_RE.is_match("// line comment"));
997        assert!(COMMENT_RE.is_match("  \n// line \tcomment\n"));
998        assert!(!COMMENT_RE.is_match("// line \ncomment"));
999
1000        assert!(COMMENT_RE.is_match("/* block comment */"));
1001        assert!(COMMENT_RE.is_match(" \t\n  /* block \n \t comment */\n"));
1002        assert!(!COMMENT_RE.is_match("/* block \n \t comment */\nwith \tother"));
1003    }
1004
1005    #[test]
1006    fn test_address_regex() {
1007        assert!(ADDRESS_RE.is_match("0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4"));
1008        assert!(ADDRESS_RE.is_match(" 0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4 "));
1009        assert!(ADDRESS_RE.is_match("0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4,"));
1010        assert!(ADDRESS_RE.is_match("(0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4)"));
1011        assert!(!ADDRESS_RE.is_match("0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4aaa"));
1012        assert!(!ADDRESS_RE.is_match("'0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4'"));
1013        assert!(!ADDRESS_RE.is_match("'    0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4'"));
1014        assert!(!ADDRESS_RE.is_match("'0xe5f3aF50FE5d0bF402a3C6F55ccC47d4307922d4'"));
1015    }
1016}
chisel/dispatcher.rs

chisel/
dispatcher.rs
