Tidy up dst library use (#4631)

* Use AddComment if there's only one comment

* Use astbuilder.Statements instead of literal

* Use astbuilder.Returns factory method

* Use astbuilder.StatementBlock factory method

* Code gardening

Author: theunrepentantgeek

v2/tools/generator/internal/armconversion/convert_from_arm_function_builder.go

@@ -564,17 +564,11 @@ func (builder *convertFromARMBuilder) buildFlattenedAssignment(
 			astbuilder.IfNotNil(propToCheck, stmts...))
 	}
 
-	comment := []dst.Stmt{
-		&dst.EmptyStmt{
-			Decs: dst.EmptyStmtDecorations{
-				NodeDecs: dst.NodeDecs{
-					End: []string{"// copying flattened property:"},
-				},
-			},
-		},
-	}
+	astbuilder.AddComment(
+		&stmts[0].Decorations().Start,
+		"// copying flattened property:")
 
-	return handleWith(comment, stmts), nil
+	return handleWith(stmts), nil
 }
 
 func (builder *convertFromARMBuilder) propertiesByNameHandler(

v2/tools/generator/internal/armconversion/convert_to_arm_function_builder.go

@@ -135,12 +135,9 @@ func (builder *convertToARMBuilder) functionBodyStatements() ([]dst.Stmt, error)
 		return nil, eris.Wrapf(err, "unable to generate property conversions for %s", builder.methodName)
 	}
 
