chisel/
dispatcher.rs

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

chisel/
dispatcher.rs
