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

chisel/
dispatcher.rs
