1use 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
40pub const PROMPT_ARROW: char = '➜';
42pub const PROMPT_ARROW_STR: &str = "➜";
44const DEFAULT_PROMPT: &str = "➜ ";
45
46pub const COMMAND_LEADER: char = '!';
48pub const CHISEL_CHAR: &str = "⚒️";
50
51#[derive(Debug)]
53pub struct ChiselDispatcher<FEN: FoundryEvmNetwork> {
54 pub session: ChiselSession<FEN>,
55 pub helper: SolidityHelper,
56}
57
58pub 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 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 pub fn id(&self) -> Option<&str> {
73 self.session.id.as_deref()
74 }
75
76 pub const fn source(&self) -> &SessionSource<FEN> {
78 &self.session.source
79 }
80
81 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 pub fn get_prompt(&self) -> Cow<'static, str> {
95 match self.session.id.as_deref() {
96 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 None => Cow::Borrowed(DEFAULT_PROMPT),
107 }
108 }
109
110 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 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 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 pub async fn decode_traces(
161 session_config: &SessionSourceConfig<FEN>,
162 result: &mut ChiselResult,
163 ) -> 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 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 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 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 eyre::bail!("Failed to execute edited contract!");
232 }
233 }
234
235 *self.source_mut() = new_source;
237
238 Ok(())
239 }
240}
241
242impl<FEN: FoundryEvmNetwork> ChiselDispatcher<FEN> {
244 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 if let Some(id) = id {
287 self.session.id = Some(id);
289 }
290
291 self.session.write()?;
292 sh_println!("Saved session to cache with ID = {}", self.session.id.as_ref().unwrap())
293 }
294
295 pub(crate) fn load_session(&mut self, id: &str) -> Result<()> {
296 if !self.source().run_code.is_empty() {
299 self.session.write()?;
300 sh_println!("{}", "Saved current session!".green())?;
301 }
302
303 let new_session = match id {
304 "latest" => ChiselSession::<FEN>::latest(),
305 id => ChiselSession::<FEN>::load(id),
306 }
307 .wrap_err("failed to load session")?;
308
309 ensure_loaded_session_network_matches(
310 &self.session.source.config.foundry_config,
311 &new_session.source.config.foundry_config,
312 id,
313 )?;
314 new_session.source.build()?;
315 self.session = new_session;
316 sh_println!("Loaded Chisel session! (ID = {})", self.session.id.as_ref().unwrap())
317 }
318
319 pub(crate) fn list_sessions(&self) -> Result<()> {
320 let sessions = ChiselSession::<FEN>::get_sessions()?;
321 if sessions.is_empty() {
322 eyre::bail!("No sessions found. Use the `!save` command to save a session.");
323 }
324 sh_println!(
325 "{}\n{}",
326 format!("{CHISEL_CHAR} Chisel Sessions").cyan(),
327 sessions
328 .iter()
329 .map(|(time, name)| format!("{} - {}", format!("{time:?}").blue(), name))
330 .collect::<Vec<String>>()
331 .join("\n")
332 )
333 }
334
335 pub(crate) fn show_source(&self) -> Result<()> {
336 let formatted = self.format_source().wrap_err("failed to format session source")?;
337 let highlighted = self.helper.highlight(&formatted);
338 sh_println!("{highlighted}")
339 }
340
341 pub(crate) fn clear_cache(&mut self) -> Result<()> {
342 ChiselSession::<FEN>::clear_cache().wrap_err("failed to clear cache")?;
343 self.session.id = None;
344 sh_println!("Cleared chisel cache!")
345 }
346
347 pub(crate) fn set_fork(&mut self, url: Option<String>) -> Result<()> {
348 let Some(url) = url else {
349 self.source_mut().config.evm_opts.fork_url = None;
350 sh_println!("Now using local environment.")?;
351 return Ok(());
352 };
353
354 let endpoint = if let Some(endpoint) =
358 self.source_mut().config.foundry_config.rpc_endpoints.get(&url)
359 {
360 endpoint.clone()
361 } else {
362 RpcEndpointUrl::Env(url).into()
363 };
364 let fork_url = endpoint.resolve().url()?;
365
366 if let Err(e) = Url::parse(&fork_url) {
367 eyre::bail!("invalid fork URL: {e}");
368 }
369
370 sh_println!("Set fork URL to {}", fork_url.yellow())?;
371
372 self.source_mut().config.evm_opts.fork_url = Some(fork_url);
373 self.source_mut().config.backend = None;
376
377 Ok(())
378 }
379
380 pub(crate) fn toggle_traces(&mut self) -> Result<()> {
381 let t = &mut self.source_mut().config.traces;
382 *t = !*t;
383 sh_println!("{} traces!", if *t { "Enabled" } else { "Disabled" })
384 }
385
386 pub(crate) fn set_calldata(&mut self, data: Option<&str>) -> Result<()> {
387 let arg = data
389 .map(|s| s.trim_matches(|c: char| c.is_whitespace() || c == '"' || c == '\''))
390 .map(|s| s.strip_prefix("0x").unwrap_or(s))
391 .unwrap_or("");
392
393 if arg.is_empty() {
394 self.source_mut().config.calldata = None;
395 sh_println!("Calldata cleared.")?;
396 return Ok(());
397 }
398
399 let calldata = hex::decode(arg);
400 match calldata {
401 Ok(calldata) => {
402 self.source_mut().config.calldata = Some(calldata);
403 sh_println!("Set calldata to '{}'", arg.yellow())
404 }
405 Err(e) => {
406 eyre::bail!("Invalid calldata: {e}");
407 }
408 }
409 }
410
411 pub(crate) async fn show_mem_dump(&mut self) -> Result<()> {
412 let res = self.source_mut().execute().await?;
413 let Some((_, mem)) = res.state.as_ref() else {
414 eyre::bail!("Run function is empty.");
415 };
416 for i in (0..mem.len()).step_by(32) {
417 let _ = sh_println!(
418 "{}: {}",
419 format!("[0x{:02x}:0x{:02x}]", i, i + 32).yellow(),
420 hex::encode_prefixed(&mem[i..i + 32]).cyan()
421 );
422 }
423 Ok(())
424 }
425
426 pub(crate) async fn show_stack_dump(&mut self) -> Result<()> {
427 let res = self.source_mut().execute().await?;
428 let Some((stack, _)) = res.state.as_ref() else {
429 eyre::bail!("Run function is empty.");
430 };
431 for i in (0..stack.len()).rev() {
432 let _ = sh_println!(
433 "{}: {}",
434 format!("[{}]", stack.len() - i - 1).yellow(),
435 format!("0x{:02x}", stack[i]).cyan()
436 );
437 }
438 Ok(())
439 }
440
441 pub(crate) fn export(&self) -> Result<()> {
442 if !Path::new("foundry.toml").exists() {
444 eyre::bail!("Must be in a foundry project to export source to script.");
445 }
446
447 if !Path::new("script").exists() {
449 std::fs::create_dir_all("script")?;
450 }
451
452 let formatted_source = self.format_source()?;
453 std::fs::write(PathBuf::from("script/REPL.s.sol"), formatted_source)?;
454 sh_println!("Exported session source to script/REPL.s.sol!")
455 }
456
457 pub(crate) async fn fetch_interface(&mut self, address: Address, name: String) -> Result<()> {
459 let abis = fetch_abi_from_etherscan(address, &self.source().config.foundry_config)
460 .await
461 .wrap_err("Failed to fetch ABI from Etherscan")?;
462 let (abi, _) = abis
463 .into_iter()
464 .next()
465 .ok_or_else(|| eyre::eyre!("No ABI found for address {address} on Etherscan"))?;
466 let code = forge_fmt::format(&abi.to_sol(&name, None), FormatterConfig::default())
467 .into_result()?;
468 self.source_mut().add_global_code(&code);
469 sh_println!("Added {address}'s interface to source as `{name}`")
470 }
471
472 pub(crate) fn exec_command(&self, command: String, args: Vec<String>) -> Result<()> {
473 let mut cmd = Command::new(command);
474 cmd.args(args);
475 let _ = cmd.status()?;
476 Ok(())
477 }
478
479 pub(crate) async fn edit_session(&mut self) -> Result<()> {
480 let mut tmp = Builder::new()
482 .prefix("chisel-")
483 .suffix(".sol")
484 .tempfile()
485 .wrap_err("Could not create temporary file")?;
486 tmp.as_file_mut()
487 .write_all(self.source().run_code.as_bytes())
488 .wrap_err("Could not write to temporary file")?;
489
490 let editor = std::env::var("EDITOR").unwrap_or_else(|_| "vim".to_string());
492 let mut cmd = Command::new(editor);
493 cmd.arg(tmp.path());
494 let st = cmd.status()?;
495 if !st.success() {
496 eyre::bail!("Editor exited with {st}");
497 }
498
499 let edited_code = std::fs::read_to_string(tmp.path())?;
500 let mut new_source = self.source().clone();
501 new_source.clear_run();
502 new_source.add_run_code(&edited_code);
503
504 self.execute_and_replace(new_source).await?;
506 sh_println!("Successfully edited `run()` function's body!")
507 }
508
509 pub(crate) async fn show_raw_stack(&mut self, var: String) -> Result<()> {
510 let source = self.source_mut();
511 let line = format!("bytes32 __raw__; assembly {{ __raw__ := {var} }}");
512 if let Ok((new_source, _)) = source.clone_with_new_line(line)
513 && let (_, Some(res)) = new_source.inspect("__raw__").await?
514 {
515 sh_println!("{res}")?;
516 return Ok(());
517 }
518
519 eyre::bail!("Variable must exist within `run()` function.");
520 }
521}
522
523fn config_network_name(config: &Config) -> &'static str {
524 config.networks.active_network_name().unwrap_or("ethereum")
525}
526
527fn ensure_loaded_session_network_matches(
528 current: &Config,
529 loaded: &Config,
530 id: &str,
531) -> Result<()> {
532 let current_network = config_network_name(current);
533 let loaded_network = config_network_name(loaded);
534 if current_network != loaded_network {
535 eyre::bail!(
536 "Chisel session `{id}` was saved for network `{loaded_network}`, but the current \
537 network is `{current_network}`. Rerun with `--network {loaded_network}` to load it.",
538 );
539 }
540 Ok(())
541}
542
543fn preprocess(input: &str) -> (bool, Cow<'_, str>) {
546 let mut only_trivia = true;
547 let mut new_input = Cow::Borrowed(input);
548 for (pos, token) in solar::parse::Cursor::new(input).with_position() {
549 use RawTokenKind::{BlockComment, LineComment, Literal, Whitespace};
550
551 if matches!(token.kind, Whitespace | LineComment { .. } | BlockComment { .. }) {
552 continue;
553 }
554 only_trivia = false;
555
556 if let Literal { kind: RawLiteralKind::Int { base: Base::Hexadecimal, .. } } = token.kind
558 && token.len == 42
559 {
560 let range = pos..pos + 42;
561 if let Ok(addr) = input[range.clone()].parse::<Address>() {
562 new_input.to_mut().replace_range(range, addr.to_checksum_buffer(None).as_str());
563 }
564 }
565 }
566 (only_trivia, new_input)
567}
568
569#[cfg(test)]
570mod tests {
571 use super::*;
572
573 fn config_with_network(network: Option<&str>) -> Config {
574 let mut config = Config::default();
575 if let Some(network) = network {
576 config.networks = serde_json::from_value(serde_json::json!({
577 "network": network,
578 "celo": false,
579 "bypass_prevrandao": false,
580 }))
581 .unwrap();
582 }
583 config
584 }
585
586 #[test]
587 fn config_network_name_defaults_to_ethereum() {
588 assert_eq!(config_network_name(&Config::default()), "ethereum");
589 }
590
591 #[test]
592 fn ensure_loaded_session_network_matches_rejects_different_network() {
593 let current = config_with_network(None);
594 let loaded = config_with_network(Some("tempo"));
595
596 let err = ensure_loaded_session_network_matches(¤t, &loaded, "42").unwrap_err();
597 assert_eq!(
598 err.to_string(),
599 "Chisel session `42` was saved for network `tempo`, but the current network is \
600 `ethereum`. Rerun with `--network tempo` to load it."
601 );
602 }
603
604 #[test]
605 fn ensure_loaded_session_network_matches_accepts_same_network() {
606 let current = config_with_network(Some("tempo"));
607 let loaded = config_with_network(Some("tempo"));
608
609 ensure_loaded_session_network_matches(¤t, &loaded, "42").unwrap();
610 }
611
612 #[test]
613 fn test_trivia() {
614 fn only_trivia(s: &str) -> bool {
615 let (only_trivia, _new_input) = preprocess(s);
616 only_trivia
617 }
618 assert!(only_trivia("// line comment"));
619 assert!(only_trivia(" \n// line \tcomment\n"));
620 assert!(!only_trivia("// line \ncomment"));
621
622 assert!(only_trivia("/* block comment */"));
623 assert!(only_trivia(" \t\n /* block \n \t comment */\n"));
624 assert!(!only_trivia("/* block \n \t comment */\nwith \tother"));
625 }
626}