edb_engine/instrumentation/
modification.rs

1// EDB - Ethereum Debugger
2// Copyright (C) 2024 Zhuo Zhang and Wuqi Zhang
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU Affero General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8//
9// This program is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU Affero General Public License for more details.
13//
14// You should have received a copy of the GNU Affero General Public License
15// along with this program. If not, see <https://www.gnu.org/licenses/>.
16
17use std::{collections::BTreeMap, fmt::Display, sync::Arc};
18
19use crate::{
20    analysis::{stmt_src, SourceAnalysis, VariableRef},
21    find_index_of_first_statement_in_block, find_index_of_first_statement_in_block_or_statement,
22    find_next_index_of_last_statement_in_block, find_next_index_of_source_location,
23    find_next_index_of_statement,
24    instrumentation::codegen,
25    USID, UVID,
26};
27
28use eyre::Result;
29use foundry_compilers::artifacts::{ast::SourceLocation, BlockOrStatement, Statement};
30use semver::Version;
31
32const LEFT_BRACKET_PRIORITY: u8 = 255; // used for the left bracket of the block
33const FUNCTION_ENTRY_PRIORITY: u8 = 191; // used for the before step hook of function and modifier entry
34const VISIBILITY_PRIORITY: u8 = 128; // used for the visibility of state variables and functions
35const VARIABLE_UPDATE_PRIORITY: u8 = 127; // used for the variable update hook
36const BEFORE_STEP_PRIORITY: u8 = 63; // used for the before step hook of statements other than function and modifier entry.
37const RIGHT_BRACKET_PRIORITY: u8 = 0; // used for the right bracket of the block
38
39/// A reference to a version.
40pub type VersionRef = Arc<Version>;
41
42/// The collections of modifications on a source file.
43pub struct SourceModifications {
44    source_id: u32,
45    /// The modifications on the source file. The key is the location of modification in the original source code.
46    modifications: BTreeMap<usize, Modification>,
47}
48
49impl SourceModifications {
50    /// Creates a new source modifications.
51    pub fn new(source_id: u32) -> Self {
52        Self { source_id, modifications: BTreeMap::new() }
53    }
54
55    /// Adds a modification to the source modifications.
56    ///
57    /// # Panics
58    ///
59    /// Panics if the modification overlaps with the previous or next modification.
60    pub fn add_modification(&mut self, modification: Modification) {
61        assert_eq!(modification.source_id(), self.source_id, "modification source id mismatch");
62
63        let loc = modification.loc();
64        // Check if the modification overlaps with the previous modification
65        if let Some((immediate_prev_loc, immediate_prev)) =
66            self.modifications.range_mut(..loc).next_back()
67        {
68            assert!(
69                immediate_prev_loc + immediate_prev.modified_length() <= loc,
70                "modification location overlaps with previous modification"
71            );
72        }
73        // Check if the modification overlaps with the next modification
74        if let Some((immediate_next_loc, immediate_next)) =
75            self.modifications.range_mut(loc..).next()
76        {
77            assert!(
78                loc + modification.modified_length() <= *immediate_next_loc,
79                "modification location overlaps with next modification"
80            );
81            // if both of them are instrument actions and instrument at the same location, merge them. The later comming modification will be appended after the earlier one.
82            if immediate_next.is_instrument()
83                && modification.is_instrument()
84                && *immediate_next_loc == loc
85            {
86                immediate_next.modify_instrument_action(|act| {
87                    act.content = if act.priority >= modification.as_instrument_action().priority {
88                        InstrumentContent::Plain(format!(
89                            "{} {}",
90                            act.content,
91                            modification.as_instrument_action().content,
92                        ))
93                    } else {
94                        InstrumentContent::Plain(format!(
95                            "{} {}",
96                            modification.as_instrument_action().content,
97                            act.content,
98                        ))
99                    };
100                });
101                return;
102            }
103        }
104        // Insert the modification
105        self.modifications.insert(loc, modification);
106    }
107
108    /// Extends the modifications with the given modifications.
109    pub fn extend_modifications(&mut self, modifications: Vec<Modification>) {
110        for modification in modifications {
111            self.add_modification(modification);
112        }
113    }
114
115    /// Modifies the source code with the modifications.
116    pub fn modify_source(&self, source: &str) -> String {
117        let mut modified_source = source.to_string();
118        // Apply the modifications in reverse order to avoid index shifting
119        for (_, modification) in self.modifications.iter().rev() {
120            match modification {
121                Modification::Instrument(instrument_action) => {
122                    modified_source.insert_str(
123                        instrument_action.loc,
124                        instrument_action.content.to_string().as_str(),
125                    );
126                }
127                Modification::Remove(remove_action) => {
128                    modified_source.replace_range(remove_action.start()..remove_action.end(), "");
129                }
130            }
131        }
132        modified_source
133    }
134}
135
136/// The modifications on a source file.
137#[allow(clippy::large_enum_variant)]
138#[derive(Debug, Clone, derive_more::From)]
139pub enum Modification {
140    /// An action to instrument a code in the source file.
141    Instrument(#[from] InstrumentAction),
142    /// An action to remove a code in the source file.
143    Remove(#[from] RemoveAction),
144}
145
146impl Modification {
147    /// Gets the source ID of the modification.
148    pub fn source_id(&self) -> u32 {
149        match self {
150            Self::Instrument(instrument_action) => instrument_action.source_id,
151            Self::Remove(remove_action) => {
152                remove_action.src.index.expect("remove action index not found") as u32
153            }
154        }
155    }
156
157    /// Gets the location of the modification.
158    pub fn loc(&self) -> usize {
159        match self {
160            Self::Instrument(instrument_action) => instrument_action.loc,
161            Self::Remove(remove_action) => remove_action.src.start.unwrap_or(0),
162        }
163    }
164
165    /// Gets the length of the original code that is modified.
166    pub const fn modified_length(&self) -> usize {
167        match self {
168            Self::Instrument(_) => 0,
169            Self::Remove(remove_action) => {
170                remove_action.src.length.expect("remove action length not found")
171            }
172        }
173    }
174
175    /// Checks if the modification is an instrument action.
176    pub const fn is_instrument(&self) -> bool {
177        matches!(self, Self::Instrument(_))
178    }
179
180    /// Checks if the modification is a remove action.
181    pub const fn is_remove(&self) -> bool {
182        matches!(self, Self::Remove(_))
183    }
184
185    /// Gets the instrument action if it is an instrument action.
186    pub fn as_instrument_action(&self) -> &InstrumentAction {
187        match self {
188            Self::Instrument(instrument_action) => instrument_action,
189            Self::Remove(_) => panic!("cannot get instrument action from remove action"),
190        }
191    }
192
193    /// Gets the remove action if it is a remove action.
194    pub fn as_remove_action(&self) -> &RemoveAction {
195        match self {
196            Self::Instrument(_) => panic!("cannot get remove action from instrument action"),
197            Self::Remove(remove_action) => remove_action,
198        }
199    }
200
201    /// Modifies the remove action if it is a remove action.
202    pub fn modify_remove_action(&mut self, f: impl FnOnce(&mut RemoveAction)) {
203        match self {
204            Self::Instrument(_) => {}
205            Self::Remove(remove_action) => {
206                f(remove_action);
207            }
208        }
209    }
210
211    /// Modifies the instrument action if it is an instrument action.
212    pub fn modify_instrument_action(&mut self, f: impl FnOnce(&mut InstrumentAction)) {
213        match self {
214            Self::Instrument(instrument_action) => {
215                f(instrument_action);
216            }
217            Self::Remove(_) => {}
218        }
219    }
220}
221
222/// An action to instrument a code in the source file.
223#[derive(Debug, Clone)]
224pub struct InstrumentAction {
225    /// The source ID of the source file to instrument
226    pub source_id: u32,
227    /// The location of the code to instrument. This is the offset of the code at which the instrumented code should be inserted.
228    pub loc: usize,
229    /// The code to instrument
230    pub content: InstrumentContent,
231    /// The priority of the instrument action. If two `InstrumentAction`s have the same `loc`, the one with higher priority will be applied first.
232    pub priority: u8,
233}
234
235/// An action to remove a code in the source file.
236#[derive(Debug, Clone)]
237pub struct RemoveAction {
238    /// The source location of the code to remove
239    pub src: SourceLocation,
240}
241
242impl RemoveAction {
243    /// Gets the start index of the code to remove.
244    pub fn start(&self) -> usize {
245        self.src.start.expect("remove action start not found")
246    }
247
248    /// Gets the end index of the code to remove (exclusive).
249    pub fn end(&self) -> usize {
250        self.start() + self.src.length.expect("remove action length not found")
251    }
252}
253
254/// The content to instrument.
255#[derive(Debug, Clone)]
256pub enum InstrumentContent {
257    /// The code to instrument. The plain code can be directly inserted into the source code as a string.
258    Plain(String),
259    /// View method for state variables
260    ViewMethod {
261        /// The state variable being accessed.
262        variable: VariableRef,
263    },
264    /// A `before_step` hook. The debugger will pause here during step-by-step execution.
265    BeforeStepHook {
266        /// Compiler Version
267        version: VersionRef,
268        /// The USID of the step.
269        usid: USID,
270        /// The number of function calls made in the step.
271        function_calls: usize,
272    },
273    /// A `variable_update` hook. The debugger will record the value of the variable when it is updated.
274    VariableUpdateHook {
275        /// Compiler Version
276        version: VersionRef,
277        /// The UVID of the variable.
278        uvid: UVID,
279        /// The variable that is updated.
280        variable: VariableRef,
281    },
282}
283
284impl Display for InstrumentContent {
285    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
286        let content = match self {
287            Self::Plain(content) => content.clone(),
288            Self::ViewMethod { variable } => {
289                codegen::generate_view_method(variable).unwrap_or_default()
290            }
291            Self::BeforeStepHook { version, usid, .. } => {
292                codegen::generate_step_hook(version, *usid).unwrap_or_default()
293            }
294            Self::VariableUpdateHook { version, uvid, variable } => {
295                codegen::generate_variable_update_hook(version, *uvid, variable).unwrap_or_default()
296            }
297        };
298        write!(f, "{content}")
299    }
300}
301
302impl SourceModifications {
303    /// Collects the modifications on the source code given the analysis result.
304    pub fn collect_modifications(
305        &mut self,
306        compiler_version: VersionRef,
307        source: &str,
308        analysis: &SourceAnalysis,
309    ) -> Result<()> {
310        // Collect the modifications to generate view methods for state variables.
311        self.collect_view_method_modifications(analysis);
312
313        // Collect the modifications to patch single-statement if/for/while/try/catch/etc.
314        self.collect_statement_to_block_modifications(source, analysis)?;
315
316        // Collect the before step hook modifications for each step.
317        self.collect_before_step_hook_modifications(compiler_version.clone(), source, analysis)?;
318
319        // Collect the variable update hook modifications for each step.
320        self.collect_variable_update_hook_modifications(compiler_version, source, analysis)?;
321
322        Ok(())
323    }
324
325    /// Collects the modifications to generate view methods for state variables.
326    fn collect_view_method_modifications(&mut self, analysis: &SourceAnalysis) {
327        let source_id = self.source_id;
328        for state_variable in &analysis.state_variables {
329            let src = &state_variable.declaration().src;
330            let loc = src.start.unwrap_or(0) + src.length.unwrap_or(0) + 1; // XXX (ZZ): we may need to check last char
331            let instrument_action = InstrumentAction {
332                source_id,
333                loc,
334                content: InstrumentContent::ViewMethod { variable: state_variable.clone() },
335                priority: VISIBILITY_PRIORITY,
336            };
337            self.add_modification(instrument_action.into());
338        }
339    }
340
341    /// Collects the modifications to convert a statement to a block. Some control flow structures, such as if/for/while/try/catch/etc., may have their body as a single statement. We need to convert them to a block.
342    fn collect_statement_to_block_modifications(
343        &mut self,
344        source: &str,
345        analysis: &SourceAnalysis,
346    ) -> Result<()> {
347        let source_id = self.source_id;
348
349        let left_bracket = |loc: usize| -> InstrumentAction {
350            InstrumentAction {
351                source_id,
352                loc,
353                content: InstrumentContent::Plain("{".to_string()),
354                priority: LEFT_BRACKET_PRIORITY,
355            }
356        };
357        let right_bracket = |loc: usize| -> InstrumentAction {
358            InstrumentAction {
359                source_id,
360                loc,
361                content: InstrumentContent::Plain("}".to_string()),
362                priority: RIGHT_BRACKET_PRIORITY,
363            }
364        };
365        let wrap_statement_as_block = |stmt: &Statement| -> Vec<Modification> {
366            let stmt_src = stmt_src(stmt);
367            // The left bracket is inserted just before the statement.
368            let start_pos = stmt_src.start.expect_with_context(
369                "statement start location not found",
370                source_id,
371                source,
372                &stmt_src,
373            );
374            let left_bracket = left_bracket(start_pos);
375
376            // The right bracket is inserted just after the statement.
377            let end_pos =
378                find_next_index_of_statement(source, stmt).expect("statement end not found");
379            let right_bracket = right_bracket(end_pos);
380
381            vec![left_bracket.into(), right_bracket.into()]
382        };
383
384        fn indeed_statement(block_or_stmt: &BlockOrStatement) -> Option<&Statement> {
385            match block_or_stmt {
386                BlockOrStatement::Statement(stmt) => match stmt {
387                    Statement::Block(_) => None,
388                    _ => Some(stmt),
389                },
390                BlockOrStatement::Block(_) => None,
391            }
392        }
393
394        for step in &analysis.steps {
395            match &step.variant() {
396                crate::analysis::StepVariant::IfCondition(if_stmt) => {
397                    // modify the true body if needed
398                    if let Some(stmt) = indeed_statement(&if_stmt.true_body) {
399                        let modifications = wrap_statement_as_block(stmt);
400                        self.extend_modifications(modifications);
401                    }
402
403                    // modify the false body if needed
404                    if let Some(stmt) =
405                        if_stmt.false_body.as_ref().and_then(|body| indeed_statement(body))
406                    {
407                        let modifications = wrap_statement_as_block(stmt);
408                        self.extend_modifications(modifications);
409                    }
410                }
411                crate::analysis::StepVariant::ForLoop(for_stmt) => {
412                    // modify the body if needed
413                    if let Some(stmt) = indeed_statement(&for_stmt.body) {
414                        let modifications = wrap_statement_as_block(stmt);
415                        self.extend_modifications(modifications);
416                    }
417                }
418                crate::analysis::StepVariant::WhileLoop(while_stmt) => {
419                    // modify the body if needed
420                    if let Some(stmt) = indeed_statement(&while_stmt.body) {
421                        let modifications = wrap_statement_as_block(stmt);
422                        self.extend_modifications(modifications);
423                    }
424                }
425                _ => {}
426            }
427        }
428        Ok(())
429    }
430
431    fn collect_before_step_hook_modifications(
432        &mut self,
433        compiler_version: VersionRef,
434        source: &str,
435        analysis: &SourceAnalysis,
436    ) -> Result<()> {
437        let source_id = self.source_id;
438        for step in &analysis.steps {
439            let usid = step.usid();
440            let variant = step.variant();
441            let function_calls = step.function_calls();
442            let (loc, priority) = match variant {
443                crate::analysis::StepVariant::FunctionEntry(function_definition) => {
444                    // the before step hook should be instrumented before the first statement of the function
445                    let Some(body) = &function_definition.body else {
446                        // skip the step if the function has no body
447                        continue;
448                    };
449                    // the first char of function body is the '{', so we insert after that.
450                    let loc = find_index_of_first_statement_in_block(body)
451                        .expect("function body start location not found");
452                    (loc, FUNCTION_ENTRY_PRIORITY)
453                }
454                crate::analysis::StepVariant::ModifierEntry(modifier_definition) => {
455                    // the before step hook should be instrumented before the first statement of the modifier
456                    let Some(body) = &modifier_definition.body else {
457                        // skip the step if the modifier has no body
458                        continue;
459                    };
460                    let loc = find_index_of_first_statement_in_block(body)
461                        .expect("modifier body start location not found");
462                    (loc, FUNCTION_ENTRY_PRIORITY)
463                }
464                crate::analysis::StepVariant::Statement(statement) => {
465                    // the before step hook should be instrumented before the statement
466                    let loc =
467                        stmt_src(statement).start.expect("statement start location not found");
468                    (loc, BEFORE_STEP_PRIORITY)
469                }
470                crate::analysis::StepVariant::Statements(statements) => {
471                    // the before step hook should be instrumented before the first statement
472                    let loc =
473                        stmt_src(&statements[0]).start.expect("statement start location not found");
474                    (loc, BEFORE_STEP_PRIORITY)
475                }
476                crate::analysis::StepVariant::IfCondition(if_statement) => {
477                    // the before step hook should be instrumented before the if statement
478                    let loc =
479                        if_statement.src.start.expect("if statement start location not found");
480                    (loc, BEFORE_STEP_PRIORITY)
481                }
482                crate::analysis::StepVariant::ForLoop(for_statement) => {
483                    // the before step hook should be instrumented before the for statement
484                    let loc =
485                        for_statement.src.start.expect("for statement start location not found");
486                    (loc, BEFORE_STEP_PRIORITY)
487                }
488                crate::analysis::StepVariant::WhileLoop(while_statement) => {
489                    // the before step hook should be instrumented before the while statement
490                    let loc = while_statement
491                        .src
492                        .start
493                        .expect("while statement start location not found");
494                    (loc, BEFORE_STEP_PRIORITY)
495                }
496                crate::analysis::StepVariant::DoWhileLoop(do_while_statement) => {
497                    // the before step hook should be instrumented after the last statement of the do-while statement
498                    let loc = find_next_index_of_last_statement_in_block(
499                        source,
500                        &do_while_statement.body,
501                    )
502                    .expect("do-while statement last statement location not found");
503                    (loc, BEFORE_STEP_PRIORITY)
504                }
505                crate::analysis::StepVariant::Try(try_statement) => {
506                    // the before step hook should be instrumented before the try statement
507                    let loc =
508                        try_statement.src.start.expect("try statement start location not found");
509                    (loc, BEFORE_STEP_PRIORITY)
510                }
511            };
512            let instrument_action = InstrumentAction {
513                source_id,
514                loc,
515                content: InstrumentContent::BeforeStepHook {
516                    version: compiler_version.clone(),
517                    usid,
518                    function_calls,
519                },
520                priority,
521            };
522            self.add_modification(instrument_action.into());
523        }
524
525        Ok(())
526    }
527
528    fn collect_variable_update_hook_modifications(
529        &mut self,
530        compiler_version: VersionRef,
531        source: &str,
532        analysis: &SourceAnalysis,
533    ) -> Result<()> {
534        let source_id = self.source_id;
535        for step in &analysis.steps {
536            let updated_variables = &step.read().updated_variables;
537            let locs: Vec<usize> = match step.variant() {
538                crate::analysis::StepVariant::FunctionEntry(function_definition) => {
539                    // the variable update hook should be instrumented before the first statement of the function
540                    let Some(body) = &function_definition.body else {
541                        // skip the step if the function has no body
542                        continue;
543                    };
544                    // the first char of function body is the '{', so we insert after that.
545                    vec![body.src.start.expect("function body start location not found") + 1]
546                }
547                crate::analysis::StepVariant::ModifierEntry(modifier_definition) => {
548                    // the variable update hook should be instrumented before the first statement of the modifier
549                    let Some(body) = &modifier_definition.body else {
550                        // skip the step if the modifier has no body
551                        continue;
552                    };
553                    vec![body.src.start.expect("modifier body start location not found") + 1]
554                }
555                crate::analysis::StepVariant::Statement(statement) => {
556                    match statement {
557                        Statement::Block(_)
558                        | Statement::UncheckedBlock(_)
559                        | Statement::DoWhileStatement(_)
560                        | Statement::ForStatement(_)
561                        | Statement::IfStatement(_)
562                        | Statement::TryStatement(_)
563                        | Statement::WhileStatement(_) => {
564                            unreachable!("should not be a statement step")
565                        }
566                        Statement::Break(_)
567                        | Statement::Continue(_)
568                        | Statement::PlaceholderStatement(_)
569                        | Statement::Return(_)
570                        | Statement::RevertStatement(_) => {
571                            // these statement does not have any variable update, so we skip it
572                            vec![]
573                        }
574                        Statement::EmitStatement(_)
575                        | Statement::ExpressionStatement(_)
576                        | Statement::InlineAssembly(_)
577                        | Statement::VariableDeclarationStatement(_) => {
578                            // the variable update hook should be instrumented after the emit statement
579                            find_next_index_of_statement(source, statement)
580                                .map(|loc| vec![loc])
581                                .unwrap_or_default()
582                        }
583                    }
584                }
585                crate::analysis::StepVariant::Statements(statements) => {
586                    // the variable update hook should be instrumented after the statments
587                    statements
588                        .last()
589                        .and_then(|stmt| {
590                            find_next_index_of_statement(source, stmt).map(|loc| vec![loc])
591                        })
592                        .unwrap_or_default()
593                }
594                crate::analysis::StepVariant::IfCondition(if_statement) => {
595                    // the variable update hook should be instrumented before the first statement of both the true and false bodies
596                    let mut locs = find_index_of_first_statement_in_block_or_statement(
597                        &if_statement.true_body,
598                    )
599                    .map(|loc| vec![loc])
600                    .unwrap_or_default();
601                    if let Some(false_loc) = if_statement.false_body.as_ref().and_then(|body| {
602                        find_index_of_first_statement_in_block_or_statement(body)
603                            .map(|loc| vec![loc])
604                    }) {
605                        locs.extend(false_loc.into_iter());
606                    }
607                    locs
608                }
609                crate::analysis::StepVariant::ForLoop(for_statement) => {
610                    // the variable update hook should be instrumented before the first statement of the for statement
611                    find_index_of_first_statement_in_block_or_statement(&for_statement.body)
612                        .map(|loc| vec![loc])
613                        .unwrap_or_default()
614                }
615                crate::analysis::StepVariant::WhileLoop(while_statement) => {
616                    // the variable update hook should be instrumented before the first statement of the while statement
617                    find_index_of_first_statement_in_block_or_statement(&while_statement.body)
618                        .map(|loc| vec![loc])
619                        .unwrap_or_default()
620                }
621                crate::analysis::StepVariant::DoWhileLoop(do_while_statement) => {
622                    // the variable update hook should be instrumented before the do-while statement
623                    find_next_index_of_source_location(&do_while_statement.src)
624                        .map(|loc| vec![loc])
625                        .unwrap_or_default()
626                }
627                crate::analysis::StepVariant::Try(try_statement) => {
628                    // the variable update hook should be instrumented before the first statement in all catch blocks
629                    try_statement
630                        .clauses
631                        .iter()
632                        .filter_map(|clause| find_index_of_first_statement_in_block(&clause.block))
633                        .collect()
634                }
635            };
636            for loc in locs {
637                for updated_variable in updated_variables {
638                    let uvid = updated_variable.id();
639                    let instrument_action = InstrumentAction {
640                        source_id,
641                        loc,
642                        content: InstrumentContent::VariableUpdateHook {
643                            version: compiler_version.clone(),
644                            uvid,
645                            variable: updated_variable.clone(),
646                        },
647                        priority: VARIABLE_UPDATE_PRIORITY,
648                    };
649
650                    self.add_modification(instrument_action.into());
651                }
652            }
653        }
654        Ok(())
655    }
656}
657
658/// Trait to extend Option with better context for instrumentation failures
659trait ExpectWithContext<T> {
660    fn expect_with_context(
661        self,
662        error_msg: &str,
663        source_id: u32,
664        source: &str,
665        src_loc: &SourceLocation,
666    ) -> T;
667}
668
669impl<T> ExpectWithContext<T> for Option<T> {
670    fn expect_with_context(
671        self,
672        error_msg: &str,
673        source_id: u32,
674        source: &str,
675        src_loc: &SourceLocation,
676    ) -> T {
677        match self {
678            Some(value) => value,
679            None => {
680                // Simple source dump to temp file
681                let temp_dir = std::env::temp_dir();
682                let dump_path = temp_dir.join(format!("edb_fail_source_{source_id}.sol"));
683                let _ = std::fs::write(&dump_path, source);
684
685                // Extract context around the error location with line numbers
686                let context = if let Some(start) = src_loc.start {
687                    let context_start = start.saturating_sub(200);
688                    let context_end = (start + 200).min(source.len());
689                    let context_slice = &source[context_start..context_end];
690
691                    // Find line number of the error position
692                    let lines_before_context: Vec<&str> =
693                        source[..context_start].split('\n').collect();
694                    let context_lines: Vec<&str> = context_slice.split('\n').collect();
695                    let start_line_num = lines_before_context.len();
696                    let error_pos_in_context = start - context_start;
697
698                    // Find which line contains the error
699                    let mut current_pos = 0;
700                    let mut error_line_idx = 0;
701                    let mut error_col = 0;
702
703                    for (idx, line) in context_lines.iter().enumerate() {
704                        let line_end = current_pos + line.len();
705                        if error_pos_in_context >= current_pos && error_pos_in_context <= line_end {
706                            error_line_idx = idx;
707                            error_col = error_pos_in_context - current_pos;
708                            break;
709                        }
710                        current_pos = line_end + 1; // +1 for newline
711                    }
712
713                    let mut formatted_context = format!(
714                        "\n  Source context around line {}:",
715                        start_line_num + error_line_idx
716                    );
717
718                    for (idx, line) in context_lines.iter().enumerate() {
719                        if line.trim().is_empty() && idx != error_line_idx {
720                            continue; // Skip empty lines except the error line
721                        }
722
723                        let line_num = start_line_num + idx;
724                        let marker = if idx == error_line_idx { " --> " } else { "     " };
725                        formatted_context.push_str(&format!("\n{marker}{line_num:4} | {line}"));
726
727                        // Add error pointer for the error line
728                        if idx == error_line_idx {
729                            let pointer = format!(
730                                "\n     {} | {}{}^ error here",
731                                " ".repeat(4),
732                                "_".repeat(error_col),
733                                ""
734                            );
735                            formatted_context.push_str(&pointer);
736                        }
737                    }
738
739                    formatted_context
740                } else {
741                    String::new()
742                };
743
744                panic!(
745                    "{}\n  Source ID: {}\n  Source location: {:?}{}\n  Full source dumped to: {}",
746                    error_msg,
747                    source_id,
748                    src_loc,
749                    context,
750                    dump_path.display()
751                );
752            }
753        }
754    }
755}
756
757#[cfg(test)]
758mod tests {
759    use crate::analysis::{self, tests::compile_and_analyze};
760
761    use super::*;
762
763    #[test]
764    fn test_collect_statement_to_block_modifications() {
765        let source = r#"
766        contract C {
767            function a() public returns (uint256) {
768                if (false )return 0;
769
770                if (true)return 0;
771                else    return 1 ;
772            }
773
774            function b() public returns (uint256 x) {
775                for (uint256 i = 0; i < 10; i++)x += i
776                ;
777            }
778
779            function c() public returns (uint256) {
780                while (true) return 0;
781            }
782        }
783        "#;
784
785        let (_sources, analysis) = analysis::tests::compile_and_analyze(source);
786
787        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
788        modifications.collect_statement_to_block_modifications(source, &analysis).unwrap();
789        assert_eq!(modifications.modifications.len(), 10);
790        let modified_source = modifications.modify_source(source);
791
792        // The modified source should be able to be compiled and analyzed.
793        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
794    }
795
796    #[test]
797    fn test_do_while_loop_step_modification() {
798        let source = r#"
799        contract C {
800            function a() public returns (uint256) {
801                do {
802                    uint x = 1;
803                    return 0;
804                } while (false);
805            }
806        }
807        "#;
808
809        let (_sources, analysis) = analysis::tests::compile_and_analyze(source);
810
811        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
812        let version = Arc::new(Version::parse("0.8.0").unwrap());
813        modifications.collect_before_step_hook_modifications(version, source, &analysis).unwrap();
814        assert_eq!(modifications.modifications.len(), 4);
815        let modified_source = modifications.modify_source(source);
816
817        // The modified source should be able to be compiled and analyzed.
818        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
819    }
820
821    #[test]
822    fn test_collect_function_entry_step_hook_modifications() {
823        let source = r#"
824        abstract contract C {
825            function v() public virtual returns (uint256);
826
827            function a() public returns (uint256) {
828                uint x = 1;
829            }
830        }
831        "#;
832
833        let (_sources, analysis) = analysis::tests::compile_and_analyze(source);
834
835        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
836        let version = Arc::new(Version::parse("0.8.0").unwrap());
837        modifications.collect_before_step_hook_modifications(version, source, &analysis).unwrap();
838        // assert_eq!(modifications.modifications.len(), 1);
839        let modified_source = modifications.modify_source(source);
840
841        // The modified source should be able to be compiled and analyzed.
842        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
843    }
844
845    #[test]
846    fn test_collect_before_step_hook_modifications() {
847        let source = r#"
848        abstract contract C {
849            function a() public returns (uint256) {
850                if (false) {return 0;}
851                else    {return 1;}
852                for (uint256 i = 0; i < 10; i++) {
853                    return 0;
854                }
855                while (true) {
856                    return 0;
857                }
858                do {
859                    return 0;
860                } while (false);
861                try this.a() {
862                    return 0;
863                }
864                catch {}
865                return 0;
866            }
867        }
868        "#;
869
870        let (_sources, analysis) = analysis::tests::compile_and_analyze(source);
871
872        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
873        let version = Arc::new(Version::parse("0.8.0").unwrap());
874        modifications.collect_before_step_hook_modifications(version, source, &analysis).unwrap();
875        assert_eq!(modifications.modifications.len(), 13);
876        let modified_source = modifications.modify_source(source);
877
878        // The modified source should be able to be compiled and analyzed.
879        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
880    }
881
882    #[test]
883    fn test_modifier_is_not_step() {
884        let source = r#"
885        contract C {
886            modifier m(uint x) {
887                _;
888            }
889
890            function a() public m(1) {}
891        }
892        "#;
893        let (_sources, analysis) = analysis::tests::compile_and_analyze(source);
894
895        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
896        let version = Arc::new(Version::parse("0.8.0").unwrap());
897        modifications.collect_before_step_hook_modifications(version, source, &analysis).unwrap();
898        assert_eq!(modifications.modifications.len(), 1);
899        let modified_source = modifications.modify_source(source);
900
901        // The modified source should be able to be compiled and analyzed.
902        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
903    }
904
905    #[test]
906    fn test_else_if_statement_to_block() {
907        let source = r#"
908contract TestContract {
909    function foo() public {
910        if (true)
911            revert();
912        else if (false)
913            return;
914        else {
915            require(true, "error");
916        }
917    }
918}
919"#;
920        let (_sources, analysis) = compile_and_analyze(source);
921
922        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
923        modifications.collect_statement_to_block_modifications(source, &analysis).unwrap();
924        assert_eq!(modifications.modifications.len(), 6);
925        let modified_source = modifications.modify_source(source);
926
927        // The modified source should be able to be compiled and analyzed.
928        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
929    }
930
931    #[test]
932    fn test_if_for_statement_to_block() {
933        let source = r#"
934contract TestContract {
935    function foo() public {
936        if (true)
937            for (uint256 i = 0; i < 10; i++)
938                return;
939        else
940            while (true)
941                return;
942    }
943}
944"#;
945        let (_sources, analysis) = compile_and_analyze(source);
946
947        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
948        modifications.collect_statement_to_block_modifications(source, &analysis).unwrap();
949        assert_eq!(modifications.modifications.len(), 6);
950        let modified_source = modifications.modify_source(source);
951
952        // The modified source should be able to be compiled and analyzed.
953        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
954    }
955
956    #[test]
957    fn test_variable_update_hook_modification_for_for_loop() {
958        let source = r#"
959        contract C {
960            function a() public returns (uint256) {
961                for (uint i = 0; i < 10; i++) {
962                    return i;
963                }
964            }
965        }
966        "#;
967        let (_sources, analysis) = compile_and_analyze(source);
968
969        let mut modifications = SourceModifications::new(analysis::tests::TEST_CONTRACT_SOURCE_ID);
970        modifications
971            .collect_variable_update_hook_modifications(
972                Arc::new(Version::parse("0.8.0").unwrap()),
973                source,
974                &analysis,
975            )
976            .unwrap();
977        assert_eq!(modifications.modifications.len(), 1);
978        let modified_source = modifications.modify_source(source);
979
980        // The modified source should be able to be compiled and analyzed.
981        let (_sources, _analysis2) = analysis::tests::compile_and_analyze(&modified_source);
982    }
983}
edb_engine/instrumentation/modification.rs

edb_engine/instrumentation/
modification.rs
