diff --git a/_fixtures/chained_args.go b/_fixtures/chained_args.go
new file mode 100644
index 0000000..4975291
--- /dev/null
+++ b/_fixtures/chained_args.go
@@ -0,0 +1,22 @@
+package fixtures
+
+import "strings"
+
+type Logger struct{}
+
+func (l Logger) Str(k, v string) Logger { return l }
+func (l Logger) Msgf(format string, args ...interface{}) {}
+
+func Map[T, U any](collection []T, fn func(T, int) U) []U {
+	return nil
+}
+
+func testChainedArgs() {
+	messages := []string{}
+	l := Logger{}
+	
+	// Long argument inside a chained method call - should split the arguments
+	l.Str("key", strings.Join(Map(messages, func(message string, _ int) string { return message + message + message }), ",")).
+		Msgf("Message with format %s", "arg")
+}
+
diff --git a/_fixtures/chained_args__exp.go b/_fixtures/chained_args__exp.go
new file mode 100644
index 0000000..91e120f
--- /dev/null
+++ b/_fixtures/chained_args__exp.go
@@ -0,0 +1,21 @@
+package fixtures
+
+import "strings"
+
+type Logger struct{}
+
+func (l Logger) Str(k, v string) Logger                  { return l }
+func (l Logger) Msgf(format string, args ...interface{}) {}
+
+func Map[T, U any](collection []T, fn func(T, int) U) []U {
+	return nil
+}
+
+func testChainedArgs() {
+	messages := []string{}
+	l := Logger{}
+
+	// Long argument inside a chained method call - should split the arguments
+	l.Str("key", strings.Join(Map(messages, func(message string, _ int) string { return message + message + message }), ",")).
+		Msgf("Message with format %s", "arg")
+}
diff --git a/_fixtures/chained_calls__exp.go b/_fixtures/chained_calls__exp.go
index 7fa5d04..84026e5 100644
--- a/_fixtures/chained_calls__exp.go
+++ b/_fixtures/chained_calls__exp.go
@@ -20,9 +20,16 @@ func ChainedCalls() {
 	NewChain().ChainCall(
 		"a really really really really really long argument4",
 		"another really really really really really long argument4",
-		fmt.Sprintf("%v", "this is a long method"),
+		fmt.Sprintf(
+			"%v",
+			"this is a long method",
+		),
 	).
-		ChainCall("a really really really really really long argument5", "another really really really really really long argument5", "a third really really really really really long argument5").
+		ChainCall(
+			"a really really really really really long argument5",
+			"another really really really really really long argument5",
+			"a third really really really really really long argument5",
+		).
 		ChainCall("a", "b", fmt.Sprintf("%v", "this is a long method"))
 	NewChain().ChainCall("a", "b", "c").ChainCall("d", "e", "f")
 	NewChain().ChainCall("a", "b", "c").
diff --git a/_fixtures/composite_lit.go b/_fixtures/composite_lit.go
new file mode 100644
index 0000000..7e6e430
--- /dev/null
+++ b/_fixtures/composite_lit.go
@@ -0,0 +1,24 @@
+package fixtures
+
+type Request struct {
+	Entity       string
+	Fields       []string
+	VersionCheck bool
+}
+
+func update(ctx string, req Request) error {
+	return nil
+}
+
+func testCompositeLit() {
+	// Struct literal with long field line that should be split
+	update(
+		"ctx",
+		Request{
+			Entity: "enrollment", Fields: []string{"field1", "field2", "field3", "field4"}, VersionCheck: true,
+		},
+	)
+
+	// Struct literal with nested long slice
+	update("ctx", Request{Entity: "test", Fields: []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}, VersionCheck: false})
+}
diff --git a/_fixtures/composite_lit__exp.go b/_fixtures/composite_lit__exp.go
new file mode 100644
index 0000000..7f5b3d9
--- /dev/null
+++ b/_fixtures/composite_lit__exp.go
@@ -0,0 +1,33 @@
+package fixtures
+
+type Request struct {
+	Entity       string
+	Fields       []string
+	VersionCheck bool
+}
+
+func update(ctx string, req Request) error {
+	return nil
+}
+
+func testCompositeLit() {
+	// Struct literal with long field line that should be split
+	update(
+		"ctx",
+		Request{
+			Entity:       "enrollment",
+			Fields:       []string{"field1", "field2", "field3", "field4"},
+			VersionCheck: true,
+		},
+	)
+
+	// Struct literal with nested long slice
+	update(
+		"ctx",
+		Request{
+			Entity:       "test",
+			Fields:       []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"},
+			VersionCheck: false,
+		},
+	)
+}
diff --git a/_fixtures/for_init.go b/_fixtures/for_init.go
new file mode 100644
index 0000000..79f946b
--- /dev/null
+++ b/_fixtures/for_init.go
@@ -0,0 +1,18 @@
+package fixtures
+
+func getValue(a, b, c, d, e string) int {
+	return 0
+}
+
+func testForInit() {
+	// For statement with long init clause
+	for i := getValue("argument1", "argument2", "argument3", "argument4", "argument5"); i < 10; i++ {
+		println(i)
+	}
+
+	// For with long condition
+	for j := 0; j < getValue("value1", "value2", "value3", "value4", "value5"); j++ {
+		println(j)
+	}
+}
+
diff --git a/_fixtures/for_init__exp.go b/_fixtures/for_init__exp.go
new file mode 100644
index 0000000..e6e73ce
--- /dev/null
+++ b/_fixtures/for_init__exp.go
@@ -0,0 +1,23 @@
+package fixtures
+
+func getValue(a, b, c, d, e string) int {
+	return 0
+}
+
+func testForInit() {
+	// For statement with long init clause
+	for i := getValue(
+		"argument1",
+		"argument2",
+		"argument3",
+		"argument4",
+		"argument5",
+	); i < 10; i++ {
+		println(i)
+	}
+
+	// For with long condition
+	for j := 0; j < getValue("value1", "value2", "value3", "value4", "value5"); j++ {
+		println(j)
+	}
+}
diff --git a/_fixtures/func_lit_body.go b/_fixtures/func_lit_body.go
new file mode 100644
index 0000000..aa940e5
--- /dev/null
+++ b/_fixtures/func_lit_body.go
@@ -0,0 +1,24 @@
+package fixtures
+
+type LongTypeName struct {
+	Value int
+}
+
+func (l LongTypeName) Compare(other LongTypeName) int {
+	return l.Value - other.Value
+}
+
+func SortFunc[T any](collection []T, fn func(T, T) int) {}
+
+func testFuncLitBody() {
+	// Function literal with long body that should be expanded
+	items := []LongTypeName{}
+	SortFunc(
+		items,
+		func(a, b LongTypeName) int { return a.Compare(b) + a.Compare(b) + a.Compare(b) + a.Compare(b) },
+	)
+
+	// Single line function literal that fits - should stay as is
+	SortFunc(items, func(a, b LongTypeName) int { return a.Compare(b) })
+}
+
diff --git a/_fixtures/func_lit_body__exp.go b/_fixtures/func_lit_body__exp.go
new file mode 100644
index 0000000..d3be176
--- /dev/null
+++ b/_fixtures/func_lit_body__exp.go
@@ -0,0 +1,25 @@
+package fixtures
+
+type LongTypeName struct {
+	Value int
+}
+
+func (l LongTypeName) Compare(other LongTypeName) int {
+	return l.Value - other.Value
+}
+
+func SortFunc[T any](collection []T, fn func(T, T) int) {}
+
+func testFuncLitBody() {
+	// Function literal with long body that should be expanded
+	items := []LongTypeName{}
+	SortFunc(
+		items,
+		func(a, b LongTypeName) int {
+			return a.Compare(b) + a.Compare(b) + a.Compare(b) + a.Compare(b)
+		},
+	)
+
+	// Single line function literal that fits - should stay as is
+	SortFunc(items, func(a, b LongTypeName) int { return a.Compare(b) })
+}
diff --git a/_fixtures/if_init.go b/_fixtures/if_init.go
new file mode 100644
index 0000000..498117d
--- /dev/null
+++ b/_fixtures/if_init.go
@@ -0,0 +1,20 @@
+package fixtures
+
+func doSomething(a, b, c, d, e string) error {
+	return nil
+}
+
+func testIfInit() {
+	// If statement with long init clause
+	if err := doSomething("argument1", "argument2", "argument3", "argument4", "argument5"); err != nil {
+		println(err)
+	}
+
+	// Nested if with long init
+	if x := doSomething("value1", "value2", "value3", "value4", "value5"); x == nil {
+		if y := doSomething("nested1", "nested2", "nested3", "nested4", "nested5"); y == nil {
+			println("both nil")
+		}
+	}
+}
+
diff --git a/_fixtures/if_init__exp.go b/_fixtures/if_init__exp.go
new file mode 100644
index 0000000..afece17
--- /dev/null
+++ b/_fixtures/if_init__exp.go
@@ -0,0 +1,25 @@
+package fixtures
+
+func doSomething(a, b, c, d, e string) error {
+	return nil
+}
+
+func testIfInit() {
+	// If statement with long init clause
+	if err := doSomething(
+		"argument1",
+		"argument2",
+		"argument3",
+		"argument4",
+		"argument5",
+	); err != nil {
+		println(err)
+	}
+
+	// Nested if with long init
+	if x := doSomething("value1", "value2", "value3", "value4", "value5"); x == nil {
+		if y := doSomething("nested1", "nested2", "nested3", "nested4", "nested5"); y == nil {
+			println("both nil")
+		}
+	}
+}
diff --git a/_fixtures/switch_init.go b/_fixtures/switch_init.go
new file mode 100644
index 0000000..baca312
--- /dev/null
+++ b/_fixtures/switch_init.go
@@ -0,0 +1,18 @@
+package fixtures
+
+func getType(a, b, c, d, e, f string) int {
+	return 0
+}
+
+func testSwitchInit() {
+	// Switch statement with long init clause
+	switch v := getType("argument1", "argument2", "argument3", "argument4", "argument5", "argument6"); v {
+	case 1:
+		println("one")
+	case 2:
+		println("two")
+	default:
+		println("other")
+	}
+}
+
diff --git a/_fixtures/switch_init__exp.go b/_fixtures/switch_init__exp.go
new file mode 100644
index 0000000..6a9b278
--- /dev/null
+++ b/_fixtures/switch_init__exp.go
@@ -0,0 +1,24 @@
+package fixtures
+
+func getType(a, b, c, d, e, f string) int {
+	return 0
+}
+
+func testSwitchInit() {
+	// Switch statement with long init clause
+	switch v := getType(
+		"argument1",
+		"argument2",
+		"argument3",
+		"argument4",
+		"argument5",
+		"argument6",
+	); v {
+	case 1:
+		println("one")
+	case 2:
+		println("two")
+	default:
+		println("other")
+	}
+}
diff --git a/annotation.go b/annotation.go
index 0790a30..96e2ad2 100644
--- a/annotation.go
+++ b/annotation.go
@@ -79,6 +79,12 @@ func HasAnnotationRecursive(node dst.Node) bool {
 				return true
 			}
 		}
