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

chisel/
dispatcher.rs
