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>

8 files changed, 114 insertions, 25 deletions

diff --git a/src/cmd/internal/goobj/read.go b/src/cmd/internal/goobj/read.go
index e61e95dcc8..48537d2b1c 100644
--- a/src/cmd/internal/goobj/read.go
+++ b/src/cmd/internal/goobj/read.go
@@ -95,6 +95,7 @@ type Var struct {
 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 @@ func (r *objReader) parseObject(prefix []byte) error {
 			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 e8b092a2a8..8e5b598084 100644
--- a/src/cmd/internal/obj/arm64/asm7.go
+++ b/src/cmd/internal/obj/arm64/asm7.go
@@ -886,25 +886,10 @@ const OP_NOOP = 0xd503201f
 
 // 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 @@ func span7(ctxt *obj.Link, cursym *obj.LSym, newprog obj.ProgAlloc) {
 		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 @@ func span7(ctxt *obj.Link, cursym *obj.LSym, newprog obj.ProgAlloc) {
 			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 1691828739..9efdb0217f 100644
--- a/src/cmd/internal/obj/arm64/asm_test.go
+++ b/src/cmd/internal/obj/arm64/asm_test.go
@@ -18,7 +18,9 @@ import (
 
 // 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 @@ func TestLarge(t *testing.T) {
 		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 @@ func gen(buf *bytes.Buffer) {
 	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 d1cc536a8c..0879c611ba 100644
--- a/src/cmd/internal/obj/link.go
+++ b/src/cmd/internal/obj/link.go
@@ -398,6 +398,7 @@ type LSym struct {
 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 7fd97f7363..46e8a551ad 100644
--- a/src/cmd/internal/obj/objfile.go
+++ b/src/cmd/internal/obj/objfile.go
@@ -346,6 +346,7 @@ func (w *objWriter) writeSym(s *LSym) {
 
 	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 7ca01c8c25..31613e5cef 100644
--- a/src/cmd/link/internal/ld/data.go
+++ b/src/cmd/link/internal/ld/data.go
@@ -2119,6 +2119,10 @@ func assignAddress(ctxt *Link, sect *sym.Section, n int, s *sym.Symbol, va uint6
 		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 a15d3c3e07..295acb2d29 100644
--- a/src/cmd/link/internal/objfile/objfile.go
+++ b/src/cmd/link/internal/objfile/objfile.go
@@ -312,6 +312,7 @@ overwrite:
 
 		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 4f792bd1f1..7d87093813 100644
--- a/src/cmd/link/link_test.go
+++ b/src/cmd/link/link_test.go
@@ -447,3 +447,73 @@ func TestStrictDup(t *testing.T) {
 		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)
+	}
+}