+	case *dst.CompositeLit:
+		for _, elt := range n.Elts {
+			if HasAnnotation(elt) {
+				return true
+			}
+		}
 	}
 
 	return false
diff --git a/shortener.go b/shortener.go
index 4e8d9f3..9bca03d 100644
--- a/shortener.go
+++ b/shortener.go
@@ -241,6 +241,10 @@ func (s *Shortener) annotateLongLines(lines []string) ([]string, int) {
 				// Replace annotation with new length
 				annotatedLines[len(annotatedLines)-1] = CreateAnnotation(length)
 				linesToShorten++
+			} else {
+				// Line is still too long with no progress, but keep trying
+				// (maxRounds will prevent infinite loops)
+				linesToShorten++
 			}
 		} else if !s.isComment(line) && length > s.config.MaxLen {
 			annotatedLines = append(
@@ -454,12 +458,43 @@ func (s *Shortener) formatStmt(stmt dst.Stmt) {
 	case *dst.ExprStmt:
 		s.formatExpr(st.X, shouldShorten, false)
 	case *dst.ForStmt:
+		if st.Init != nil {
+			s.formatStmt(st.Init)
+			if shouldShorten {
+				if assignStmt, ok := st.Init.(*dst.AssignStmt); ok {
+					for _, expr := range assignStmt.Rhs {
+						s.formatExpr(expr, true, false)
+					}
+				}
+			}
+		}
+		if st.Cond != nil {
+			s.formatExpr(st.Cond, shouldShorten, false)
+		}
+		if st.Post != nil {
+			s.formatStmt(st.Post)
+		}
 		s.formatStmt(st.Body)
 	case *dst.GoStmt:
 		s.formatExpr(st.Call, shouldShorten, false)
 	case *dst.IfStmt:
+		if st.Init != nil {
+			s.formatStmt(st.Init)
+			// If the init statement is an assignment containing a call expression that
+			// should be shortened, mark it for shortening
+			if shouldShorten {
+				if assignStmt, ok := st.Init.(*dst.AssignStmt); ok {
+					for _, expr := range assignStmt.Rhs {
+						s.formatExpr(expr, true, false)
+					}
+				}
+			}
+		}
 		s.formatExpr(st.Cond, shouldShorten, false)
 		s.formatStmt(st.Body)
+		if st.Else != nil {
+			s.formatStmt(st.Else)
+		}
 	case *dst.RangeStmt:
 		s.formatStmt(st.Body)
 	case *dst.ReturnStmt:
@@ -469,6 +504,19 @@ func (s *Shortener) formatStmt(stmt dst.Stmt) {
 	case *dst.SelectStmt:
 		s.formatStmt(st.Body)
 	case *dst.SwitchStmt:
+		if st.Init != nil {
+			s.formatStmt(st.Init)
+			if shouldShorten {
+				if assignStmt, ok := st.Init.(*dst.AssignStmt); ok {
+					for _, expr := range assignStmt.Rhs {
+						s.formatExpr(expr, true, false)
+					}
+				}
+			}
+		}
+		if st.Tag != nil {
+			s.formatExpr(st.Tag, shouldShorten, false)
+		}
 		s.formatStmt(st.Body)
 	default:
 		if shouldShorten {
@@ -498,7 +546,7 @@ func (s *Shortener) formatExpr(expr dst.Expr, force bool, isChain bool) {
 			s.formatExpr(e.Y, shouldShorten, isChain)
 		}
 	case *dst.CallExpr:
-		_, ok := e.Fun.(*dst.SelectorExpr)
+		selectorExpr, ok := e.Fun.(*dst.SelectorExpr)
 
 		if ok &&
 			s.config.ChainSplitDots &&
@@ -506,8 +554,19 @@ func (s *Shortener) formatExpr(expr dst.Expr, force bool, isChain bool) {
 			(isChain || s.chainLength(e) > 1) {
 			e.Decorations().After = dst.NewLine
 
-			for _, arg := range e.Args {
-				s.formatExpr(arg, false, true)
+			// If this specific call is annotated (line still too long after chain split),
+			// also split the arguments. The annotation may be on the CallExpr, the
+			// SelectorExpr, or the method name Ident.
+			shortenArgs := HasAnnotation(e) || HasAnnotation(selectorExpr) ||
+				HasAnnotation(selectorExpr.Sel)
+			for a, arg := range e.Args {
+				if shortenArgs {
+					if a == 0 {
+						arg.Decorations().Before = dst.NewLine
+					}
+					arg.Decorations().After = dst.NewLine
+				}
+				s.formatExpr(arg, shortenArgs, true)
 			}
 
 			s.formatExpr(e.Fun, shouldShorten, true)
@@ -528,7 +587,11 @@ func (s *Shortener) formatExpr(expr dst.Expr, force bool, isChain bool) {
 			s.formatExpr(e.Fun, shouldShorten, isChain)
 		}
 	case *dst.CompositeLit:
-		if shouldShorten {
+		// Check if we should shorten - either because the parent told us to,
+		// or because one of the elements has an annotation (meaning the line
+		// with the elements is too long)
+		shortenElements := shouldShorten || HasAnnotationRecursive(e)
+		if shortenElements {
 			for i, element := range e.Elts {
 				if i == 0 {
 					element.Decorations().Before = dst.NewLine
@@ -541,6 +604,16 @@ func (s *Shortener) formatExpr(expr dst.Expr, force bool, isChain bool) {
 			s.formatExpr(element, false, isChain)
 		}
 	case *dst.FuncLit:
+		// If the function literal line is too long, expand the body
+		if shouldShorten && e.Body != nil && len(e.Body.List) > 0 {
+			// Put each statement in the body on its own line
+			for i, stmt := range e.Body.List {
+				if i == 0 {
+					stmt.Decorations().Before = dst.NewLine
+				}
+				stmt.Decorations().After = dst.NewLine
+			}
+		}
 		s.formatStmt(e.Body)
 	case *dst.FuncType:
 		if shouldShorten {