odinfmt: change return formatting and fix binary expression issues.

5 files changed, 80 insertions, 3 deletions

diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 895d2dd..bb0ba0b 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -853,11 +853,33 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 
 		document = enforce_fit_if_do(v.body, document)
 	case ^Return_Stmt:
-		document = cons(document, text("return"))
+		if v.results == nil {
+			document = cons(document, text("return"))
+			break
+		}
+
+		if len(v.results) == 1 {
+			result := v.results[0]
+
+			if paren, is_paren := result.derived.(^ast.Paren_Expr); is_paren {
+				result = paren.expr
+			}
+
+			document = cons(
+				text("return"),
+				if_break("("),
+				break_with(" "),
+				visit_expr(p, result),
+			)
 
-		if v.results != nil {
+			document = nest(document)	
+			document = group(cons(document, if_break(" \\"), break_with(""), if_break(")")))
+		} else {
+			document = cons(document, text("return"))
 			document = cons_with_nopl(document, visit_exprs(p, v.results, {.Add_Comma}))
 		}
+
+
 	case ^Defer_Stmt:
 		document = cons(document, text("defer"))
 		document = cons_with_nopl(document, visit_stmt(p, v.stmt))
@@ -1701,7 +1723,7 @@ visit_binary_expr :: proc(p: ^Printer, binary: ast.Binary_Expr, nested := false)
 		if b, ok := binary.right.derived.(^ast.Binary_Expr); ok {
 			document = cons_with_opl(document, group(nest(visit_binary_expr(p, b^, true))))
 		} else {
-			document = cons_with_opl(document, group(nest(visit_expr(p, binary.right, nested ? .Binary_Expr : .Generic))))
+			document = cons_with_opl(document, group(nest(visit_expr(p, binary.right, .Binary_Expr))))
 		}
 	}
 
diff --git a/tools/odinfmt/tests/.snapshots/assignments.odin b/tools/odinfmt/tests/.snapshots/assignments.odin
index 983087e..54bfbe3 100644
--- a/tools/odinfmt/tests/.snapshots/assignments.odin
+++ b/tools/odinfmt/tests/.snapshots/assignments.odin
@@ -10,6 +10,17 @@ assignments :: proc() {
 		&big.Int{},
 		&big.Int{}
 
+	value +=
+		b4 *
+		grad(
+			seed,
+			[3]i64{rbp.x, rbp.y - i_sign.y * PRIME_Y, rbp.z},
+			[3]f32{ri.x, ri.y + f_sign.y, ri.z},
+		)
+
+	a :=
+		(GRADIENTS_4D[gi] * delta.x + GRADIENTS_4D[gi | 1] * delta.y) +
+		(GRADIENTS_4D[gi | 2] * delta.z + GRADIENTS_4D[gi | 3] * delta.w)
 }
 
 ternary_when_assignment :: proc() {
diff --git a/tools/odinfmt/tests/.snapshots/returns.odin b/tools/odinfmt/tests/.snapshots/returns.odin
new file mode 100644
index 0000000..85996b7
--- /dev/null
+++ b/tools/odinfmt/tests/.snapshots/returns.odin
@@ -0,0 +1,19 @@
+package odinfmt_test
+
+return_with_paren_1 :: proc() {
+	return(
+		GRADIENTS_4D[gi] * delta.x +
+		GRADIENTS_4D[gi | 1] * delta.y +
+		GRADIENTS_4D[gi | 2] * delta.z +
+		GRADIENTS_4D[gi | 3] * delta.w \
+	)
+}
+
+return_without_paren_1 :: proc() {
+	return(
+		GRADIENTS_4D[gi] * delta.x +
+		GRADIENTS_4D[gi | 1] * delta.y +
+		GRADIENTS_4D[gi | 2] * delta.z +
+		GRADIENTS_4D[gi | 3] * delta.w \
+	)
+}
diff --git a/tools/odinfmt/tests/assignments.odin b/tools/odinfmt/tests/assignments.odin
index 53cc9aa..326c0ba 100644
--- a/tools/odinfmt/tests/assignments.odin
+++ b/tools/odinfmt/tests/assignments.odin
@@ -3,6 +3,15 @@ package odinfmt_test
 assignments :: proc() {
 	a, b, c, d, e, f, res := &big.Int{}, &big.Int{}, &big.Int{}, &big.Int{}, &big.Int{}, &big.Int{}, &big.Int{}
 
+	value += b4 * grad(
+            seed,
+            [3]i64{rbp.x, rbp.y - i_sign.y * PRIME_Y, rbp.z},
+            [3]f32{ri.x, ri.y + f_sign.y, ri.z},
+        )
+
+		a := (GRADIENTS_4D[gi] * delta.x +
+			GRADIENTS_4D[gi | 1] *
+			delta.y) + (GRADIENTS_4D[gi | 2] * delta.z + GRADIENTS_4D[gi | 3] * delta.w)
 }
 
 ternary_when_assignment :: proc() {
diff --git a/tools/odinfmt/tests/returns.odin b/tools/odinfmt/tests/returns.odin
new file mode 100644
index 0000000..f4ff400
--- /dev/null
+++ b/tools/odinfmt/tests/returns.odin
@@ -0,0 +1,16 @@
+package odinfmt_test
+
+return_with_paren_1 :: proc() {
+	return(
+        GRADIENTS_4D[gi] * delta.x +
+        GRADIENTS_4D[gi | 1] * delta.y +
+        GRADIENTS_4D[gi | 2] * delta.z +
+        GRADIENTS_4D[gi | 3] * delta.w \
+    )
+}
+
+return_without_paren_1 :: proc() {
+	return GRADIENTS_4D[gi] * delta.x +
+        GRADIENTS_4D[gi | 1] * delta.y +
+        GRADIENTS_4D[gi | 2] * delta.z + GRADIENTS_4D[gi | 3] * delta.w    
+}
\ No newline at end of file