cmd/asm: align an instruction or a function's address

Recently, the gVisor project needs an instruction's address
with 128 bytes alignment to fit the architecture requirement
for interrupt table.

This patch allows aligning an instruction's address to be
aligned to a specific value (2^n and in the range [8, 2048])

The main changes include:

1. Adds a new element in the FuncInfo structure defined in
cmd/internal/obj/link.go file to record the alignment
information.

2. Adds a new element in the Func structure defined in
cmd/internal/goobj/read.go file to read the alignment
information.

3. Adds the assembler support to align an intruction's offset
with a specific value (2^n and in the range [8, 2048]).
e.g. "PCALIGN $256" indicates that the next instruction should
be aligned to 256 bytes.

4. An instruction's alignment is relative to the start of the
function where this instruction is located, so the function's
address must be aligned to the same or coarser boundary.

This CL also adds a test.

Change-Id: I9b365c111b3a12f767728f1b45aa0c00f073c37d
Reviewed-on: https://go-review.googlesource.com/c/go/+/226997
Reviewed-by: Bryan C. Mills <bcmills@google.com>
Reviewed-by: Cherry Zhang <cherryyz@google.com>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
diff --git a/src/cmd/internal/goobj/read.go b/src/cmd/internal/goobj/read.go
index e61e95d..48537d2 100644
--- a/src/cmd/internal/goobj/read.go
+++ b/src/cmd/internal/goobj/read.go
@@ -95,6 +95,7 @@
 type Func struct {
 	Args     int64      // size in bytes of argument frame: inputs and outputs
 	Frame    int64      // size in bytes of local variable frame
+	Align    uint32     // alignment requirement in bytes for the address of the function
 	Leaf     bool       // function omits save of link register (ARM)
 	NoSplit  bool       // function omits stack split prologue
 	TopFrame bool       // function is the top of the call stack
@@ -590,6 +591,7 @@
 			s.Func = f
 			f.Args = r.readInt()
 			f.Frame = r.readInt()
+			f.Align = uint32(r.readInt())
 			flags := r.readInt()
 			f.Leaf = flags&(1<<0) != 0
 			f.TopFrame = flags&(1<<4) != 0
diff --git a/src/cmd/internal/obj/arm64/asm7.go b/src/cmd/internal/obj/arm64/asm7.go
index e8b092a..8e5b598 100644
--- a/src/cmd/internal/obj/arm64/asm7.go
+++ b/src/cmd/internal/obj/arm64/asm7.go
@@ -886,25 +886,10 @@
 
 // align code to a certain length by padding bytes.
 func pcAlignPadLength(pc int64, alignedValue int64, ctxt *obj.Link) int {
-	switch alignedValue {
-	case 8:
-		if pc%8 == 4 {
-			return 4
-		}
-	case 16:
-		switch pc % 16 {
-		case 4:
-			return 12
-		case 8:
-			return 8
-		case 12:
-			return 4
-		}
-	default:
-		ctxt.Diag("Unexpected alignment: %d for PCALIGN directive\n", alignedValue)
+	if !((alignedValue&(alignedValue-1) == 0) && 8 <= alignedValue && alignedValue <= 2048) {
+		ctxt.Diag("alignment value of an instruction must be a power of two and in the range [8, 2048], got %d\n", alignedValue)
 	}
-
-	return 0
+	return int(-pc & (alignedValue - 1))
 }
 
 func span7(ctxt *obj.Link, cursym *obj.LSym, newprog obj.ProgAlloc) {
@@ -940,8 +925,12 @@
 		if m == 0 {
 			switch p.As {
 			case obj.APCALIGN:
-				a := p.From.Offset
-				m = pcAlignPadLength(pc, a, ctxt)
+				alignedValue := p.From.Offset
+				m = pcAlignPadLength(pc, alignedValue, ctxt)
+				// Update the current text symbol alignment value.
+				if int32(alignedValue) > cursym.Func.Align {
+					cursym.Func.Align = int32(alignedValue)
+				}
 				break
 			case obj.ANOP, obj.AFUNCDATA, obj.APCDATA:
 				continue
@@ -1017,8 +1006,8 @@
 			if m == 0 {
 				switch p.As {
 				case obj.APCALIGN:
-					a := p.From.Offset
-					m = pcAlignPadLength(pc, a, ctxt)
+					alignedValue := p.From.Offset
+					m = pcAlignPadLength(pc, alignedValue, ctxt)
 					break
 				case obj.ANOP, obj.AFUNCDATA, obj.APCDATA:
 					continue
diff --git a/src/cmd/internal/obj/arm64/asm_test.go b/src/cmd/internal/obj/arm64/asm_test.go
index 1691828..9efdb02 100644
--- a/src/cmd/internal/obj/arm64/asm_test.go
+++ b/src/cmd/internal/obj/arm64/asm_test.go
@@ -18,7 +18,9 @@
 
 // TestLarge generates a very large file to verify that large
 // program builds successfully, in particular, too-far
-// conditional branches are fixed.
+// conditional branches are fixed, and also verify that the
+// instruction's pc can be correctly aligned even when branches
+// need to be fixed.
 func TestLarge(t *testing.T) {
 	if testing.Short() {
 		t.Skip("Skip in short mode")
@@ -41,11 +43,28 @@
 		t.Fatalf("can't write output: %v\n", err)
 	}
 
-	// build generated file
-	cmd := exec.Command(testenv.GoToolPath(t), "tool", "asm", "-o", filepath.Join(dir, "x.o"), tmpfile)
+	pattern := `0x0080\s00128\s\(.*\)\tMOVD\t\$3,\sR3`
+
+	// assemble generated file
+	cmd := exec.Command(testenv.GoToolPath(t), "tool", "asm", "-S", "-o", filepath.Join(dir, "test.o"), tmpfile)
 	cmd.Env = append(os.Environ(), "GOARCH=arm64", "GOOS=linux")
 	out, err := cmd.CombinedOutput()
 	if err != nil {
+		t.Errorf("Assemble failed: %v, output: %s", err, out)
+	}
+	matched, err := regexp.MatchString(pattern, string(out))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !matched {
+		t.Errorf("The alignment is not correct: %t, output:%s\n", matched, out)
+	}
+
+	// build generated file
+	cmd = exec.Command(testenv.GoToolPath(t), "tool", "asm", "-o", filepath.Join(dir, "x.o"), tmpfile)
+	cmd.Env = append(os.Environ(), "GOARCH=arm64", "GOOS=linux")
+	out, err = cmd.CombinedOutput()
+	if err != nil {
 		t.Errorf("Build failed: %v, output: %s", err, out)
 	}
 }
@@ -56,6 +75,8 @@
 	fmt.Fprintln(buf, "TBZ $5, R0, label")
 	fmt.Fprintln(buf, "CBZ R0, label")
 	fmt.Fprintln(buf, "BEQ label")
+	fmt.Fprintln(buf, "PCALIGN $128")
+	fmt.Fprintln(buf, "MOVD $3, R3")
 	for i := 0; i < 1<<19; i++ {
 		fmt.Fprintln(buf, "MOVD R0, R1")
 	}
diff --git a/src/cmd/internal/obj/link.go b/src/cmd/internal/obj/link.go
index d1cc536..0879c61 100644
--- a/src/cmd/internal/obj/link.go
+++ b/src/cmd/internal/obj/link.go
@@ -398,6 +398,7 @@
 type FuncInfo struct {
 	Args     int32
 	Locals   int32
+	Align    int32
 	Text     *Prog
 	Autot    map[*LSym]struct{}
 	Pcln     Pcln
diff --git a/src/cmd/internal/obj/objfile.go b/src/cmd/internal/obj/objfile.go
index 7fd97f7..46e8a55 100644
--- a/src/cmd/internal/obj/objfile.go
+++ b/src/cmd/internal/obj/objfile.go
@@ -346,6 +346,7 @@
 
 	w.writeInt(int64(s.Func.Args))
 	w.writeInt(int64(s.Func.Locals))
+	w.writeInt(int64(s.Func.Align))
 	w.writeBool(s.NoSplit())
 	flags = int64(0)
 	if s.Leaf() {
diff --git a/src/cmd/link/internal/ld/data.go b/src/cmd/link/internal/ld/data.go
index 7ca01c8..31613e5 100644
--- a/src/cmd/link/internal/ld/data.go
+++ b/src/cmd/link/internal/ld/data.go
@@ -2119,6 +2119,10 @@
 		funcsize = uint64(s.Size)
 	}
 
+	if sect.Align < s.Align {
+		sect.Align = s.Align
+	}
+
 	// On ppc64x a text section should not be larger than 2^26 bytes due to the size of
 	// call target offset field in the bl instruction.  Splitting into smaller text
 	// sections smaller than this limit allows the GNU linker to modify the long calls
diff --git a/src/cmd/link/internal/objfile/objfile.go b/src/cmd/link/internal/objfile/objfile.go
index a15d3c3..295acb2 100644
--- a/src/cmd/link/internal/objfile/objfile.go
+++ b/src/cmd/link/internal/objfile/objfile.go
@@ -312,6 +312,7 @@
 
 		pc.Args = r.readInt32()
 		pc.Locals = r.readInt32()
+		s.Align = r.readInt32()
 		if r.readUint8() != 0 {
 			s.Attr |= sym.AttrNoSplit
 		}
diff --git a/src/cmd/link/link_test.go b/src/cmd/link/link_test.go
index 4f792bd..7d87093 100644
--- a/src/cmd/link/link_test.go
+++ b/src/cmd/link/link_test.go
@@ -447,3 +447,73 @@
 		t.Errorf("unexpected output:\n%s", out)
 	}
 }
+
+const testFuncAlignSrc = `
+package main
+import (
+	"fmt"
+	"reflect"
+)
+func alignPc()
+
+func main() {
+	addr := reflect.ValueOf(alignPc).Pointer()
+	if (addr % 512) != 0 {
+		fmt.Printf("expected 512 bytes alignment, got %v\n", addr)
+	} else {
+		fmt.Printf("PASS")
+	}
+}
+`
+
+const testFuncAlignAsmSrc = `
+#include "textflag.h"
+
+TEXT	·alignPc(SB),NOSPLIT, $0-0
+	MOVD	$2, R0
+	PCALIGN	$512
+	MOVD	$3, R1
+	RET
+`
+
+// TestFuncAlign verifies that the address of a function can be aligned
+// with a specfic value on arm64.
+func TestFuncAlign(t *testing.T) {
+	if runtime.GOARCH != "arm64" || runtime.GOOS != "linux" {
+		t.Skip("skipping on non-linux/arm64 platform")
+	}
+	testenv.MustHaveGoBuild(t)
+
+	tmpdir, err := ioutil.TempDir("", "TestFuncAlign")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpdir)
+
+	src := filepath.Join(tmpdir, "falign.go")
+	err = ioutil.WriteFile(src, []byte(testFuncAlignSrc), 0666)
+	if err != nil {
+		t.Fatal(err)
+	}
+	src = filepath.Join(tmpdir, "falign.s")
+	err = ioutil.WriteFile(src, []byte(testFuncAlignAsmSrc), 0666)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Build and run with old object file format.
+	cmd := exec.Command(testenv.GoToolPath(t), "build", "-o", "falign")
+	cmd.Dir = tmpdir
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		t.Errorf("build failed: %v", err)
+	}
+	cmd = exec.Command(tmpdir + "/falign")
+	out, err = cmd.CombinedOutput()
+	if err != nil {
+		t.Errorf("failed to run with err %v, output: %s", err, out)
+	}
+	if string(out) != "PASS" {
+		t.Errorf("unexpected output: %s\n", out)
+	}
+}