4 files changed, 37 insertions, 25 deletions
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() {