From ac02c83a86f4930412d36fa02a991d588221c062 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Petr=20Balv=C3=ADn?= <opensource@petrbalvin.org>
Date: Fri, 21 Aug 2026 00:17:38 +0200
Subject: [PATCH] feat(lint): add stack-imbalance rule

Assisted-by: MiMo V2.5 Pro
---
 lint/lint.go      | 121 ++++++++++++++++++++++++++++++++++++++++++++++
 lint/lint_test.go |  27 +++++++++++
 2 files changed, 148 insertions(+)

diff --git a/lint/lint.go b/lint/lint.go
index 5f3656a..e42605e 100644
--- a/lint/lint.go
+++ b/lint/lint.go
@@ -76,6 +76,7 @@ const (
 	CodeFuncdata        = "funcdata-pcdata"
 	CodeUnusedLabel     = "unused-label"
 	CodeInvalidFlag     = "invalid-textflag"
+	CodeStackImbalance  = "stack-imbalance"
 )
 
 // knownTextFlags are the flags recognised by the Go assembler's textflag.h.
@@ -350,6 +351,27 @@ func lintText(t *ast.Text, tab *arch.Table, archKnown bool, cfg Config, macros m
 		})
 	}
 
+	// Stack imbalance: track SP changes and flag if the net delta at RET
+	// does not match the declared frame size.  Only checked for functions
+	// with a declared frame, no macros, and no indirect branches.
+	if doLabelChecks && !cfg.Disable[CodeStackImbalance] && !hasMacro && !hasIndirect {
+		frameSize := int64(0)
+		if t.Frame != nil && t.Frame.Imm.HasVal {
+			frameSize = t.Frame.Imm.Val
+		}
+		if frameSize > 0 {
+			delta := stackDelta(t, cfg.Arch)
+			if delta != 0 && delta != -frameSize {
+				out = append(out, Diagnostic{
+					Pos:      t.Keyword.Pos,
+					Severity: Warning,
+					Code:     CodeStackImbalance,
+					Message:  fmt.Sprintf("function %q has net SP delta %d (frame size %d)", t.Name.Name, delta, frameSize),
+				})
+			}
+		}
+	}
+
 	// ABI conformance: the argument area declared in the TEXT directive should
 	// match the size computed from the // func signature in the doc comment.
 	// Only applies to stack-argument (ABI0) functions, which reference their
@@ -594,6 +616,105 @@ func isUnconditionalJump(a arch.Arch, upper string) bool {
 	}
 }
 
+// stackDelta computes the net SP change across a function body.
+// It tracks PUSH/POP and SUB/ADD on SP.  Returns the net delta (negative
+// means SP decreased, which is the normal direction for stack growth).
+func stackDelta(t *ast.Text, a arch.Arch) int64 {
+	var delta int64
+	for _, s := range t.Body {
+		in, ok := s.(*ast.Instr)
+		if !ok {
+			continue
+		}
+		upper := strings.ToUpper(in.Mnemonic.Text)
+		switch a {
+		case arch.AMD64:
+			switch upper {
+			case "PUSHQ", "PUSHL", "PUSHW":
+				delta -= 8
+			case "POPQ", "POPL", "POPW":
+				delta += 8
+			case "SUBQ", "SUBL":
+				if len(in.Operands) >= 2 && isSPReg(in.Operands[1], a) {
+					if in.Operands[0].Imm.HasVal {
+						delta -= in.Operands[0].Imm.Val
+					}
+				}
+			case "ADDQ", "ADDL":
+				if len(in.Operands) >= 2 && isSPReg(in.Operands[1], a) {
+					if in.Operands[0].Imm.HasVal {
+						delta += in.Operands[0].Imm.Val
+					}
+				}
+			}
+		case arch.ARM64:
+			switch upper {
+			case "STP":
+				// STP with pre-index: STP Xt1, Xt2, [SP, #imm]!
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					// Could be pre-index decrement; skip for simplicity.
+				}
+			case "SUB":
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					if in.Operands[1].Imm.HasVal {
+						delta -= in.Operands[1].Imm.Val
+					}
+				}
+			case "ADD":
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					if in.Operands[1].Imm.HasVal {
+						delta += in.Operands[1].Imm.Val
+					}
+				}
+			}
+		case arch.RISCV:
+			switch upper {
+			case "ADDI":
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					if in.Operands[1].Imm.HasVal {
+						delta += in.Operands[1].Imm.Val
+					}
+				}
+			}
+		case arch.LOONG64:
+			switch upper {
+			case "ADDI.D", "ADDI.W":
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					if in.Operands[1].Imm.HasVal {
+						delta += in.Operands[1].Imm.Val
+					}
+				}
+			case "ADD.D", "ADD.W":
+				if len(in.Operands) >= 3 && isSPReg(in.Operands[2], a) {
+					if in.Operands[1].Imm.HasVal {
+						delta += in.Operands[1].Imm.Val
+					}
+				}
+			}
+		}
+	}
+	return delta
+}
+
+// isSPReg reports whether the operand is the stack pointer register.
+func isSPReg(op *ast.Operand, a arch.Arch) bool {
+	if op == nil || op.Addr.Sym == nil {
+		return false
+	}
+	name := strings.ToUpper(op.Addr.Sym.Name)
+	switch a {
+	case arch.AMD64:
+		return name == "RSP" || name == "ESP" || name == "SP"
+	case arch.ARM64:
+		return name == "R31" || name == "SP"
+	case arch.RISCV:
+		return name == "X2" || name == "SP"
+	case arch.LOONG64:
+		return name == "R3" || name == "SP"
+	}
+	return false
+}
+
 func countRange(min, max int) string {
 	if min == max {
 		return fmt.Sprintf("%d operand(s)", min)
diff --git a/lint/lint_test.go b/lint/lint_test.go
index a816546..917fa14 100644
--- a/lint/lint_test.go
+++ b/lint/lint_test.go
@@ -282,3 +282,30 @@ TEXT ·f(SB), NOSPLIT|NOFRAME|DUPOK, $0
 		t.Fatalf("valid flags must not be flagged: %+v", diags)
 	}
 }
+
+func TestStackImbalance(t *testing.T) {
+	// Function with frame size 16 but only SUB 8, SP — imbalance.
+	diags := lintSrc(t, `
+#include "textflag.h"
+TEXT ·f(SB), NOSPLIT, $16-0
+	SUBQ $8, SP
+	RET
+`)
+	if codes(diags)[CodeStackImbalance] != 1 {
+		t.Fatalf("want one stack-imbalance, got %+v", diags)
+	}
+}
+
+func TestStackBalanced(t *testing.T) {
+	// Function with frame size 16 and matching SUB/ADD — balanced.
+	diags := lintSrc(t, `
+#include "textflag.h"
+TEXT ·f(SB), NOSPLIT, $16-0
+	SUBQ $16, SP
+	ADDQ $16, SP
+	RET
+`)
+	if codes(diags)[CodeStackImbalance] != 0 {
+		t.Fatalf("balanced stack must not be flagged: %+v", diags)
+	}
+}