fix(cmd): Handle imbalanced LHS and RHS in defer assignment

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 ### Fixed
 - cmd/errtrace: Don't exit with a non-zero status when `-h` is used.
+- cmd/errtrace: Don't panic on imbalanced assignments inside defer blocks.
 
 ## v0.3.0 - 2023-12-22
 

cmd/errtrace/main.go

@@ -909,7 +909,6 @@ func (t *walker) deferStmt(n *ast.DeferStmt) {
 		if !ok {
 			return true
 		}
-
 		for i, lhs := range assign.Lhs {
 			ident, ok := lhs.(*ast.Ident)
 			if !ok {
@@ -920,9 +919,33 @@ func (t *walker) deferStmt(n *ast.DeferStmt) {
 				continue // not an error assignment
 			}
 
-			// Assigning to a named error return value.
-			// Wrap the rhs in-place.
-			t.wrapExpr(1, assign.Rhs[i])
+			// Assignment to an error return value.
+			// This will take one of the following forms:
+			//
+			//  (1) x, y, err = f1(), f2(), f3()
+			//  (2) x, y, err = f() // returns multiple values
+			//  (3) x, err, z = f() // returns multiple values
+			//
+			// For (1), we can wrap just the function
+			// that returns the error. (f3 in this case)
+			//
+			// For (2), we can use a WrapN function
+			// to wrap the entire function call.
+			//
+			// For (3), we could use an inline function call,
+			// but that's not implemented yet.
+
+			if i < len(assign.Rhs) && len(assign.Lhs) == len(assign.Rhs) {
+				// Case (1):
+				// Wrap the function that returns the error.
+				t.wrapExpr(1, assign.Rhs[i])
+			} else if i == len(assign.Lhs)-1 && len(assign.Rhs) == 1 {
+				// Case (2):
+				// Wrap the entire function call.
+				t.wrapExpr(len(assign.Lhs), assign.Rhs[0])
+			} else {
+				t.logf(assign.Pos(), "skipping assignment: error is not the last return value")
+			}
 		}
 
 		return true

cmd/errtrace/testdata/golden/tuple_rhs.go

@@ -0,0 +1,20 @@
+package main
+
+import "fmt"
+
+func multipleValueErrAssignment() (err error) {
+	defer func() {
+		_, err = fmt.Println("Hello, World!")
+
+		// Handles too few lhs variables
+		err = fmt.Println("Hello, World!")
+
+		// Handles too many lhs variables
+		_, err, _ = fmt.Println("Hello, World!") // want:"skipping assignment: error is not the last return value"
+
+		// Handles misplaced err
+		err, _ = fmt.Println("Hello, World!") // want:"skipping assignment: error is not the last return value"
+	}()
+
+	return nil
+}

cmd/errtrace/testdata/golden/tuple_rhs.go.golden

@@ -0,0 +1,20 @@
+package main
+
+import "fmt"; import "braces.dev/errtrace"
+
+func multipleValueErrAssignment() (err error) {
+	defer func() {
+		_, err = errtrace.Wrap2(fmt.Println("Hello, World!"))
+
+		// Handles too few lhs variables
+		err = errtrace.Wrap(fmt.Println("Hello, World!"))
+
+		// Handles too many lhs variables
+		_, err, _ = fmt.Println("Hello, World!") // want:"skipping assignment: error is not the last return value"
+
+		// Handles misplaced err
+		err, _ = fmt.Println("Hello, World!") // want:"skipping assignment: error is not the last return value"
+	}()
+
+	return nil
+}