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

crates/config/src/lint.rs

@@ -45,6 +45,7 @@ pub enum Severity {
     Low,
     Info,
     Gas,
+    CodeSize,
 }
 
 impl Severity {
@@ -55,6 +56,7 @@ impl Severity {
             Self::Low => Paint::yellow(message).bold().to_string(),
             Self::Info => Paint::cyan(message).bold().to_string(),
             Self::Gas => Paint::green(message).bold().to_string(),
+            Self::CodeSize => Paint::green(message).bold().to_string(),
         }
     }
 }
@@ -63,7 +65,7 @@ impl From<Severity> for Level {
     fn from(severity: Severity) -> Self {
         match severity {
             Severity::High | Severity::Med | Severity::Low => Self::Warning,
-            Severity::Info | Severity::Gas => Self::Note,
+            Severity::Info | Severity::Gas | Severity::CodeSize => Self::Note,
         }
     }
 }
@@ -76,6 +78,7 @@ impl fmt::Display for Severity {
             Self::Low => self.color("Low"),
             Self::Info => self.color("Info"),
             Self::Gas => self.color("Gas"),
+            Self::CodeSize => self.color("CodeSize"),
         };
         write!(f, "{colored}")
     }
@@ -102,8 +105,9 @@ impl FromStr for Severity {
             "low" => Ok(Self::Low),
             "info" => Ok(Self::Info),
             "gas" => Ok(Self::Gas),
+            "size" | "codesize" | "code-size" => Ok(Self::CodeSize),
             _ => Err(format!(
-                "unknown variant: found `{s}`, expected `one of `High`, `Med`, `Low`, `Info`, `Gas``"
+                "unknown variant: found `{s}`, expected `one of `High`, `Med`, `Low`, `Info`, `Gas`, `CodeSize`"
             )),
         }
     }

crates/lint/src/linter/mod.rs

@@ -93,13 +93,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 }
@@ -122,47 +122,89 @@ 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_span_indentation(&self, span: Span) -> usize {
+        if !span.is_dummy() {
+            // 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 trimmed 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_span_indentation(span);
+                    let diag_msg = |line: &str, prefix: &str, style: Style| {
+                        let content = if trim { Self::trim_start_limited(line, ind) } else { line };
+                        (DiagMsg::from(format!("{prefix}{content}\n")), style)
+                    };
+                    output.extend(rmv.lines().map(|line| diag_msg(line, "- ", Style::Removal)));
+                    output.extend(add.lines().map(|line| diag_msg(line, "+ ", 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));
@@ -175,4 +217,18 @@ impl Snippet {
             Self::Block { desc, .. } => *desc,
         }
     }
+
+    /// Removes up to `max_chars` whitespaces from the start of the string.
+    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/codesize/mod.rs

@@ -0,0 +1,6 @@
+use crate::sol::{EarlyLintPass, LateLintPass, SolLint};
+
+mod unwrapped_modifier_logic;
+use unwrapped_modifier_logic::UNWRAPPED_MODIFIER_LOGIC;
+
+register_lints!((UnwrappedModifierLogic, late, (UNWRAPPED_MODIFIER_LOGIC)));

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

@@ -0,0 +1,168 @@
+use super::UnwrappedModifierLogic;
+use crate::{
+    linter::{LateLintPass, LintContext, Snippet},
+    sol::{Severity, SolLint},
+};
+use solar_ast::{self as 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) {
+            (ast::FunctionKind::Modifier, Some(body), Some(name)) => (body, name),
+            _ => return,
+        };
+
+        // Split statements into before and after the placeholder `_`.
+        let stmts = body.stmts[..].as_ref();
+        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 {
+    /// Returns `true` if an expr is not a built-in ('require' or 'assert') call or a lib function.
+    fn is_valid_expr(&self, hir: &hir::Hir<'_>, expr: &hir::Expr<'_>) -> bool {
+        if let hir::ExprKind::Call(func_expr, _, _) = &expr.kind {
+            if let hir::ExprKind::Ident(resolutions) = &func_expr.kind {
+                return !resolutions.iter().any(|r| matches!(r, Res::Builtin(_)));
+            }
+
+            if let hir::ExprKind::Member(base, _) = &func_expr.kind
+                && let hir::ExprKind::Ident(resolutions) = &base.kind
+            {
+                return resolutions.iter().any(|r| {
+                    matches!(r, Res::Item(hir::ItemId::Contract(id)) if hir.contract(*id).kind == ast::ContractKind::Library)
+                });
+            }
+        }
+
+        false
+    }
+
+    /// Checks if a block of statements is complex and should be wrapped in a helper function.
+    ///
+    /// This is true if the block contains:
+    /// 1. Any statement that is not a placeholder or a valid expression.
+    /// 2. More than one simple call expression.
+    fn stmts_require_wrapping(&self, hir: &hir::Hir<'_>, stmts: &[hir::Stmt<'_>]) -> bool {
+        let mut has_valid_stmt = false;
+        for stmt in stmts {
+            match &stmt.kind {
+                hir::StmtKind::Placeholder => continue,
+                hir::StmtKind::Expr(expr) => {
+                    if !self.is_valid_expr(hir, expr) || has_valid_stmt {
+                        return true;
+                    }
+                    has_valid_stmt = true;
+                }
+                _ => return 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.stmts_require_wrapping(hir, before);
+        let wrap_after = !after.is_empty() && self.stmts_require_wrapping(hir, after);
+
+        if !(wrap_before || wrap_after) {
+            return None;
+        }
+
+        let binding = func.name.unwrap();
+        let modifier_name = binding.name.as_str();
+        let mut param_list = vec![];
+        let mut param_decls = vec![];
+
+        for var_id in func.parameters {
+            let var = hir.variable(*var_id);
+            let ty = ctx
+                .span_to_snippet(var.ty.span)
+                .unwrap_or_else(|| "/* unknown type */".to_string());
+
+            // solidity functions should always have named parameters
+            if let Some(ident) = var.name {
+                param_list.push(ident.to_string());
+                param_decls.push(format!("{ty} {}", ident.to_string()));
+            }
+        }
+
+        let param_list = param_list.join(", ");
+        let param_decls = param_decls.join(", ");
+
+        let body_indent = " ".repeat(ctx.get_span_indentation(
+            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_span_indentation(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,
+        })
+    }
+}

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

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