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

chisel/
dispatcher.rs
