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>

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 {
+			switch repl.kind {
+			case inlineFunc:
 				if toDecl, ok := funcDecls[repl.to]; ok {
 					stamped = inlineFunction(stamped, repl.from, toDecl)
-			} else {
+				}
+			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

src/runtime/malloc_stubs.go

@@ -37,7 +37,7 @@ const elemsize_ = 8
 const sizeclass_ = 0
 const noscanint_ = 0
 const size_ = 0
-const isTiny_ = 0
+const isTiny_ = false
 
 func malloc0(size uintptr, typ *_type, needzero bool) unsafe.Pointer {
 	if doubleCheckMalloc {
@@ -58,16 +58,17 @@ func mallocPanic(size uintptr, typ *_type, needzero bool) unsafe.Pointer {
 // to steer out of this codepath early if sanitizers are enabled.
 func mallocStub(size uintptr, typ *_type, needzero bool) unsafe.Pointer {
 
+	if isTiny_ {
 		// secret code, need to avoid the tiny allocator since it might keep
 		// co-located values alive longer and prevent timely zero-ing
 		//
 		// Call directly into the NoScan allocator.
 		// See go.dev/issue/76356
-	const isTiny = isTiny_ == 1
 		gp := getg()
-	if goexperiment.RuntimeSecret && isTiny && gp.secret > 0 {
+		if goexperiment.RuntimeSecret && gp.secret > 0 {
 			return mallocgcSmallNoScanSC2(size, typ, needzero)
 		}
+	}
 	if doubleCheckMalloc {
 		if gcphase == _GCmarktermination {
 			throw("mallocgc called with gcphase == _GCmarktermination")
@@ -95,11 +96,14 @@ func mallocStub(size uintptr, typ *_type, needzero bool) unsafe.Pointer {
 	// Actually do the allocation.
 	x, elemsize := inlinedMalloc(size, typ, needzero)
 
+	if !isTiny_ {
+		gp := getg()
 		if goexperiment.RuntimeSecret && gp.secret > 0 {
 			// Mark any object allocated while in secret mode as secret.
 			// This ensures we zero it immediately when freeing it.
 			addSecret(x)
 		}
+	}
 
 	// Notify valgrind, if enabled.
 	// To allow the compiler to not know about valgrind, we do valgrind instrumentation