runtime: update mkmalloc to make generated code look nicer

This cl adds a new operation that can remove an if statement or replace
it with its body if its condition is var or !var for some variable var
that's being replaced with a constant.

Change-Id: I864abf1f023b2a66b2299ca65d4f837d6a6a6964
Reviewed-on: https://go-review.googlesource.com/c/go/+/724940
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Reviewed-by: Michael Matloob <matloob@google.com>
Auto-Submit: Michael Matloob <matloob@google.com>

1 files changed, 51 insertions, 6 deletions

diff --git a/src/runtime/_mkmalloc/mkmalloc.go b/src/runtime/_mkmalloc/mkmalloc.go
index 46c50d6661..8032983da8 100644
--- a/src/runtime/_mkmalloc/mkmalloc.go
+++ b/src/runtime/_mkmalloc/mkmalloc.go
@@ -107,6 +107,7 @@ type replacementKind int
 const (
 	inlineFunc = replacementKind(iota)
 	subBasicLit
+	foldCondition
 )
 
 // op is a single inlining operation for the inliner. Any calls to the function
@@ -171,7 +172,7 @@ func specializedMallocConfig(classes []class, sizeToSizeClass []uint8) generator
 					{subBasicLit, "elemsize_", str(elemsize)},
 					{subBasicLit, "sizeclass_", str(sc)},
 					{subBasicLit, "noscanint_", str(noscan)},
-					{subBasicLit, "isTiny_", str(0)},
+					{foldCondition, "isTiny_", str(false)},
 				},
 			})
 		}
@@ -199,7 +200,7 @@ func specializedMallocConfig(classes []class, sizeToSizeClass []uint8) generator
 					{subBasicLit, "sizeclass_", str(tinySizeClass)},
 					{subBasicLit, "size_", str(s)},
 					{subBasicLit, "noscanint_", str(noscan)},
-					{subBasicLit, "isTiny_", str(1)},
+					{foldCondition, "isTiny_", str(true)},
 				},
 			})
 		}
@@ -217,7 +218,7 @@ func specializedMallocConfig(classes []class, sizeToSizeClass []uint8) generator
 					{subBasicLit, "elemsize_", str(elemsize)},
 					{subBasicLit, "sizeclass_", str(sc)},
 					{subBasicLit, "noscanint_", str(noscan)},
-					{subBasicLit, "isTiny_", str(0)},
+					{foldCondition, "isTiny_", str(false)},
 				},
 			})
 		}
@@ -277,10 +278,17 @@ func inline(config generatorConfig) []byte {
 		// Apply each of the ops given by the specs
 		stamped := ast.Node(containingFuncCopy)
 		for _, repl := range spec.ops {
-			if toDecl, ok := funcDecls[repl.to]; ok {
-				stamped = inlineFunction(stamped, repl.from, toDecl)
-			} else {
+			switch repl.kind {
+			case inlineFunc:
+				if toDecl, ok := funcDecls[repl.to]; ok {
+					stamped = inlineFunction(stamped, repl.from, toDecl)
+				}
+			case subBasicLit:
 				stamped = substituteWithBasicLit(stamped, repl.from, repl.to)
+			case foldCondition:
+				stamped = foldIfCondition(stamped, repl.from, repl.to)
+			default:
+				log.Fatal("unknown op kind %v", repl.kind)
 			}
 		}
 
@@ -310,6 +318,43 @@ func substituteWithBasicLit(node ast.Node, from, to string) ast.Node {
 	}, nil)
 }
 
+// foldIfCondition looks for if statements with a single boolean variable from, or
+// the negation of from and either replaces it with its body or nothing,
+// depending on whether the to value is true or false.
+func foldIfCondition(node ast.Node, from, to string) ast.Node {
+	var isTrue bool
+	switch to {
+	case "true":
+		isTrue = true
+	case "false":
+		isTrue = false
+	default:
+		log.Fatalf("op 'to' expr %q is not true or false", to)
+	}
+	return astutil.Apply(node, func(cursor *astutil.Cursor) bool {
+		var foldIfTrue bool
+		ifexpr, ok := cursor.Node().(*ast.IfStmt)
+		if !ok {
+			return true
+		}
+		if isIdentWithName(ifexpr.Cond, from) {
+			foldIfTrue = true
+		} else if unaryexpr, ok := ifexpr.Cond.(*ast.UnaryExpr); ok && unaryexpr.Op == token.NOT && isIdentWithName(unaryexpr.X, from) {
+			foldIfTrue = false
+		} else {
+			// not an if with from or !from.
+			return true
+		}
+		if foldIfTrue == isTrue {
+			for _, stmt := range ifexpr.Body.List {
+				cursor.InsertBefore(stmt)
+			}
+		}
+		cursor.Delete()
+		return true
+	}, nil)
+}
+
 // inlineFunction recursively replaces calls to the function 'from' with the body of the function
 // 'toDecl'. All calls to 'from' must appear in assignment statements.
 // The replacement is very simple: it doesn't substitute the arguments for the parameters, so the