feat(forge-lint): [claude] check for unwrapped modifiers

crates/lint/src/linter/mod.rs

@@ -89,13 +89,13 @@ impl<'s> LintContext<'s> {
 
         // Convert the snippet to ensure we have the appropriate type
         let snippet = match snippet {
-            Snippet::Diff { desc, span: diff_span, add } => {
+            Snippet::Diff { desc, span: diff_span, add, trim_code } => {
                 // Use the provided span or fall back to the lint span
                 let target_span = diff_span.unwrap_or(span);
 
                 // Check if we can get the original code
                 if self.span_to_snippet(target_span).is_some() {
-                    Snippet::Diff { desc, span: Some(target_span), add }
+                    Snippet::Diff { desc, span: Some(target_span), add, trim_code }
                 } else {
                     // Fall back to Block if we can't get the original code
                     Snippet::Block { desc, code: add }
@@ -118,47 +118,93 @@ impl<'s> LintContext<'s> {
         diag.emit();
     }
 
+    /// Gets the "raw" source code (snippet) of the given span.
     pub fn span_to_snippet(&self, span: Span) -> Option<String> {
         self.sess.source_map().span_to_snippet(span).ok()
     }
+
+    /// Gets the number of leading whitespaces (indentation) of the line where the span begins.
+    pub fn get_ind_for_span(&self, span: Span) -> usize {
+        if span.is_dummy() {
+            return 0;
+        }
+
+        // Get the line text and compute the indentation prior to the span's position.
+        let loc = self.sess.source_map().lookup_char_pos(span.lo());
+        if let Some(line_text) = loc.file.get_line(loc.line) {
+            let col_offset = loc.col.to_usize();
+            if col_offset <= line_text.len() {
+                let prev_text = &line_text[..col_offset];
+                return prev_text.len() - prev_text.trim().len();
+            }
+        }
+
+        0
+    }
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum Snippet {
-    /// Represents a code block. Can have an optional description.
-    Block { desc: Option<&'static str>, code: String },
-    /// Represents a code diff. Can have an optional description and a span for the code to remove.
-    Diff { desc: Option<&'static str>, span: Option<Span>, add: String },
+    /// A standalone block of code. Used for showing examples without suggesting a fix.
+    Block {
+        /// An optional description displayed above the code block.
+        desc: Option<&'static str>,
+        /// The source code to display. Multi-line strings should include newlines.
+        code: String,
+    },
+
+    /// A proposed code change, displayed as a diff. Used to suggest replacements, showing the code
+    /// to be removed (from `span`) and the code to be added (from `add`).
+    Diff {
+        /// An optional description displayed above the diff.
+        desc: Option<&'static str>,
+        /// The `Span` of the source code to be removed. Note that, if uninformed,
+        /// `fn emit_with_fix()` falls back to the lint span.
+        span: Option<Span>,
+        /// The fix.
+        add: String,
+        /// If `true`, the leading whitespaces of the first line will be trimed from the whole code
+        /// block. Applies to both, the added and removed code. This is useful for aligning
+        /// the indentation of multi-line replacements.
+        trim_code: bool,
+    },
 }
 
 impl Snippet {
     pub fn to_note(self, ctx: &LintContext<'_>) -> Vec<(DiagMsg, Style)> {
-        let mut output = Vec::new();
-        match self.desc() {
-            Some(desc) => {
-                output.push((DiagMsg::from(desc), Style::NoStyle));
-                output.push((DiagMsg::from("\n\n"), Style::NoStyle));
-            }
-            None => output.push((DiagMsg::from(" \n"), Style::NoStyle)),
-        }
+        let mut output = if let Some(desc) = self.desc() {
+            vec![(DiagMsg::from(desc), Style::NoStyle), (DiagMsg::from("\n\n"), Style::NoStyle)]
+        } else {
+            vec![(DiagMsg::from(" \n"), Style::NoStyle)]
+        };
+
         match self {
-            Self::Diff { span, add, .. } => {
-                // Get the original code from the span if provided
+            Self::Diff { span, add, trim_code: trim, .. } => {
+                // Get the original code from the span if provided and normalize its indentation
                 if let Some(span) = span
                     && let Some(rmv) = ctx.span_to_snippet(span)
                 {
-                    for line in rmv.lines() {
-                        output.push((DiagMsg::from(format!("- {line}\n")), Style::Removal));
-                    }
-                }
-                for line in add.lines() {
-                    output.push((DiagMsg::from(format!("+ {line}\n")), Style::Addition));
+                    let ind = ctx.get_ind_for_span(span);
+                    output.extend(rmv.lines().map(|line| {
+                        let content = if trim { Self::trim_start_limited(line, ind) } else { line };
+                        (DiagMsg::from(format!("- {content}\n")), Style::Removal)
+                    }));
+                    output.extend(add.lines().map(|line| {
+                        let content = if trim { Self::trim_start_limited(line, ind) } else { line };
+                        (DiagMsg::from(format!("+ {content}\n")), Style::Addition)
+                    }));
+                } else {
+                    // Should never happen, but fall back to `Self::Block` behavior.
+                    output.extend(
+                        add.lines()
+                            .map(|line| (DiagMsg::from(format!("{line}\n")), Style::NoStyle)),
+                    );
                 }
             }
             Self::Block { code, .. } => {
-                for line in code.lines() {
-                    output.push((DiagMsg::from(format!("- {line}\n")), Style::NoStyle));
-                }
+                output.extend(
+                    code.lines().map(|line| (DiagMsg::from(format!("{line}\n")), Style::NoStyle)),
+                );
             }
         }
         output.push((DiagMsg::from("\n"), Style::NoStyle));
@@ -171,4 +217,17 @@ impl Snippet {
             Self::Block { desc, .. } => *desc,
         }
     }
+
+    fn trim_start_limited(s: &str, max_chars: usize) -> &str {
+        let (mut chars, mut byte_offset) = (0, 0);
+        for c in s.chars() {
+            if chars >= max_chars || !c.is_whitespace() {
+                break;
+            }
+            chars += 1;
+            byte_offset += c.len_utf8();
+        }
+
+        &s[byte_offset..]
+    }
 }

crates/lint/src/sol/gas/mod.rs

@@ -3,4 +3,10 @@ use crate::sol::{EarlyLintPass, LateLintPass, SolLint};
 mod keccak;
 use keccak::ASM_KECCAK256;
 
-register_lints!((AsmKeccak256, late, (ASM_KECCAK256)));
+mod unwrapped_modifier_logic;
+use unwrapped_modifier_logic::UNWRAPPED_MODIFIER_LOGIC;
+
+register_lints!(
+    (AsmKeccak256, late, (ASM_KECCAK256)),
+    (UnwrappedModifierLogic, late, (UNWRAPPED_MODIFIER_LOGIC))
+);

crates/lint/src/sol/gas/unwrapped_modifier_logic.rs

@@ -0,0 +1,179 @@
+use super::UnwrappedModifierLogic;
+use crate::{
+    linter::{LateLintPass, LintContext, Snippet},
+    sol::{Severity, SolLint},
+};
+use solar_ast::Span;
+use solar_sema::hir::{self, Res};
+
+declare_forge_lint!(
+    UNWRAPPED_MODIFIER_LOGIC,
+    Severity::Gas,
+    "unwrapped-modifier-logic",
+    "wrap modifier logic to reduce code size"
+);
+
+impl<'hir> LateLintPass<'hir> for UnwrappedModifierLogic {
+    fn check_function(
+        &mut self,
+        ctx: &LintContext<'_>,
+        hir: &'hir hir::Hir<'hir>,
+        func: &'hir hir::Function<'hir>,
+    ) {
+        // Only check modifiers with a body and a name
+        let (body, name) = match (func.kind, &func.body, func.name) {
+            (solar_ast::FunctionKind::Modifier, Some(body), Some(name)) => (body, name),
+            _ => return,
+        };
+
+        // Split statements into before and after the placeholder `_`.
+        let stmts = &body.stmts[..];
+        let (before, after) = stmts
+            .iter()
+            .position(|s| matches!(s.kind, hir::StmtKind::Placeholder))
+            .map_or((stmts, &[][..]), |idx| (&stmts[..idx], &stmts[idx + 1..]));
+
+        // Generate a fix snippet if the modifier logic should be wrapped.
+        if let Some(snippet) = self.get_snippet(ctx, hir, func, before, after) {
+            ctx.emit_with_fix(&UNWRAPPED_MODIFIER_LOGIC, name.span, snippet);
+        }
+    }
+}
+
+impl UnwrappedModifierLogic {
+    fn is_valid_expr(&self, hir: &hir::Hir<'_>, expr: &hir::Expr<'_>) -> bool {
+        match &expr.kind {
+            hir::ExprKind::Call(func_expr, _, _) => match &func_expr.kind {
+                hir::ExprKind::Ident(resolutions) => {
+                    !resolutions.iter().any(|r| matches!(r, Res::Builtin(_)))
+                }
+                hir::ExprKind::Member(base, _) => {
+                    if let hir::ExprKind::Ident(resolutions) = &base.kind {
+                        resolutions.iter().any(|r| matches!(r, Res::Item(hir::ItemId::Contract(id)) if hir.contract(*id).kind == solar_ast::ContractKind::Library))
+                    } else {
+                        false
+                    }
+                }
+                _ => false,
+            },
+            _ => false,
+        }
+    }
+
+    fn is_valid_stmt(&self, hir: &hir::Hir<'_>, stmt: &hir::Stmt<'_>) -> bool {
+        match &stmt.kind {
+            hir::StmtKind::Expr(expr) => self.is_valid_expr(hir, expr),
+            hir::StmtKind::Placeholder => true,
+            _ => false,
+        }
+    }
+
+    fn check_stmts(&self, hir: &hir::Hir<'_>, stmts: &[hir::Stmt<'_>]) -> bool {
+        let mut has_valid = false;
+        for stmt in stmts {
+            if !self.is_valid_stmt(hir, stmt) {
+                return true;
+            }
+            if let hir::StmtKind::Expr(expr) = &stmt.kind
+                && self.is_valid_expr(hir, expr)
+            {
+                if has_valid {
+                    return true;
+                }
+                has_valid = true;
+            }
+        }
+        false
+    }
+
+    fn get_snippet<'a>(
+        &self,
+        ctx: &LintContext<'_>,
+        hir: &hir::Hir<'_>,
+        func: &hir::Function<'_>,
+        before: &'a [hir::Stmt<'a>],
+        after: &'a [hir::Stmt<'a>],
+    ) -> Option<Snippet> {
+        let wrap_before = !before.is_empty() && self.check_stmts(hir, before);
+        let wrap_after = !after.is_empty() && self.check_stmts(hir, after);
+
+        if !(wrap_before || wrap_after) {
+            return None;
+        }
+
+        let binding = func.name.unwrap();
+        let modifier_name = binding.name.as_str();
+
+        // Get parameter names from the function
+        let param_list = func
+            .parameters
+            .iter()
+            .filter_map(|&var_id| {
+                let var = hir.variable(var_id);
+                var.name.map(|n| n.name.to_string())
+            })
+            .collect::<Vec<_>>()
+            .join(", ");
+
+        // Get parameter declarations with types
+        let param_decls = func
+            .parameters
+            .iter()
+            .map(|&var_id| {
+                let var = hir.variable(var_id);
+                let name = var.name.as_ref().map(|n| n.name.as_str()).unwrap_or("");
+                let ty = ctx
+                    .span_to_snippet(var.ty.span)
+                    .unwrap_or_else(|| "/* unknown type */".to_string());
+                format!("{ty} {name}")
+            })
+            .collect::<Vec<_>>()
+            .join(", ");
+
+        let body_indent = " ".repeat(ctx.get_ind_for_span(
+            before.first().or(after.first()).map(|stmt| stmt.span).unwrap_or(func.span),
+        ));
+        let body = match (wrap_before, wrap_after) {
+            (true, true) => format!(
+                "{body_indent}_{modifier_name}Before({param_list});\n{body_indent}_;\n{body_indent}_{modifier_name}After({param_list});"
+            ),
+            (true, false) => {
+                format!("{body_indent}_{modifier_name}({param_list});\n{body_indent}_;")
+            }
+            (false, true) => {
+                format!("{body_indent}_;\n{body_indent}_{modifier_name}({param_list});")
+            }
+            _ => unreachable!(),
+        };
+
+        let mod_indent = " ".repeat(ctx.get_ind_for_span(func.span));
+        let mut replacement = format!(
+            "{mod_indent}modifier {modifier_name}({param_decls}) {{\n{body}\n{mod_indent}}}"
+        );
+
+        let build_func = |stmts: &[hir::Stmt<'_>], suffix: &str| {
+            let body_stmts = stmts
+                .iter()
+                .filter_map(|s| ctx.span_to_snippet(s.span))
+                .map(|code| format!("\n{body_indent}{code}"))
+                .collect::<String>();
+            format!(
+                "\n\n{mod_indent}function _{modifier_name}{suffix}({param_decls}) internal {{{body_stmts}\n{mod_indent}}}"
+            )
+        };
+
+        if wrap_before {
+            replacement.push_str(&build_func(before, if wrap_after { "Before" } else { "" }));
+        }
+        if wrap_after {
+            replacement.push_str(&build_func(after, if wrap_before { "After" } else { "" }));
+        }
+
+        Some(Snippet::Diff {
+            desc: Some("wrap modifier logic to reduce code size"),
+            span: Some(Span::new(func.span.lo(), func.body_span.hi())),
+            add: replacement,
+            trim_code: true,
+        })
+    }
+}