src/cmd/internal/obj/wasm: optimize blocks in wasm binary

This change optimizes the blocks in the wasm binary by generating the
entryPointLoop only if it is used and adding an unwindExit block to
be able to use the short BrIf instruction for unwinding the stack.
These changes were suggested by the wasm-opt tool and reduce the
wasm binary size of "hello world" by 1.5%.

Change-Id: Ie52db2fa2d9b8482f9a78b7c189231750811fe97
Reviewed-on: https://go-review.googlesource.com/c/go/+/167937
Run-TryBot: Richard Musiol <neelance@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

src/cmd/internal/obj/wasm/wasmobj.go

@@ -341,47 +341,20 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 		p = appendp(p, AEnd)
 	}
 
-	// Add Block instructions for resume points and BrTable to jump to selected resume point.
-	if numResumePoints > 0 {
-		p := s.Func.Text
-		p = appendp(p, ALoop) // entryPointLoop, used to jump between basic blocks
-
-		for i := 0; i < numResumePoints+1; i++ {
-			p = appendp(p, ABlock)
-		}
-		p = appendp(p, AGet, regAddr(REG_PC_B)) // read next basic block from PC_B
-		p = appendp(p, ABrTable, obj.Addr{Val: tableIdxs})
-		p = appendp(p, AEnd) // end of Block
-
-		for p.Link != nil {
-			p = p.Link
-		}
-
-		p = appendp(p, AEnd) // end of entryPointLoop
-		p = appendp(p, obj.AUNDEF)
-	}
-
-	p := s.Func.Text
+	// record the branches targeting the entry loop and the unwind exit,
+	// their targets with be filled in later
+	var entryPointLoopBranches []*obj.Prog
+	var unwindExitBranches []*obj.Prog
 	currentDepth := 0
-	blockDepths := make(map[*obj.Prog]int)
-	for p != nil {
+	for p := s.Func.Text; p != nil; p = p.Link {
 		switch p.As {
 		case ABlock, ALoop, AIf:
 			currentDepth++
-			blockDepths[p] = currentDepth
 		case AEnd:
 			currentDepth--
 		}
 
 		switch p.As {
-		case ABr, ABrIf:
-			if p.To.Type == obj.TYPE_BRANCH {
-				blockDepth, ok := blockDepths[p.To.Val.(*obj.Prog)]
-				if !ok {
-					panic("label not at block")
-				}
-				p.To = constAddr(int64(currentDepth - blockDepth))
-			}
 		case obj.AJMP:
 			jmp := *p
 			p.As = obj.ANOP
@@ -390,7 +363,8 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 				// jump to basic block
 				p = appendp(p, AI32Const, constAddr(jmp.To.Val.(*obj.Prog).Pc))
 				p = appendp(p, ASet, regAddr(REG_PC_B)) // write next basic block to PC_B
-				p = appendp(p, ABr, constAddr(int64(currentDepth-1))) // jump to beginning of entryPointLoop
+				p = appendp(p, ABr)                     // jump to beginning of entryPointLoop
+				entryPointLoopBranches = append(entryPointLoopBranches, p)
 				break
 			}
 
@@ -478,16 +452,16 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 			}
 
 			// return value of call is on the top of the stack, indicating whether to unwind the WebAssembly stack
-			p = appendp(p, AIf)
 			if call.As == ACALLNORESUME && call.To.Sym != sigpanic { // sigpanic unwinds the stack, but it never resumes
 				// trying to unwind WebAssembly stack but call has no resume point, terminate with error
+				p = appendp(p, AIf)
 				p = appendp(p, obj.AUNDEF)
+				p = appendp(p, AEnd)
 			} else {
 				// unwinding WebAssembly stack to switch goroutine, return 1
-				p = appendp(p, AI32Const, constAddr(1))
-				p = appendp(p, AReturn)
+				p = appendp(p, ABrIf)
+				unwindExitBranches = append(unwindExitBranches, p)
 			}
-			p = appendp(p, AEnd)
 
 			// jump to before the call if jmpdefer has reset the return address to the call's PC
 			if call.To.Sym == deferreturn {
@@ -550,12 +524,9 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 			p = appendp(p, AI32Const, constAddr(0))
 			p = appendp(p, AReturn)
 		}
-
-		p = p.Link
 	}
 
-	p = s.Func.Text
-	for p != nil {
+	for p := s.Func.Text; p != nil; p = p.Link {
 		switch p.From.Name {
 		case obj.NAME_AUTO:
 			p.From.Offset += int64(framesize)
@@ -702,8 +673,65 @@ func preprocess(ctxt *obj.Link, s *obj.LSym, newprog obj.ProgAlloc) {
 			p = appendp(p, ACall, obj.Addr{Type: obj.TYPE_MEM, Name: obj.NAME_EXTERN, Sym: s})
 			p.Mark = WasmImport
 		}
+	}
 
-		p = p.Link
+	{
+		p := s.Func.Text
+		if len(unwindExitBranches) > 0 {
+			p = appendp(p, ABlock) // unwindExit, used to return 1 when unwinding the stack
+			for _, b := range unwindExitBranches {
+				b.To = obj.Addr{Type: obj.TYPE_BRANCH, Val: p}
+			}
+		}
+		if len(entryPointLoopBranches) > 0 {
+			p = appendp(p, ALoop) // entryPointLoop, used to jump between basic blocks
+			for _, b := range entryPointLoopBranches {
+				b.To = obj.Addr{Type: obj.TYPE_BRANCH, Val: p}
+			}
+		}
+		if numResumePoints > 0 {
+			// Add Block instructions for resume points and BrTable to jump to selected resume point.
+			for i := 0; i < numResumePoints+1; i++ {
+				p = appendp(p, ABlock)
+			}
+			p = appendp(p, AGet, regAddr(REG_PC_B)) // read next basic block from PC_B
+			p = appendp(p, ABrTable, obj.Addr{Val: tableIdxs})
+			p = appendp(p, AEnd) // end of Block
+		}
+		for p.Link != nil {
+			p = p.Link // function instructions
+		}
+		if len(entryPointLoopBranches) > 0 {
+			p = appendp(p, AEnd) // end of entryPointLoop
+		}
+		p = appendp(p, obj.AUNDEF)
+		if len(unwindExitBranches) > 0 {
+			p = appendp(p, AEnd) // end of unwindExit
+			p = appendp(p, AI32Const, constAddr(1))
+		}
+	}
+
+	currentDepth = 0
+	blockDepths := make(map[*obj.Prog]int)
+	for p := s.Func.Text; p != nil; p = p.Link {
+		switch p.As {
+		case ABlock, ALoop, AIf:
+			currentDepth++
+			blockDepths[p] = currentDepth
+		case AEnd:
+			currentDepth--
+		}
+
+		switch p.As {
+		case ABr, ABrIf:
+			if p.To.Type == obj.TYPE_BRANCH {
+				blockDepth, ok := blockDepths[p.To.Val.(*obj.Prog)]
+				if !ok {
+					panic("label not at block")
+				}
+				p.To = constAddr(int64(currentDepth - blockDepth))
+			}
+		}
 	}
 }