-	returnStatement := &dst.ReturnStmt{
-		Results: []dst.Expr{
-			dst.NewIdent(builder.resultIdent),
-			astbuilder.Nil(),
-		},
-	}
+	returnStatement := astbuilder.Returns(
+		dst.NewIdent(builder.resultIdent),
+		astbuilder.Nil())
 
 	return astbuilder.Statements(
 		returnIfNil,

v2/tools/generator/internal/astmodel/conversion_function_builder.go

@@ -354,9 +354,9 @@ func IdentityConvertComplexArrayProperty(
 
 		emptySlice := astbuilder.SliceLiteral(destinationTypeExpr)
 		assignEmpty := astbuilder.SimpleAssignment(params.GetDestination(), emptySlice)
-		astbuilder.AddComments(
+		astbuilder.AddComment(
 			&assignEmpty.Decs.Start,
-			[]string{"// Set property to empty map, as this resource is set to serialize all collections explicitly"})
+			"// Set property to empty map, as this resource is set to serialize all collections explicitly")
 		assignEmpty.Decs.Before = dst.NewLine
 
 		ifNil := astbuilder.IfNil(
@@ -471,20 +471,20 @@ func IdentityConvertComplexMapProperty(
 		emptyMap := astbuilder.MakeMap(keyTypeExpr, valueTypeExpr)
 
 		assignEmpty := astbuilder.SimpleAssignment(params.GetDestination(), emptyMap)
-		astbuilder.AddComments(
+		astbuilder.AddComment(
 			&assignEmpty.Decs.Start,
-			[]string{"// Set property to empty map, as this resource is set to serialize all collections explicitly"})
+			"// Set property to empty map, as this resource is set to serialize all collections explicitly")
 		assignEmpty.Decs.Before = dst.NewLine
 
 		result = astbuilder.SimpleIfElse(
 			astbuilder.NotNil(params.GetSource()),
-			[]dst.Stmt{
+			astbuilder.Statements(
 				makeMapStatement,
 				rangeStatement,
-			},
-			[]dst.Stmt{
+			),
+			astbuilder.Statements(
 				assignEmpty,
-			},
+			),
 		)
 	} else {
 		result = astbuilder.IfNotNil(
@@ -500,7 +500,7 @@ func IdentityConvertComplexMapProperty(
 		result.Body.List = append(result.Body.List, params.AssignmentHandler(params.GetDestination(), dst.Clone(destination).(dst.Expr)))
 	}
 
-	return []dst.Stmt{result}, nil
+	return astbuilder.Statements(result), nil
 }
 
 // IdentityAssignTypeName handles conversion for TypeName's that are the same

v2/tools/generator/internal/astmodel/file_definition.go

@@ -227,21 +227,19 @@ func (file *FileDefinition) AsAst() (result *dst.File, err error) {
 			&dst.FuncDecl{
 				Type: &dst.FuncType{Params: &dst.FieldList{}},
 				Name: dst.NewIdent("init"),
-				Body: &dst.BlockStmt{
-					List: []dst.Stmt{
-						&dst.ExprStmt{
-							Decs: dst.ExprStmtDecorations{
-								NodeDecs: dst.NodeDecs{
-									Before: dst.NewLine,
-								},
-							},
-							X: &dst.CallExpr{
-								Fun:  dst.NewIdent("SchemeBuilder.Register"), // HACK
-								Args: exprs,
+				Body: astbuilder.StatementBlock(
+					&dst.ExprStmt{
+						Decs: dst.ExprStmtDecorations{
+							NodeDecs: dst.NodeDecs{
+								Before: dst.NewLine,
 							},
 						},
+						X: &dst.CallExpr{
+							Fun:  dst.NewIdent("SchemeBuilder.Register"), // HACK
+							Args: exprs,
+						},
 					},
-				},
+				),
 			})
 	}
 

v2/tools/generator/internal/codegen/pipeline/report_resource_versions.go

@@ -680,7 +680,9 @@ func (report *ResourceVersionsReport) createItem(
 }
 
 // generateAPILink returns a link to the API definition for the given resource
-func (report *ResourceVersionsReport) generateAPILink(name astmodel.InternalTypeName) string {
+func (report *ResourceVersionsReport) generateAPILink(
+	name astmodel.InternalTypeName,
+) string {
 	crdKind := name.Name()
 	linkTemplate := report.reportConfiguration.ResourceUrlTemplate
 	pathTemplate := report.reportConfiguration.ResourcePathTemplate

v2/tools/generator/internal/codegen/pipeline/resource_registration_file.go

@@ -370,11 +370,7 @@ func (r *ResourceRegistrationFile) createGetKnownStorageTypesFunc(
 		resourceAppendStatements = append(resourceAppendStatements, appendStmt)
 	}
 
-	returnStmt := &dst.ReturnStmt{
-		Results: []dst.Expr{
-			resultIdent,
-		},
-	}
+	returnStmt := astbuilder.Returns(resultIdent)
 
 	body := astbuilder.Statements(resultVar, resourceAppendStatements, returnStmt)
 
@@ -494,11 +490,7 @@ func (r *ResourceRegistrationFile) createCreateSchemeFunc(codeGenerationContext
 		groupVersionAssignments = append(groupVersionAssignments, groupSchemeAssign)
 	}
 
-	returnStmt := &dst.ReturnStmt{
-		Results: []dst.Expr{
-			dst.NewIdent(scheme),
-		},
-	}
+	returnStmt := astbuilder.Returns(dst.NewIdent(scheme))
 
 	body := astbuilder.Statements(initSchemeVar, clientGoSchemeAssign, groupVersionAssignments, returnStmt)
 

v2/tools/generator/internal/conversions/property_conversions.go

@@ -1337,14 +1337,14 @@ func assignArrayFromArray(
 			astbuilder.MakeSlice(destinationArrayExpr, astbuilder.CallFunc("len", actualReader)))
 
 		writeToElement := func(expr dst.Expr) []dst.Stmt {
-			return []dst.Stmt{
+			return astbuilder.Statements(
 				astbuilder.SimpleAssignment(
 					&dst.IndexExpr{
 						X:     dst.NewIdent(tempId),
 						Index: dst.NewIdent(indexId),
 					},
 					expr),
-			}
+			)
 		}
 
 		avoidAliasing := astbuilder.ShortDeclaration(itemId, dst.NewIdent(itemId))
@@ -1499,14 +1499,14 @@ func assignMapFromMap(
 				astbuilder.CallFunc("len", actualReader)))
 
 		assignToItem := func(expr dst.Expr) []dst.Stmt {
-			return []dst.Stmt{
+			return astbuilder.Statements(
 				astbuilder.SimpleAssignment(
 					&dst.IndexExpr{
 						X:     dst.NewIdent(tempId),
 						Index: dst.NewIdent(keyId),
 					},
 					expr),
-			}
+			)
 		}
 
 		avoidAliasing := astbuilder.ShortDeclaration(itemId, dst.NewIdent(itemId))
@@ -1649,9 +1649,8 @@ func assignUserAssignedIdentityMapFromArray(
 		uaiBuilder.AddField("Reference", dst.NewIdent(intermediateDestination))
 
 		writeToElement := func(expr dst.Expr) []dst.Stmt {
-			return []dst.Stmt{
-				astbuilder.ShortDeclaration(intermediateDestination, expr),
-			}
+			return astbuilder.Statements(
+				astbuilder.ShortDeclaration(intermediateDestination, expr))
 		}
 
 		//	for key, _ := range source.UserAssignedIdentities {

v2/tools/generator/internal/conversions/writable_conversion_endpoint.go

@@ -40,11 +40,10 @@ func NewWritableConversionEndpointWritingProperty(
 	return &WritableConversionEndpoint{
 		endpoint: endpoint,
 		writer: func(destination dst.Expr, value dst.Expr) []dst.Stmt {
-			return []dst.Stmt{
+			return astbuilder.Statements(
 				astbuilder.SimpleAssignment(
 					astbuilder.Selector(destination, name),
-					value),
-			}
+					value))
 		},
 		description: fmt.Sprintf("write to property %s", name),
 	}
@@ -59,11 +58,10 @@ func NewWritableConversionEndpointWritingPropertyBagMember(
 	return &WritableConversionEndpoint{
 		endpoint: NewTypedConversionEndpoint(NewPropertyBagMemberType(itemType), itemName),
 		writer: func(destination dst.Expr, value dst.Expr) []dst.Stmt {
-			return []dst.Stmt{
+			return astbuilder.Statements(
 				astbuilder.SimpleAssignment(
 					astbuilder.Selector(destination, itemName),
-					value),
-			}
+					value))
 		},
 		description: fmt.Sprintf("write %s to property bag", itemName),
 	}

v2/tools/generator/internal/functions/conditions_functions.go

@@ -38,9 +38,9 @@ func GetConditionsFunction(
 		Name:          methodName,
 		ReceiverIdent: receiverIdent,
 		ReceiverType:  astbuilder.PointerTo(receiverExpr),
-		Body: []dst.Stmt{
+		Body: astbuilder.Statements(
 			astbuilder.Returns(astbuilder.Selector(status, astmodel.ConditionsProperty)),
-		},
+		),
 	}
 
 	fn.AddComments("returns the conditions of the resource")
@@ -79,9 +79,9 @@ func SetConditionsFunction(
 		Name:          methodName,
 		ReceiverIdent: receiverIdent,
 		ReceiverType:  astbuilder.PointerTo(receiverExpr),
-		Body: []dst.Stmt{
+		Body: astbuilder.Statements(
 			astbuilder.QualifiedAssignment(status, "Conditions", token.ASSIGN, dst.NewIdent(conditionsParameterName)),
-		},
+		),
 	}
 
 	conditionsTypeExpr, err := astmodel.ConditionsType.AsTypeExpr(codeGenerationContext)

v2/tools/generator/internal/functions/kubernetes_admissions_validator.go

@@ -441,9 +441,9 @@ func (v *ValidatorBuilder) localValidationFuncBody(
 	}
 
 	if len(elements) == 0 {
-		return []dst.Stmt{
+		return astbuilder.Statements(
 			astbuilder.Returns(astbuilder.Nil()),
-		}, nil
+		), nil
 	}
 
 	returnStmt := astbuilder.Returns(&dst.CompositeLit{
@@ -493,11 +493,9 @@ func (v *ValidatorBuilder) makeLocalValidationElement(
 					},
 				},
 				Type: getValidationFuncType(kind, codeGenerationContext),
-				Body: &dst.BlockStmt{
-					List: []dst.Stmt{
-						astbuilder.Returns(astbuilder.CallQualifiedFunc(receiverIdent, validation.name)),
-					},
-				},
+				Body: astbuilder.StatementBlock(
+					astbuilder.Returns(astbuilder.CallQualifiedFunc(receiverIdent, validation.name)),
+				),
 			}, nil
 		}
 	}

v2/tools/generator/internal/functions/one_of_json_marshal_function.go

@@ -106,9 +106,7 @@ func (f *OneOfJSONMarshalFunction) AsFunc(
 		statements = append(statements, ifStatement)
 	}
 
-	finalReturnStatement := &dst.ReturnStmt{
-		Results: []dst.Expr{astbuilder.Nil(), astbuilder.Nil()},
-	}
+	finalReturnStatement := astbuilder.Returns(astbuilder.Nil(), astbuilder.Nil())
 
 	fn := &astbuilder.FuncDetails{
 		Name:          f.Name(),

v2/tools/generator/internal/functions/one_of_json_unmarshal_function.go

@@ -77,7 +77,7 @@ func (f *OneOfJSONUnmarshalFunction) AsFunc(
 	discrimName := "discriminator"
 	errName := "err"
 
-	statements := []dst.Stmt{
+	statements := astbuilder.Statements(
 		astbuilder.LocalVariableDeclaration(mapName, &dst.MapType{Key: dst.NewIdent("string"), Value: dst.NewIdent("interface{}")}, ""),
 		astbuilder.ShortDeclaration(errName,
 			astbuilder.CallQualifiedFunc(jsonPackage, "Unmarshal", dst.NewIdent(paramName), astbuilder.AddrOf(dst.NewIdent(mapName)))),
@@ -86,7 +86,7 @@ func (f *OneOfJSONUnmarshalFunction) AsFunc(
 			X:     dst.NewIdent(mapName),
 			Index: astbuilder.StringLiteral(discrimJSONName),
 		}),
-	}
+	)
 
 	// must order this for consistent output
 	values := make([]string, 0, len(valuesMapping))

v2/tools/generator/internal/interfaces/kubernetes_resource_interface.go

@@ -627,13 +627,13 @@ func setStringAzureNameFunction(
 		Name:          methodName,
 		ReceiverIdent: receiverIdent,
 		ReceiverType:  astbuilder.PointerTo(receiverTypeExpr),
-		Body: []dst.Stmt{
+		Body: astbuilder.Statements(
 			astbuilder.QualifiedAssignment(
 				dst.NewIdent(receiverIdent),
 				astmodel.AzureNameProperty,
 				token.ASSIGN,
 				dst.NewIdent("azureName")),
-		},
+		),
 	}
 
 	fn.AddComments("sets the Azure name of the resource")

v2/tools/generator/internal/testcases/json_serialization_test_case.go

@@ -487,7 +487,9 @@ func (o *JSONSerializationTestCase) createGeneratorMethod(ctx *astmodel.CodeGene
 			Ellipsis: true,
 		}
 
-		oneOfStmts = []dst.Stmt{possibleGenerators, initGopters}
+		oneOfStmts = astbuilder.Statements(
+			possibleGenerators,
+			initGopters)
 	}
 
 	// If we have already cached our builder, return it immediately