chisel/
source.rs

1//! Session Source
2//!
3//! This module contains the `SessionSource` struct, which is a minimal wrapper around
4//! the REPL contract's source code. It provides simple compilation, parsing, and
5//! execution helpers.
6
7use eyre::Result;
8use foundry_compilers::{
9    Artifact, ProjectCompileOutput,
10    artifacts::{ConfigurableContractArtifact, Source, Sources},
11    project::ProjectCompiler,
12    solc::Solc,
13};
14use foundry_config::{Config, SolcReq};
15use foundry_evm::{
16    backend::Backend,
17    core::{bytecode::InstIter, evm::FoundryEvmNetwork},
18    opts::EvmOpts,
19};
20use semver::Version;
21use serde::{Deserialize, Serialize};
22use solar::{
23    ast::{ItemKind, StmtKind as AstStmtKind, yul},
24    interface::{Span, diagnostics::EmittedDiagnostics},
25    sema::{
26        CompilerRef,
27        hir::{Block, Contract, EventId, ItemId, Stmt, StmtKind as HirStmtKind},
28        ty::Gcx,
29    },
30};
31use std::{cell::OnceCell, fmt};
32use walkdir::WalkDir;
33
34/// The minimum Solidity version of the `Vm` interface.
35pub const MIN_VM_VERSION: Version = Version::new(0, 6, 2);
36
37/// Solidity source for the `Vm` interface in [forge-std](https://github.com/foundry-rs/forge-std)
38static VM_SOURCE: &str = include_str!("../../../testdata/utils/Vm.sol");
39
40/// [`SessionSource`] build output.
41pub struct GeneratedOutput {
42    output: ProjectCompileOutput,
43}
44
45impl fmt::Debug for GeneratedOutput {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        f.debug_struct("GeneratedOutput").finish_non_exhaustive()
48    }
49}
50
51impl GeneratedOutput {
52    /// Enters the solar compiler context, providing access to the HIR and `Gcx`.
53    pub fn enter<R: Send>(
54        &self,
55        f: impl for<'a, 'b, 'gcx> FnOnce(GeneratedOutputRef<'a, 'b, 'gcx>) -> R + Send,
56    ) -> R {
57        self.output
58            .parser()
59            .solc()
60            .compiler()
61            .enter(|c| f(GeneratedOutputRef { output: &self.output, compiler: c }))
62    }
63}
64
65/// A scoped reference to a [`GeneratedOutput`] together with an entered solar compiler.
66pub struct GeneratedOutputRef<'a, 'b, 'gcx> {
67    output: &'a ProjectCompileOutput,
68    pub(crate) compiler: &'b CompilerRef<'gcx>,
69}
70
71impl<'gcx> GeneratedOutputRef<'_, '_, 'gcx> {
72    pub fn gcx(&self) -> Gcx<'gcx> {
73        self.compiler.gcx()
74    }
75
76    pub fn repl_contract(&self) -> Option<&ConfigurableContractArtifact> {
77        self.output.find_first("REPL")
78    }
79
80    /// Looks up the REPL contract in the HIR.
81    pub fn repl_contract_hir(&self) -> Option<&'gcx Contract<'gcx>> {
82        self.gcx().hir.contracts().find(|c| c.name.as_str() == "REPL")
83    }
84
85    /// Returns the body block of the REPL `run()` function.
86    pub fn run_func_body(&self) -> Block<'gcx> {
87        let hir = &self.gcx().hir;
88        let c = self.repl_contract_hir().expect("REPL contract not found in HIR");
89        let f = c
90            .functions()
91            .find(|&f| hir.function(f).name.as_ref().map(|n| n.as_str()) == Some("run"))
92            .expect("`run()` function not found in REPL contract");
93        hir.function(f).body.expect("`run()` function does not have a body")
94    }
95
96    /// Returns the [`EventId`] of an event named `input` in the REPL contract, if any.
97    pub fn get_event(&self, input: &str) -> Option<EventId> {
98        let hir = &self.gcx().hir;
99        let c = self.repl_contract_hir()?;
100        c.items.iter().find_map(|id| {
101            if let ItemId::Event(eid) = id
102                && hir.event(*eid).name.as_str() == input
103            {
104                Some(*eid)
105            } else {
106                None
107            }
108        })
109    }
110
111    pub fn final_pc(&self, contract: &ConfigurableContractArtifact) -> Result<Option<usize>> {
112        let deployed_bytecode = contract
113            .get_deployed_bytecode()
114            .ok_or_else(|| eyre::eyre!("No deployed bytecode found for `REPL` contract"))?;
115        let deployed_bytecode_bytes = deployed_bytecode
116            .bytes()
117            .ok_or_else(|| eyre::eyre!("No deployed bytecode found for `REPL` contract"))?;
118
119        // Fetch the run function's body statement
120        let run_body = self.run_func_body();
121
122        // Record loc of first yul block return statement (if any).
123        // This is used to decide which is the final statement within the `run()` method.
124        // see <https://github.com/foundry-rs/foundry/issues/4617>.
125        //
126        // Walk the AST of the REPL source to find a top-level `return(...)` call
127        // inside any `assembly { ... }` block in `run()`. This lets us pick the
128        // meaningful Yul return span even when HIR represents the block coarsely.
129        let last_yul_return_span: Option<Span> = self.first_yul_return_span();
130
131        // Find the last statement within the "run()" method and get the program
132        // counter via the source map.
133        let Some(last_stmt) = run_body.last() else { return Ok(None) };
134
135        // If the final statement is some type of block (unchecked or regular),
136        // we need to find the final statement within that block. Otherwise, default to
137        // the source loc of the final statement of the `run()` function's block.
138        //
139        // Inline assembly blocks are handled separately via
140        // `trailing_assembly_last_stmt_span`, which walks the AST to recover the last
141        // meaningful Yul statement.
142        let source_stmt = match &last_stmt.kind {
143            HirStmtKind::UncheckedBlock(stmts) | HirStmtKind::Block(stmts) => {
144                if let Some(stmt) = stmts.last() {
145                    stmt
146                } else {
147                    // In the case where the block is empty, attempt to grab the statement
148                    // before the block. Because we use saturating sub to get the second to
149                    // last index, this can always be safely unwrapped.
150                    &run_body[run_body.len().saturating_sub(2)]
151                }
152            }
153            _ => last_stmt,
154        };
155        // If the trailing statement is an assembly block, prefer the last meaningful
156        // (non-`let`) Yul statement's span as the source location for `final_pc`.
157        // See <https://github.com/foundry-rs/foundry/issues/4938>.
158        //
159        // `trailing_assembly_last_stmt_span` verifies via the AST that the HIR node
160        // corresponds to an assembly block and supplies the concrete Yul span to use.
161        let mut source_span =
162            if matches!(last_stmt.kind, HirStmtKind::AssemblyBlock(_) | HirStmtKind::Err(_))
163                && let Some(span) = self.trailing_assembly_last_stmt_span()
164            {
165                span
166            } else {
167                self.stmt_span_without_semicolon(source_stmt)
168            };
169
170        // Consider yul return statement as final statement (if it's loc is lower).
171        if let Some(yul_return_span) = last_yul_return_span
172            && yul_return_span.hi() < source_span.lo()
173        {
174            source_span = yul_return_span;
175        }
176
177        // Map the source location of the final statement of the `run()` function to its
178        // corresponding runtime program counter
179        let result = self
180            .compiler
181            .sess()
182            .source_map()
183            .span_to_source(source_span)
184            .map_err(|e| eyre::eyre!("failed to resolve span: {e:?}"))?;
185        let range = result.data;
186        let offset = range.start as u32;
187        let length = range.len() as u32;
188        trace!(%offset, %length, "find pc");
189        let final_pc = contract
190            .get_source_map_deployed()
191            .ok_or_else(|| eyre::eyre!("No source map found for `REPL` contract"))??
192            .into_iter()
193            .zip(InstIter::new(deployed_bytecode_bytes).with_pc().map(|(pc, _)| pc))
194            .filter(|(s, _)| s.offset() == offset && s.length() == length)
195            .map(|(_, pc)| pc)
196            .max();
197        trace!(?final_pc);
198        Ok(final_pc)
199    }
200
201    /// Statements' ranges in the solc source map do not include the semicolon.
202    fn stmt_span_without_semicolon(&self, stmt: &Stmt<'_>) -> Span {
203        match stmt.kind {
204            HirStmtKind::DeclSingle(id) => {
205                let decl = self.gcx().hir.variable(id);
206                if let Some(expr) = decl.initializer {
207                    stmt.span.with_hi(expr.span.hi())
208                } else {
209                    stmt.span
210                }
211            }
212            HirStmtKind::DeclMulti(_, expr) => stmt.span.with_hi(expr.span.hi()),
213            HirStmtKind::Expr(expr) => expr.span,
214            _ => stmt.span,
215        }
216    }
217
218    /// Returns the AST `run()` body of the REPL contract, if any.
219    ///
220    /// Returns the AST `run()` body so inline assembly blocks can be inspected at
221    /// Yul-statement granularity.
222    fn repl_run_ast_body(&self) -> Option<&'gcx solar::ast::Block<'gcx>> {
223        let contract = self.repl_contract_hir()?;
224        let source = self.gcx().sources.get(contract.source)?;
225        let ast = source.ast.as_ref()?;
226
227        let contract_ast = ast.items.iter().find_map(|i| match &i.kind {
228            ItemKind::Contract(c) if c.name.as_str() == "REPL" => Some(c),
229            _ => None,
230        })?;
231        contract_ast.body.iter().find_map(|i| match &i.kind {
232            ItemKind::Function(f) if f.header.name.is_some_and(|n| n.as_str() == "run") => {
233                f.body.as_ref()
234            }
235            _ => None,
236        })
237    }
238
239    /// Returns the span of the first top-level `return(...)` call inside any
240    /// `assembly { ... }` block in the REPL `run()` function, if any.
241    fn first_yul_return_span(&self) -> Option<Span> {
242        let run_body = self.repl_run_ast_body()?;
243        for stmt in run_body.stmts.iter() {
244            let AstStmtKind::Assembly(asm) = &stmt.kind else { continue };
245            for ystmt in asm.block.stmts.iter() {
246                if let yul::StmtKind::Expr(e) = &ystmt.kind
247                    && let yul::ExprKind::Call(call) = &e.kind
248                    && call.name.as_str() == "return"
249                {
250                    return Some(ystmt.span);
251                }
252            }
253        }
254        None
255    }
256
257    /// If the last statement of the REPL `run()` function is an `assembly { ... }` block,
258    /// returns the span of its last non-`let` (i.e. non-VarDecl) Yul statement.
259    ///
260    /// This mirrors the legacy behavior used to pick a meaningful end-of-function PC when
261    /// the trailing statement is inline assembly.
262    fn trailing_assembly_last_stmt_span(&self) -> Option<Span> {
263        let run_body = self.repl_run_ast_body()?;
264        let AstStmtKind::Assembly(asm) = &run_body.stmts.last()?.kind else { return None };
265        asm.block
266            .stmts
267            .iter()
268            .rev()
269            .find(|s| !matches!(s.kind, yul::StmtKind::VarDecl(_, _)))
270            .map(|s| s.span)
271    }
272}
273
274/// Configuration for the [SessionSource]
275#[derive(Clone, Debug, Default, Serialize, Deserialize)]
276#[serde(bound = "")]
277pub struct SessionSourceConfig<FEN: FoundryEvmNetwork> {
278    /// Foundry configuration
279    pub foundry_config: Config,
280    /// EVM Options
281    pub evm_opts: EvmOpts,
282    /// Disable the default `Vm` import.
283    pub no_vm: bool,
284    /// In-memory REVM db for the session's runner.
285    #[serde(skip)]
286    pub backend: Option<Backend<FEN>>,
287    /// Optionally enable traces for the REPL contract execution
288    pub traces: bool,
289    /// Optionally set calldata for the REPL contract execution
290    pub calldata: Option<Vec<u8>>,
291    /// Enable viaIR with minimum optimization
292    ///
293    /// This can fix most of the "stack too deep" errors while resulting a
294    /// relatively accurate source map.
295    pub ir_minimum: bool,
296}
297
298impl<FEN: FoundryEvmNetwork> SessionSourceConfig<FEN> {
299    /// Detect the solc version to know if VM can be injected.
300    pub fn detect_solc(&mut self) -> Result<()> {
301        if self.foundry_config.solc.is_none() {
302            let version = Solc::ensure_installed(&"*".parse().unwrap())?;
303            self.foundry_config.solc = Some(SolcReq::Version(version));
304        }
305        if !self.no_vm
306            && let Some(version) = self.foundry_config.solc_version()
307            && version < MIN_VM_VERSION
308        {
309            info!(%version, minimum=%MIN_VM_VERSION, "Disabling VM injection");
310            self.no_vm = true;
311        }
312        Ok(())
313    }
314}
315
316/// REPL Session Source wrapper
317///
318/// Heavily based on soli's [`ConstructedSource`](https://github.com/jpopesculian/soli/blob/master/src/main.rs#L166)
319#[derive(Debug, Serialize, Deserialize)]
320#[serde(bound = "")]
321pub struct SessionSource<FEN: FoundryEvmNetwork> {
322    /// The file name
323    pub file_name: String,
324    /// The contract name
325    pub contract_name: String,
326
327    /// Session Source configuration
328    pub config: SessionSourceConfig<FEN>,
329
330    /// Global level Solidity code.
331    ///
332    /// Above and outside all contract declarations, in the global context.
333    pub global_code: String,
334    /// Top level Solidity code.
335    ///
336    /// Within the contract declaration, but outside of the `run()` function.
337    pub contract_code: String,
338    /// The code to be executed in the `run()` function.
339    pub run_code: String,
340
341    /// Cached VM source code.
342    #[serde(skip, default = "vm_source")]
343    vm_source: Source,
344    /// The generated output
345    #[serde(skip)]
346    output: OnceCell<GeneratedOutput>,
347}
348
349fn vm_source() -> Source {
350    Source::new(VM_SOURCE)
351}
352
353impl<FEN: FoundryEvmNetwork> Clone for SessionSource<FEN> {
354    fn clone(&self) -> Self {
355        Self {
356            file_name: self.file_name.clone(),
357            contract_name: self.contract_name.clone(),
358            global_code: self.global_code.clone(),
359            contract_code: self.contract_code.clone(),
360            run_code: self.run_code.clone(),
361            config: self.config.clone(),
362            vm_source: self.vm_source.clone(),
363            output: Default::default(),
364        }
365    }
366}
367
368impl<FEN: FoundryEvmNetwork> SessionSource<FEN> {
369    /// Creates a new source given a solidity compiler version
370    ///
371    /// # Panics
372    ///
373    /// If no Solc binary is set, cannot be found or the `--version` command fails
374    ///
375    /// ### Takes
376    ///
377    /// - An instance of [Solc]
378    /// - An instance of [SessionSourceConfig]
379    ///
380    /// ### Returns
381    ///
382    /// A new instance of [SessionSource]
383    pub fn new(mut config: SessionSourceConfig<FEN>) -> Result<Self> {
384        config.detect_solc()?;
385        Ok(Self {
386            file_name: "ReplContract.sol".to_string(),
387            contract_name: "REPL".to_string(),
388            config,
389            global_code: Default::default(),
390            contract_code: Default::default(),
391            run_code: Default::default(),
392            vm_source: vm_source(),
393            output: Default::default(),
394        })
395    }
396
397    /// Clones the [SessionSource] and appends a new line of code.
398    ///
399    /// Returns `true` if the new line was added to `run()`.
400    pub fn clone_with_new_line(&self, mut content: String) -> Result<(Self, bool)> {
401        if let Some((new_source, fragment)) = self
402            .parse_fragment(&content)
403            .or_else(|| {
404                content.push(';');
405                self.parse_fragment(&content)
406            })
407            .or_else(|| {
408                content = content.trim_end().trim_end_matches(';').to_string();
409                self.parse_fragment(&content)
410            })
411        {
412            Ok((new_source, matches!(fragment, ParseTreeFragment::Function)))
413        } else {
414            eyre::bail!("\"{}\"", content.trim());
415        }
416    }
417
418    /// Parses a fragment of Solidity code in memory and assigns it a scope within the
419    /// [`SessionSource`].
420    fn parse_fragment(&self, buffer: &str) -> Option<(Self, ParseTreeFragment)> {
421        #[track_caller]
422        fn debug_errors(errors: &EmittedDiagnostics) {
423            debug!("{errors}");
424        }
425
426        let mut this = self.clone();
427        match this.add_run_code(buffer).parse() {
428            Ok(()) => return Some((this, ParseTreeFragment::Function)),
429            Err(e) => debug_errors(&e),
430        }
431        this = self.clone();
432        match this.add_contract_code(buffer).parse() {
433            Ok(()) => return Some((this, ParseTreeFragment::Contract)),
434            Err(e) => debug_errors(&e),
435        }
436        this = self.clone();
437        match this.add_global_code(buffer).parse() {
438            Ok(()) => return Some((this, ParseTreeFragment::Source)),
439            Err(e) => debug_errors(&e),
440        }
441        None
442    }
443
444    /// Append global-level code to the source.
445    pub fn add_global_code(&mut self, content: &str) -> &mut Self {
446        self.global_code.push_str(content.trim());
447        self.global_code.push('\n');
448        self.clear_output();
449        self
450    }
451
452    /// Append contract-level code to the source.
453    pub fn add_contract_code(&mut self, content: &str) -> &mut Self {
454        self.contract_code.push_str(content.trim());
455        self.contract_code.push('\n');
456        self.clear_output();
457        self
458    }
459
460    /// Append code to the `run()` function of the REPL contract.
461    pub fn add_run_code(&mut self, content: &str) -> &mut Self {
462        self.run_code.push_str(content.trim());
463        self.run_code.push('\n');
464        self.clear_output();
465        self
466    }
467
468    /// Clears all source code.
469    pub fn clear(&mut self) {
470        String::clear(&mut self.global_code);
471        String::clear(&mut self.contract_code);
472        String::clear(&mut self.run_code);
473        self.clear_output();
474    }
475
476    /// Clear the `run()` function code.
477    pub fn clear_run(&mut self) -> &mut Self {
478        String::clear(&mut self.run_code);
479        self.clear_output();
480        self
481    }
482
483    fn clear_output(&mut self) {
484        self.output.take();
485    }
486
487    /// Compiles the source if necessary.
488    pub fn build(&self) -> Result<&GeneratedOutput> {
489        // TODO: mimics `get_or_try_init`
490        if let Some(output) = self.output.get() {
491            return Ok(output);
492        }
493        let output = self.compile()?;
494        let output = GeneratedOutput { output };
495        Ok(self.output.get_or_init(|| output))
496    }
497
498    /// Compiles the source.
499    #[cold]
500    fn compile(&self) -> Result<ProjectCompileOutput> {
501        let sources = self.get_sources();
502
503        let mut project = self.config.foundry_config.ephemeral_project()?;
504        self.config.foundry_config.disable_optimizations(&mut project, self.config.ir_minimum);
505        let mut output = ProjectCompiler::with_sources(&project, sources)?.compile()?;
506
507        if output.has_compiler_errors() {
508            eyre::bail!("{output}");
509        }
510
511        // Drive HIR lowering and analysis so that subsequent `enter` queries can use them.
512        // Chisel inspects expression values, so enable Solar's expression type table.
513        let compiler = output.parser_mut().solc_mut().compiler_mut();
514        compiler.sess_mut().opts.unstable.typeck = true;
515        compiler.enter_mut(|c| {
516            let _ = c.lower_asts();
517            let _ = c.analysis();
518        });
519
520        Ok(output)
521    }
522
523    fn get_sources(&self) -> Sources {
524        let mut sources = Sources::new();
525
526        let src = self.to_repl_source();
527        sources.insert(self.file_name.clone().into(), Source::new(src));
528
529        // Include Vm.sol if forge-std remapping is not available.
530        if !self.config.no_vm
531            && !self
532                .config
533                .foundry_config
534                .get_all_remappings()
535                .any(|r| r.name.starts_with("forge-std"))
536        {
537            sources.insert("forge-std/Vm.sol".into(), self.vm_source.clone());
538        }
539
540        sources
541    }
542
543    /// Construct the REPL source.
544    pub fn to_repl_source(&self) -> String {
545        let Self {
546            contract_name,
547            global_code,
548            contract_code: top_level_code,
549            run_code,
550            config,
551            ..
552        } = self;
553        let (mut vm_import, mut vm_constant) = (String::new(), String::new());
554        // Check if there's any `forge-std` remapping and determine proper path to it by
555        // searching remapping path.
556        if !config.no_vm
557            && let Some(remapping) = config
558                .foundry_config
559                .remappings
560                .iter()
561                .find(|remapping| remapping.name == "forge-std/")
562            && let Some(vm_path) = WalkDir::new(&remapping.path.path)
563                .into_iter()
564                .filter_map(|e| e.ok())
565                .find(|e| e.file_name() == "Vm.sol")
566        {
567            vm_import = format!(
568                "import {{Vm}} from \"{}\";\n",
569                vm_path.path().to_string_lossy().replace('\\', "/")
570            );
571            vm_constant = "Vm internal constant vm = Vm(address(uint160(uint256(keccak256(\"hevm cheat code\")))));\n".to_string();
572        }
573
574        format!(
575            r#"
576// SPDX-License-Identifier: UNLICENSED
577pragma solidity 0;
578
579{vm_import}
580{global_code}
581
582contract {contract_name} {{
583    {vm_constant}
584    {top_level_code}
585
586    /// @notice REPL contract entry point
587    function run() public {{
588        {run_code}
589    }}
590}}"#,
591        )
592    }
593
594    /// Parse the current source in memory using Solar.
595    pub(crate) fn parse(&self) -> Result<(), EmittedDiagnostics> {
596        let sess =
597            solar::interface::Session::builder().with_buffer_emitter(Default::default()).build();
598        let _ = sess.enter_sequential(|| -> solar::interface::Result<()> {
599            let arena = solar::ast::Arena::new();
600            let filename = self.file_name.clone().into();
601            let src = self.to_repl_source();
602            let mut parser = solar::parse::Parser::from_source_code(&sess, &arena, filename, src)?;
603            let _ast = parser.parse_file().map_err(|e| e.emit())?;
604            Ok(())
605        });
606        sess.dcx.emitted_errors().unwrap()
607    }
608}
609
610/// A Parse Tree Fragment
611///
612/// Used to determine whether an input will go to the "run()" function,
613/// the top level of the contract, or in global scope.
614#[derive(Debug)]
615enum ParseTreeFragment {
616    /// Code for the global scope
617    Source,
618    /// Code for the top level of the contract
619    Contract,
620    /// Code for the "run()" function
621    Function,
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627    use foundry_compilers::artifacts::remappings::{RelativeRemapping, RelativeRemappingPathBuf};
628    use foundry_evm::core::evm::EthEvmNetwork;
629    use std::fs;
630
631    /// Regression test for <https://github.com/foundry-rs/foundry/issues/14711>.
632    ///
633    /// `to_repl_source()` must use forward slashes in the Vm import path regardless of OS,
634    /// because Solidity import statements require `/` as the path separator.
635    #[test]
636    fn test_vm_import_path_uses_forward_slashes() {
637        let tmp = tempfile::tempdir().unwrap();
638        let vm_sol = tmp.path().join("Vm.sol");
639        fs::write(&vm_sol, "// dummy").unwrap();
640
641        let remapping = RelativeRemapping {
642            context: None,
643            name: "forge-std/".to_string(),
644            path: RelativeRemappingPathBuf { parent: None, path: tmp.path().to_path_buf() },
645        };
646
647        let mut config: SessionSourceConfig<EthEvmNetwork> = SessionSourceConfig {
648            foundry_config: Config {
649                solc: Some(SolcReq::Version(Version::new(0, 8, 29))),
650                remappings: vec![remapping],
651                ..Default::default()
652            },
653            ..Default::default()
654        };
655        // Pre-set solc so detect_solc() skips the ensure_installed I/O.
656        config.detect_solc().unwrap();
657
658        let source = SessionSource {
659            file_name: "ReplContract.sol".to_string(),
660            contract_name: "REPL".to_string(),
661            config,
662            global_code: Default::default(),
663            contract_code: Default::default(),
664            run_code: Default::default(),
665            vm_source: vm_source(),
666            output: Default::default(),
667        };
668
669        let repl = source.to_repl_source();
670        let import_line = repl.lines().find(|l| l.contains("import {Vm}")).unwrap();
671        assert!(
672            !import_line.contains('\\'),
673            "Vm import path must not contain backslashes, got: {import_line}"
674        );
675        assert!(import_line.contains('/'), "Vm import path must use forward slashes");
676    }
677}
chisel/source.rs

chisel/
source.rs
