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::{ChiselCommand, ChiselResult, ChiselSession, SessionSourceConfig, SolidityHelper},
8    source::SessionSource,
9};
10use alloy_primitives::{Address, hex};
11use eyre::{Context, Result};
12use forge_fmt::FormatterConfig;
13use foundry_cli::utils::fetch_abi_from_etherscan;
14use foundry_config::RpcEndpointUrl;
15use foundry_evm::{
16    decode::decode_console_logs,
17    traces::{
18        CallTraceDecoder, CallTraceDecoderBuilder, TraceKind, decode_trace_arena,
19        identifier::{SignaturesIdentifier, TraceIdentifiers},
20        render_trace_arena,
21    },
22};
23use reqwest::Url;
24use solar::{
25    parse::lexer::token::{RawLiteralKind, RawTokenKind},
26    sema::ast::Base,
27};
28use std::{
29    borrow::Cow,
30    io::Write,
31    ops::ControlFlow,
32    path::{Path, PathBuf},
33    process::Command,
34};
35use tempfile::Builder;
36use yansi::Paint;
37
38/// Prompt arrow character.
39pub const PROMPT_ARROW: char = '➜';
40/// Prompt arrow string.
41pub const PROMPT_ARROW_STR: &str = "➜";
42const DEFAULT_PROMPT: &str = "➜ ";
43
44/// Command leader character
45pub const COMMAND_LEADER: char = '!';
46/// Chisel character
47pub const CHISEL_CHAR: &str = "⚒️";
48
49/// Chisel input dispatcher
50#[derive(Debug)]
51pub struct ChiselDispatcher {
52    pub session: ChiselSession,
53    pub helper: SolidityHelper,
54}
55
56/// Helper function that formats solidity source with the given [FormatterConfig]
57pub fn format_source(source: &str, config: FormatterConfig) -> eyre::Result<String> {
58    let formatted = forge_fmt::format(source, config).into_result()?;
59    Ok(formatted)
60}
61
62impl ChiselDispatcher {
63    /// Associated public function to create a new Dispatcher instance
64    pub fn new(config: SessionSourceConfig) -> eyre::Result<Self> {
65        let session = ChiselSession::new(config)?;
66        Ok(Self { session, helper: Default::default() })
67    }
68
69    /// Returns the optional ID of the current session.
70    pub fn id(&self) -> Option<&str> {
71        self.session.id.as_deref()
72    }
73
74    /// Returns the [`SessionSource`].
75    pub fn source(&self) -> &SessionSource {
76        &self.session.source
77    }
78
79    /// Returns the [`SessionSource`].
80    pub fn source_mut(&mut self) -> &mut SessionSource {
81        &mut self.session.source
82    }
83
84    fn format_source(&self) -> eyre::Result<String> {
85        format_source(
86            &self.source().to_repl_source(),
87            self.source().config.foundry_config.fmt.clone(),
88        )
89    }
90
91    /// Returns the prompt based on the current status of the Dispatcher
92    pub fn get_prompt(&self) -> Cow<'static, str> {
93        match self.session.id.as_deref() {
94            // `(ID: {id}) ➜ `
95            Some(id) => {
96                let mut prompt = String::with_capacity(DEFAULT_PROMPT.len() + id.len() + 7);
97                prompt.push_str("(ID: ");
98                prompt.push_str(id);
99                prompt.push_str(") ");
100                prompt.push_str(DEFAULT_PROMPT);
101                Cow::Owned(prompt)
102            }
103            // `➜ `
104            None => Cow::Borrowed(DEFAULT_PROMPT),
105        }
106    }
107
108    /// Dispatches an input as a command via [Self::dispatch_command] or as a Solidity snippet.
109    pub async fn dispatch(&mut self, mut input: &str) -> Result<ControlFlow<()>> {
110        if let Some(command) = input.strip_prefix(COMMAND_LEADER) {
111            return match ChiselCommand::parse(command) {
112                Ok(cmd) => self.dispatch_command(cmd).await,
113                Err(e) => eyre::bail!("unrecognized command: {e}"),
114            };
115        }
116
117        let source = self.source_mut();
118
119        input = input.trim();
120        let (only_trivia, new_input) = preprocess(input);
121        input = &*new_input;
122
123        // If the input is a comment, add it to the run code so we avoid running with empty input
124        if only_trivia {
125            debug!(?input, "matched trivia");
126            if !input.is_empty() {
127                source.add_run_code(input);
128            }
129            return Ok(ControlFlow::Continue(()));
130        }
131
132        // Create new source with exact input appended and parse
133        let (new_source, do_execute) = source.clone_with_new_line(input.to_string())?;
134
135        let (cf, res) = source.inspect(input).await?;
136        if let Some(res) = &res {
137            let _ = sh_println!("{res}");
138        }
139        if cf.is_break() {
140            debug!(%input, ?res, "inspect success");
141            return Ok(ControlFlow::Continue(()));
142        }
143
144        if do_execute {
145            self.execute_and_replace(new_source).await.map(ControlFlow::Continue)
146        } else {
147            let out = new_source.build()?;
148            debug!(%input, ?out, "skipped execute and rebuild source");
149            *self.source_mut() = new_source;
150            Ok(ControlFlow::Continue(()))
151        }
152    }
153
154    /// Decodes traces in the given [`ChiselResult`].
155    // TODO: Add `known_contracts` back in.
156    pub async fn decode_traces(
157        session_config: &SessionSourceConfig,
158        result: &mut ChiselResult,
159        // known_contracts: &ContractsByArtifact,
160    ) -> eyre::Result<CallTraceDecoder> {
161        let chain_id = session_config.evm_opts.get_remote_chain_id().await;
162
163        let mut decoder = CallTraceDecoderBuilder::new()
164            .with_labels(result.labeled_addresses.clone())
165            .with_signature_identifier(SignaturesIdentifier::from_config(
166                &session_config.foundry_config,
167            )?)
168            .with_chain_id(chain_id.map(|c| c.id()))
169            .build();
170
171        let mut identifier =
172            TraceIdentifiers::new().with_external(&session_config.foundry_config, chain_id)?;
173        if !identifier.is_empty() {
174            for (_, trace) in &mut result.traces {
175                decoder.identify(trace, &mut identifier);
176            }
177        }
178        Ok(decoder)
179    }
180
181    /// Display the gathered traces of a REPL execution.
182    pub async fn show_traces(
183        decoder: &CallTraceDecoder,
184        result: &mut ChiselResult,
185    ) -> eyre::Result<()> {
186        if result.traces.is_empty() {
187            return Ok(());
188        }
189
190        sh_println!("{}", "Traces:".green())?;
191        for (kind, trace) in &mut result.traces {
192            // Display all Setup + Execution traces.
193            if matches!(kind, TraceKind::Setup | TraceKind::Execution) {
194                decode_trace_arena(trace, decoder).await;
195                sh_println!("{}", render_trace_arena(trace))?;
196            }
197        }
198
199        Ok(())
200    }
201
202    async fn execute_and_replace(&mut self, mut new_source: SessionSource) -> Result<()> {
203        let mut res = new_source.execute().await?;
204        let failed = !res.success;
205        if new_source.config.traces || failed {
206            if let Ok(decoder) = Self::decode_traces(&new_source.config, &mut res).await {
207                Self::show_traces(&decoder, &mut res).await?;
208
209                // Show console logs, if there are any
210                let decoded_logs = decode_console_logs(&res.logs);
211                if !decoded_logs.is_empty() {
212                    let _ = sh_println!("{}", "Logs:".green());
213                    for log in decoded_logs {
214                        let _ = sh_println!("  {log}");
215                    }
216                }
217            }
218
219            if failed {
220                // If the contract execution failed, continue on without
221                // updating the source.
222                eyre::bail!("Failed to execute edited contract!");
223            }
224        }
225
226        // the code could be compiled, save it
227        *self.source_mut() = new_source;
228
229        Ok(())
230    }
231}
232
233/// [`ChiselCommand`] implementations.
234impl ChiselDispatcher {
235    /// Dispatches a [`ChiselCommand`].
236    pub async fn dispatch_command(&mut self, cmd: ChiselCommand) -> Result<ControlFlow<()>> {
237        match cmd {
238            ChiselCommand::Quit => Ok(ControlFlow::Break(())),
239            cmd => self.dispatch_command_impl(cmd).await.map(ControlFlow::Continue),
240        }
241    }
242
243    async fn dispatch_command_impl(&mut self, cmd: ChiselCommand) -> Result<()> {
244        match cmd {
245            ChiselCommand::Help => self.show_help(),
246            ChiselCommand::Quit => unreachable!(),
247            ChiselCommand::Clear => self.clear_source(),
248            ChiselCommand::Save { id } => self.save_session(id),
249            ChiselCommand::Load { id } => self.load_session(&id),
250            ChiselCommand::ListSessions => self.list_sessions(),
251            ChiselCommand::Source => self.show_source(),
252            ChiselCommand::ClearCache => self.clear_cache(),
253            ChiselCommand::Fork { url } => self.set_fork(url),
254            ChiselCommand::Traces => self.toggle_traces(),
255            ChiselCommand::Calldata { data } => self.set_calldata(data.as_deref()),
256            ChiselCommand::MemDump => self.show_mem_dump().await,
257            ChiselCommand::StackDump => self.show_stack_dump().await,
258            ChiselCommand::Export => self.export(),
259            ChiselCommand::Fetch { addr, name } => self.fetch_interface(addr, name).await,
260            ChiselCommand::Exec { command, args } => self.exec_command(command, args),
261            ChiselCommand::Edit => self.edit_session().await,
262            ChiselCommand::RawStack { var } => self.show_raw_stack(var).await,
263        }
264    }
265
266    pub(crate) fn show_help(&self) -> Result<()> {
267        sh_println!("{}", ChiselCommand::format_help())
268    }
269
270    pub(crate) fn clear_source(&mut self) -> Result<()> {
271        self.source_mut().clear();
272        sh_println!("Cleared session!")
273    }
274
275    pub(crate) fn save_session(&mut self, id: Option<String>) -> Result<()> {
276        // If a new name was supplied, overwrite the ID of the current session.
277        if let Some(id) = id {
278            // TODO: Should we delete the old cache file if the id of the session changes?
279            self.session.id = Some(id);
280        }
281
282        self.session.write()?;
283        sh_println!("Saved session to cache with ID = {}", self.session.id.as_ref().unwrap())
284    }
285
286    pub(crate) fn load_session(&mut self, id: &str) -> Result<()> {
287        // Try to save the current session before loading another.
288        // Don't save an empty session.
289        if !self.source().run_code.is_empty() {
290            self.session.write()?;
291            sh_println!("{}", "Saved current session!".green())?;
292        }
293
294        let new_session = match id {
295            "latest" => ChiselSession::latest(),
296            id => ChiselSession::load(id),
297        }
298        .wrap_err("failed to load session")?;
299
300        new_session.source.build()?;
301        self.session = new_session;
302        sh_println!("Loaded Chisel session! (ID = {})", self.session.id.as_ref().unwrap())
303    }
304
305    pub(crate) fn list_sessions(&self) -> Result<()> {
306        let sessions = ChiselSession::get_sessions()?;
307        if sessions.is_empty() {
308            eyre::bail!("No sessions found. Use the `!save` command to save a session.");
309        }
310        sh_println!(
311            "{}\n{}",
312            format!("{CHISEL_CHAR} Chisel Sessions").cyan(),
313            sessions
314                .iter()
315                .map(|(time, name)| format!("{} - {}", format!("{time:?}").blue(), name))
316                .collect::<Vec<String>>()
317                .join("\n")
318        )
319    }
320
321    pub(crate) fn show_source(&self) -> Result<()> {
322        let formatted = self.format_source().wrap_err("failed to format session source")?;
323        let highlighted = self.helper.highlight(&formatted);
324        sh_println!("{highlighted}")
325    }
326
327    pub(crate) fn clear_cache(&mut self) -> Result<()> {
328        ChiselSession::clear_cache().wrap_err("failed to clear cache")?;
329        self.session.id = None;
330        sh_println!("Cleared chisel cache!")
331    }
332
333    pub(crate) fn set_fork(&mut self, url: Option<String>) -> Result<()> {
334        let Some(url) = url else {
335            self.source_mut().config.evm_opts.fork_url = None;
336            sh_println!("Now using local environment.")?;
337            return Ok(());
338        };
339
340        // If the argument is an RPC alias designated in the
341        // `[rpc_endpoints]` section of the `foundry.toml` within
342        // the pwd, use the URL matched to the key.
343        let endpoint = if let Some(endpoint) =
344            self.source_mut().config.foundry_config.rpc_endpoints.get(&url)
345        {
346            endpoint.clone()
347        } else {
348            RpcEndpointUrl::Env(url).into()
349        };
350        let fork_url = endpoint.resolve().url()?;
351
352        if let Err(e) = Url::parse(&fork_url) {
353            eyre::bail!("invalid fork URL: {e}");
354        }
355
356        sh_println!("Set fork URL to {}", fork_url.yellow())?;
357
358        self.source_mut().config.evm_opts.fork_url = Some(fork_url);
359        // Clear the backend so that it is re-instantiated with the new fork
360        // upon the next execution of the session source.
361        self.source_mut().config.backend = None;
362
363        Ok(())
364    }
365
366    pub(crate) fn toggle_traces(&mut self) -> Result<()> {
367        let t = &mut self.source_mut().config.traces;
368        *t = !*t;
369        sh_println!("{} traces!", if *t { "Enabled" } else { "Disabled" })
370    }
371
372    pub(crate) fn set_calldata(&mut self, data: Option<&str>) -> Result<()> {
373        // remove empty space, double quotes, and 0x prefix
374        let arg = data
375            .map(|s| s.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\''))
376            .map(|s| s.strip_prefix("0x").unwrap_or(s))
377            .unwrap_or("");
378
379        if arg.is_empty() {
380            self.source_mut().config.calldata = None;
381            sh_println!("Calldata cleared.")?;
382            return Ok(());
383        }
384
385        let calldata = hex::decode(arg);
386        match calldata {
387            Ok(calldata) => {
388                self.source_mut().config.calldata = Some(calldata);
389                sh_println!("Set calldata to '{}'", arg.yellow())
390            }
391            Err(e) => {
392                eyre::bail!("Invalid calldata: {e}")
393            }
394        }
395    }
396
397    pub(crate) async fn show_mem_dump(&mut self) -> Result<()> {
398        let res = self.source_mut().execute().await?;
399        let Some((_, mem)) = res.state.as_ref() else {
400            eyre::bail!("Run function is empty.");
401        };
402        for i in (0..mem.len()).step_by(32) {
403            let _ = sh_println!(
404                "{}: {}",
405                format!("[0x{:02x}:0x{:02x}]", i, i + 32).yellow(),
406                hex::encode_prefixed(&mem[i..i + 32]).cyan()
407            );
408        }
409        Ok(())
410    }
411
412    pub(crate) async fn show_stack_dump(&mut self) -> Result<()> {
413        let res = self.source_mut().execute().await?;
414        let Some((stack, _)) = res.state.as_ref() else {
415            eyre::bail!("Run function is empty.");
416        };
417        for i in (0..stack.len()).rev() {
418            let _ = sh_println!(
419                "{}: {}",
420                format!("[{}]", stack.len() - i - 1).yellow(),
421                format!("0x{:02x}", stack[i]).cyan()
422            );
423        }
424        Ok(())
425    }
426
427    pub(crate) fn export(&self) -> Result<()> {
428        // Check if the pwd is a foundry project
429        if !Path::new("foundry.toml").exists() {
430            eyre::bail!("Must be in a foundry project to export source to script.");
431        }
432
433        // Create "script" dir if it does not already exist.
434        if !Path::new("script").exists() {
435            std::fs::create_dir_all("script")?;
436        }
437
438        let formatted_source = self.format_source()?;
439        std::fs::write(PathBuf::from("script/REPL.s.sol"), formatted_source)?;
440        sh_println!("Exported session source to script/REPL.s.sol!")
441    }
442
443    /// Fetches an interface from Etherscan
444    pub(crate) async fn fetch_interface(&mut self, address: Address, name: String) -> Result<()> {
445        let abis = fetch_abi_from_etherscan(address, &self.source().config.foundry_config)
446            .await
447            .wrap_err("Failed to fetch ABI from Etherscan")?;
448        let (abi, _) = abis
449            .into_iter()
450            .next()
451            .ok_or_else(|| eyre::eyre!("No ABI found for address {address} on Etherscan"))?;
452        let code = forge_fmt::format(&abi.to_sol(&name, None), FormatterConfig::default())
453            .into_result()?;
454        self.source_mut().add_global_code(&code);
455        sh_println!("Added {address}'s interface to source as `{name}`")
456    }
457
458    pub(crate) fn exec_command(&self, command: String, args: Vec<String>) -> Result<()> {
459        let mut cmd = Command::new(command);
460        cmd.args(args);
461        let _ = cmd.status()?;
462        Ok(())
463    }
464
465    pub(crate) async fn edit_session(&mut self) -> Result<()> {
466        // create a temp file with the content of the run code
467        let mut tmp = Builder::new()
468            .prefix("chisel-")
469            .suffix(".sol")
470            .tempfile()
471            .wrap_err("Could not create temporary file")?;
472        tmp.as_file_mut()
473            .write_all(self.source().run_code.as_bytes())
474            .wrap_err("Could not write to temporary file")?;
475
476        // open the temp file with the editor
477        let editor = std::env::var("EDITOR").unwrap_or_else(|_| "vim".to_string());
478        let mut cmd = Command::new(editor);
479        cmd.arg(tmp.path());
480        let st = cmd.status()?;
481        if !st.success() {
482            eyre::bail!("Editor exited with {st}");
483        }
484
485        let edited_code = std::fs::read_to_string(tmp.path())?;
486        let mut new_source = self.source().clone();
487        new_source.clear_run();
488        new_source.add_run_code(&edited_code);
489
490        // if the editor exited successfully, try to compile the new code
491        self.execute_and_replace(new_source).await?;
492        sh_println!("Successfully edited `run()` function's body!")
493    }
494
495    pub(crate) async fn show_raw_stack(&mut self, var: String) -> Result<()> {
496        let source = self.source_mut();
497        let line = format!("bytes32 __raw__; assembly {{ __raw__ := {var} }}");
498        if let Ok((new_source, _)) = source.clone_with_new_line(line)
499            && let (_, Some(res)) = new_source.inspect("__raw__").await?
500        {
501            sh_println!("{res}")?;
502            return Ok(());
503        }
504
505        eyre::bail!("Variable must exist within `run()` function.")
506    }
507}
508
509/// Preprocesses addresses to ensure they are correctly checksummed and returns whether the input
510/// only contained trivia (comments, whitespace).
511fn preprocess(input: &str) -> (bool, Cow<'_, str>) {
512    let mut only_trivia = true;
513    let mut new_input = Cow::Borrowed(input);
514    for (pos, token) in solar::parse::Cursor::new(input).with_position() {
515        use RawTokenKind::{BlockComment, LineComment, Literal, Whitespace};
516
517        if matches!(token.kind, Whitespace | LineComment { .. } | BlockComment { .. }) {
518            continue;
519        }
520        only_trivia = false;
521
522        // Ensure that addresses are correctly checksummed.
523        if let Literal { kind: RawLiteralKind::Int { base: Base::Hexadecimal, .. } } = token.kind
524            && token.len == 42
525        {
526            let range = pos..pos + 42;
527            if let Ok(addr) = input[range.clone()].parse::<Address>() {
528                new_input.to_mut().replace_range(range, addr.to_checksum_buffer(None).as_str());
529            }
530        }
531    }
532    (only_trivia, new_input)
533}
534
535#[cfg(test)]
536mod tests {
537    use super::*;
538
539    #[test]
540    fn test_trivia() {
541        fn only_trivia(s: &str) -> bool {
542            let (only_trivia, _new_input) = preprocess(s);
543            only_trivia
544        }
545        assert!(only_trivia("// line comment"));
546        assert!(only_trivia("  \n// line \tcomment\n"));
547        assert!(!only_trivia("// line \ncomment"));
548
549        assert!(only_trivia("/* block comment */"));
550        assert!(only_trivia(" \t\n  /* block \n \t comment */\n"));
551        assert!(!only_trivia("/* block \n \t comment */\nwith \tother"));
552    }
553}
chisel/dispatcher.rs

chisel/
dispatcher.rs
