Code gardening for code generation (#1672)

* Allocate locals only when we know we need them * Use helper method * Eliminate EnsureStatementBlock() and subsume into StatementBlock() * Allow SimpleIfElse() to handle multiple statements per branch * Use simplified SimpleIfElse() * Change SimpleAssignment() to always do assignment * Introduce SimpleDeclaration to always do variable := definition * fixup! Change SimpleAssignment() to always do assignment * fixup! Change SimpleAssignment() to always do assignment * Introduce SetVariable() * Modify SimpleDeclaration() to take string for variable name * Rename method to ShortDeclaration() * Renamed to `AssignmentStatement()`
Azure · Aug 1, 2021 · 053afb2 · 053afb2
1 parent f317fde
commit 053afb2
Show file tree

Hide file tree

Showing 15 changed files with 111 additions and 136 deletions.
diff --git a/hack/generator/pkg/astbuilder/assignments.go b/hack/generator/pkg/astbuilder/assignments.go
@@ -6,26 +6,61 @@
 package astbuilder
 
 import (
+	"fmt"
 	"go/token"
 
 	"github.com/dave/dst"
 )
 
 // SimpleAssignment performs a simple assignment like:
-//     <lhs> := <rhs>       // tok = token.DEFINE
-// or  <lhs> = <rhs>        // tok = token.ASSIGN
-func SimpleAssignment(lhs dst.Expr, tok token.Token, rhs dst.Expr) *dst.AssignStmt {
+//
+//     <lhs> = <rhs>
+//
+// See also ShortDeclaration
+func SimpleAssignment(lhs dst.Expr, rhs dst.Expr) *dst.AssignStmt {
 	return &dst.AssignStmt{
 		Lhs: []dst.Expr{
 			dst.Clone(lhs).(dst.Expr),
 		},
-		Tok: tok,
+		Tok: token.ASSIGN,
 		Rhs: []dst.Expr{
 			dst.Clone(rhs).(dst.Expr),
 		},
 	}
 }
 
+// ShortDeclaration performs a simple assignment like:
+//
+//     <id> := <rhs>
+//
+// Method naming inspired by https://tour.golang.org/basics/10
+func ShortDeclaration(id string, rhs dst.Expr) *dst.AssignStmt {
+	return &dst.AssignStmt{
+		Lhs: []dst.Expr{
+			dst.NewIdent(id),
+		},
+		Tok: token.DEFINE,
+		Rhs: []dst.Expr{
+			dst.Clone(rhs).(dst.Expr),
+		},
+	}
+}
+
+// AssignmentStatement allows for either variable declaration or assignment by passing the required token
+// Only token.DEFINE and token.ASSIGN are supported, other values will panic.
+// Use SimpleAssignment or ShortDeclaration if possible; use this method only if you must.
+func AssignmentStatement(lhs dst.Expr, tok token.Token, rhs dst.Expr) *dst.AssignStmt {
+	if tok != token.ASSIGN && tok != token.DEFINE {
+		panic(fmt.Sprintf("token %q not supported in VariableAssignment", tok))
+	}
+
+	return &dst.AssignStmt{
+		Lhs: Expressions(lhs),
+		Tok: tok,
+		Rhs: Expressions(rhs),
+	}
+}
+
 // QualifiedAssignment performs a simple assignment like:
 //     <lhs>.<lhsSel> := <rhs>       // tok = token.DEFINE
 // or  <lhs>.<lhsSel> = <rhs>        // tok = token.ASSIGN

diff --git a/hack/generator/pkg/astbuilder/builder.go b/hack/generator/pkg/astbuilder/builder.go
@@ -386,20 +386,19 @@ func Nil() *dst.Ident {
 }
 
 // StatementBlock generates a block containing the supplied statements
+// If we're given a single statement that's already a block, we won't double wrap it
 func StatementBlock(statements ...dst.Stmt) *dst.BlockStmt {
-	return &dst.BlockStmt{
-		List: Statements(statements),
-	}
-}
+	stmts := Statements(statements)
 
-// EnsureStatementBlock wraps any statement into a block safely
-// (without double wrapping an existing block)
-func EnsureStatementBlock(statement dst.Stmt) *dst.BlockStmt {
-	if block, ok := statement.(*dst.BlockStmt); ok {
-		return block
+	if len(stmts) == 1 {
+		if block, ok := stmts[0].(*dst.BlockStmt); ok {
+			return block
+		}
 	}
 
-	return StatementBlock(statement)
+	return &dst.BlockStmt{
+		List: stmts,
+	}
 }
 
 // Statements creates a sequence of statements from the provided values, each of which may be a

diff --git a/hack/generator/pkg/astbuilder/conditions.go b/hack/generator/pkg/astbuilder/conditions.go
@@ -19,11 +19,11 @@ import (
 //     <falseBranch>
 // }
 //
-func SimpleIfElse(condition dst.Expr, trueBranch dst.Stmt, falseBranch dst.Stmt) *dst.IfStmt {
+func SimpleIfElse(condition dst.Expr, trueBranch []dst.Stmt, falseBranch []dst.Stmt) *dst.IfStmt {
 	result := &dst.IfStmt{
 		Cond: condition,
-		Body: EnsureStatementBlock(trueBranch),
-		Else: EnsureStatementBlock(falseBranch),
+		Body: StatementBlock(trueBranch...),
+		Else: StatementBlock(falseBranch...),
 	}
 
 	return result

diff --git a/hack/generator/pkg/astbuilder/lists.go b/hack/generator/pkg/astbuilder/lists.go
@@ -6,8 +6,9 @@
 package astbuilder
 
 import (
-	"github.com/dave/dst"
 	"go/token"
+
+	"github.com/dave/dst"
 )
 
 // MakeList returns the call expression for making a slice
@@ -31,7 +32,6 @@ func MakeList(listType dst.Expr, len dst.Expr) *dst.CallExpr {
 func AppendList(lhs dst.Expr, rhs dst.Expr) dst.Stmt {
 	return SimpleAssignment(
 		dst.Clone(lhs).(dst.Expr),
-		token.ASSIGN,
 		CallFunc("append", dst.Clone(lhs).(dst.Expr), dst.Clone(rhs).(dst.Expr)))
 }
 

diff --git a/hack/generator/pkg/astbuilder/maps.go b/hack/generator/pkg/astbuilder/maps.go
@@ -6,8 +6,9 @@
 package astbuilder
 
 import (
-	"github.com/dave/dst"
 	"go/token"
+
+	"github.com/dave/dst"
 )
 
 // MakeMap returns the call expression for making a map
@@ -36,7 +37,6 @@ func InsertMap(mapExpr dst.Expr, key dst.Expr, rhs dst.Expr) *dst.AssignStmt {
 			X:     dst.Clone(mapExpr).(dst.Expr),
 			Index: dst.Clone(key).(dst.Expr),
 		},
-		token.ASSIGN,
 		dst.Clone(rhs).(dst.Expr))
 }
 

diff --git a/hack/generator/pkg/astmodel/armconversion/convert_from_arm_function_builder.go b/hack/generator/pkg/astmodel/armconversion/convert_from_arm_function_builder.go
@@ -431,7 +431,7 @@ func (builder *convertFromARMBuilder) convertComplexTypeNameProperty(conversionB
 	results = append(results, newVariable)
 	results = append(
 		results,
-		astbuilder.SimpleAssignment(
+		astbuilder.AssignmentStatement(
 			dst.NewIdent("err"),
 			tok,
 			astbuilder.CallQualifiedFunc(
@@ -442,7 +442,6 @@ func (builder *convertFromARMBuilder) convertComplexTypeNameProperty(conversionB
 			results,
 			astbuilder.SimpleAssignment(
 				params.GetDestination(),
-				token.ASSIGN,
 				dst.NewIdent(propertyLocalVar)))
 	} else {
 		results = append(

diff --git a/hack/generator/pkg/astmodel/conversion_function_builder.go b/hack/generator/pkg/astmodel/conversion_function_builder.go
@@ -210,7 +210,6 @@ func IdentityConvertComplexOptionalProperty(builder *ConversionFunctionBuilder,
 		innerStatements,
 		astbuilder.SimpleAssignment(
 			params.GetDestination(),
-			token.ASSIGN,
 			astbuilder.AddrOf(dst.NewIdent(tempVarIdent))))
 
 	result := &dst.IfStmt{
@@ -345,8 +344,8 @@ func IdentityConvertComplexMapProperty(builder *ConversionFunctionBuilder, param
 	keyTypeAst := destinationType.KeyType().AsType(builder.CodeGenerationContext)
 	valueTypeAst := destinationType.ValueType().AsType(builder.CodeGenerationContext)
 
-	makeMapStatement := astbuilder.SimpleAssignment(
-		dst.Clone(destination).(dst.Expr),
+	makeMapStatement := astbuilder.AssignmentStatement(
+		destination,
 		makeMapToken,
 		astbuilder.MakeMap(keyTypeAst, valueTypeAst))
 	rangeStatement := &dst.RangeStmt{
@@ -575,12 +574,12 @@ func IdentityDeepCopyJSON(builder *ConversionFunctionBuilder, params ConversionP
 
 // AssignmentHandlerDefine is an assignment handler for definitions, using :=
 func AssignmentHandlerDefine(lhs dst.Expr, rhs dst.Expr) dst.Stmt {
-	return astbuilder.SimpleAssignment(lhs, token.DEFINE, rhs)
+	return astbuilder.AssignmentStatement(lhs, token.DEFINE, rhs)
 }
 
 // AssignmentHandlerAssign is an assignment handler for standard assignments to existing variables, using =
 func AssignmentHandlerAssign(lhs dst.Expr, rhs dst.Expr) dst.Stmt {
-	return astbuilder.SimpleAssignment(lhs, token.ASSIGN, rhs)
+	return astbuilder.SimpleAssignment(lhs, rhs)
 }
 
 // CreateLocal creates an unused local variable name.

diff --git a/hack/generator/pkg/astmodel/kubernetes_admissions_defaults.go b/hack/generator/pkg/astmodel/kubernetes_admissions_defaults.go
@@ -60,7 +60,6 @@ func defaultAzureNameFunction(k *resourceFunction, codeGenerationContext *CodeGe
 				Body: astbuilder.StatementBlock(
 					astbuilder.SimpleAssignment(
 						azureNameProp,
-						token.ASSIGN,
 						nameProp)),
 			},
 		},

diff --git a/hack/generator/pkg/astmodel/kubernetes_admissions_validator.go b/hack/generator/pkg/astmodel/kubernetes_admissions_validator.go
@@ -266,7 +266,7 @@ func (v *ValidatorBuilder) validateBody(codeGenerationContext *CodeGenerationCon
 		Tok:   token.DEFINE,
 		Body: &dst.BlockStmt{
 			List: []dst.Stmt{
-				astbuilder.SimpleAssignment(dst.NewIdent("err"), token.DEFINE, astbuilder.CallFunc(validationIdent, args...)),
+				astbuilder.ShortDeclaration("err", astbuilder.CallFunc(validationIdent, args...)),
 				astbuilder.CheckErrorAndSingleStatement(astbuilder.AppendList(dst.NewIdent(errsIdent), dst.NewIdent("err"))),
 			},
 		},
@@ -279,9 +279,8 @@ func (v *ValidatorBuilder) validateBody(codeGenerationContext *CodeGenerationCon
 	appendFuncCall.Ellipsis = true
 
 	body := []dst.Stmt{
-		astbuilder.SimpleAssignment(
-			dst.NewIdent(validationsIdent),
-			token.DEFINE,
+		astbuilder.ShortDeclaration(
+			validationsIdent,
 			astbuilder.CallQualifiedFunc(receiverIdent, implFunctionName)),
 		astbuilder.AssignToInterface(tempVarIdent, dst.NewIdent(receiverIdent)),
 		&dst.IfStmt{
@@ -295,7 +294,6 @@ func (v *ValidatorBuilder) validateBody(codeGenerationContext *CodeGenerationCon
 					// Not using astbuilder.AppendList here as we want to tack on a "..." at the end
 					astbuilder.SimpleAssignment(
 						dst.NewIdent(validationsIdent),
-						token.ASSIGN,
 						appendFuncCall),
 				},
 			},

diff --git a/hack/generator/pkg/astmodel/kubernetes_resource_interface.go b/hack/generator/pkg/astmodel/kubernetes_resource_interface.go
@@ -247,7 +247,6 @@ func setEnumAzureNameFunction(enumType TypeName) objectFunctionHandler {
 			Body: []dst.Stmt{
 				astbuilder.SimpleAssignment(
 					azureNameProp,
-					token.ASSIGN,
 					&dst.CallExpr{
 						// cast from the string value to the enum
 						Fun:  enumTypeAST,

diff --git a/hack/generator/pkg/codegen/pipeline/resource_registration_file.go b/hack/generator/pkg/codegen/pipeline/resource_registration_file.go
@@ -262,29 +262,11 @@ func (r *ResourceRegistrationFile) createCreateSchemeFunc(codeGenerationContext
 	ignore := "_"
 	addToScheme := "AddToScheme"
 
-	initSchemeVar := astbuilder.SimpleAssignment(
-		dst.NewIdent(scheme),
-		token.DEFINE,
-		&dst.CallExpr{
-			Fun: &dst.SelectorExpr{
-				X:   dst.NewIdent(runtime),
-				Sel: dst.NewIdent("NewScheme"),
-			},
-			Args: []dst.Expr{},
-		})
+	initSchemeVar := astbuilder.ShortDeclaration(scheme, astbuilder.CallQualifiedFunc(runtime, "NewScheme"))
 
 	clientGoSchemeAssign := astbuilder.SimpleAssignment(
 		dst.NewIdent(ignore),
-		token.ASSIGN,
-		&dst.CallExpr{
-			Fun: &dst.SelectorExpr{
-				X:   dst.NewIdent(clientGoScheme),
-				Sel: dst.NewIdent(addToScheme),
-			},
-			Args: []dst.Expr{
-				dst.NewIdent(scheme),
-			},
-		})
+		astbuilder.CallQualifiedFunc(clientGoScheme, addToScheme, dst.NewIdent(scheme)))
 
 	var importedPackageNames []string
 	for pkg := range r.getImportedPackages() {
@@ -304,16 +286,8 @@ func (r *ResourceRegistrationFile) createCreateSchemeFunc(codeGenerationContext
 	for _, group := range importedPackageNames {
 		groupSchemeAssign := astbuilder.SimpleAssignment(
 			dst.NewIdent(ignore),
-			token.ASSIGN,
-			&dst.CallExpr{
-				Fun: &dst.SelectorExpr{
-					X:   dst.NewIdent(group),
-					Sel: dst.NewIdent(addToScheme),
-				},
-				Args: []dst.Expr{
-					dst.NewIdent(scheme),
-				},
-			})
+			astbuilder.CallQualifiedFunc(group, addToScheme, dst.NewIdent(scheme)))
+
 		groupVersionAssignments = append(groupVersionAssignments, groupSchemeAssign)
 	}