pass: remove redundant jumps and dangling labels (#81)

In jump-table-like constructs, the natural way of writing the code can sometimes produce redundant jumps or labels. Therefore some basic cleanup steps have been proposed. This diff adds two transforms:

1. Remove unconditional jumps to a label immediately following.
2. Remove labels with no references at all.

Fixes #75

ir/ir.go

@@ -212,6 +212,18 @@ func (f *Function) Instructions() []*Instruction {
 	return is
 }
 
+// Labels returns just the list of label nodes.
+func (f *Function) Labels() []Label {
+	var lbls []Label
+	for _, n := range f.Nodes {
+		lbl, ok := n.(Label)
+		if ok {
+			lbls = append(lbls, lbl)
+		}
+	}
+	return lbls
+}
+
 // Stub returns the Go function declaration.
 func (f *Function) Stub() string {
 	return "func " + f.Name + f.Signature.String()

ir/ir_test.go

@@ -0,0 +1,22 @@
+package ir
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestFunctionLabels(t *testing.T) {
+	f := NewFunction("labels")
+	f.AddInstruction(&Instruction{})
+	f.AddLabel("a")
+	f.AddInstruction(&Instruction{})
+	f.AddLabel("b")
+	f.AddInstruction(&Instruction{})
+	f.AddLabel("c")
+	f.AddInstruction(&Instruction{})
+
+	expect := []Label{"a", "b", "c"}
+	if got := f.Labels(); !reflect.DeepEqual(expect, got) {
+		t.Fatalf("f.Labels() = %v; expect %v", got, expect)
+	}
+}

pass/cleanup.go

@@ -5,6 +5,73 @@ import (
 	"github.com/mmcloughlin/avo/operand"
 )
 
+// PruneJumpToFollowingLabel removes jump instructions that target an
+// immediately following label.
+func PruneJumpToFollowingLabel(fn *ir.Function) error {
+	for i := 0; i+1 < len(fn.Nodes); i++ {
+		node := fn.Nodes[i]
+		next := fn.Nodes[i+1]
+
+		// This node is an unconditional jump.
+		inst, ok := node.(*ir.Instruction)
+		if !ok || !inst.IsBranch || inst.IsConditional {
+			continue
+		}
+
+		target := inst.TargetLabel()
+		if target == nil {
+			continue
+		}
+
+		// And the jump target is the immediately following node.
+		lbl, ok := next.(ir.Label)
+		if !ok || lbl != *target {
+			continue
+		}
+
+		// Then the jump is unnecessary and can be removed.
+		fn.Nodes = deletenode(fn.Nodes, i)
+		i--
+	}
+
+	return nil
+}
+
+// PruneDanglingLabels removes labels that are not referenced by any branches.
+func PruneDanglingLabels(fn *ir.Function) error {
+	// Count label references.
+	count := map[ir.Label]int{}
+	for _, n := range fn.Nodes {
+		i, ok := n.(*ir.Instruction)
+		if !ok || !i.IsBranch {
+			continue
+		}
+
+		target := i.TargetLabel()
+		if target == nil {
+			continue
+		}
+
+		count[*target]++
+	}
+
+	// Look for labels with no references.
+	for i := 0; i < len(fn.Nodes); i++ {
+		node := fn.Nodes[i]
+		lbl, ok := node.(ir.Label)
+		if !ok {
+			continue
+		}
+
+		if count[lbl] == 0 {
+			fn.Nodes = deletenode(fn.Nodes, i)
+			i--
+		}
+	}
+
+	return nil
+}
+
 // PruneSelfMoves removes move instructions from one register to itself.
 func PruneSelfMoves(fn *ir.Function) error {
 	return removeinstructions(fn, func(i *ir.Instruction) bool {
@@ -32,14 +99,20 @@ func removeinstructions(fn *ir.Function, predicate func(*ir.Instruction) bool) e
 			continue
 		}
 
-		copy(fn.Nodes[i:], fn.Nodes[i+1:])
+		fn.Nodes = deletenode(fn.Nodes, i)
-		fn.Nodes[len(fn.Nodes)-1] = nil
-		fn.Nodes = fn.Nodes[:len(fn.Nodes)-1]
 	}
 
 	return nil
 }
 
+// deletenode deletes node i from nodes and returns the resulting slice.
+func deletenode(nodes []ir.Node, i int) []ir.Node {
+	n := len(nodes)
+	copy(nodes[i:], nodes[i+1:])
+	nodes[n-1] = nil
+	return nodes[:n-1]
+}
+
 // invalidatecfg clears CFG structures.
 func invalidatecfg(fn *ir.Function) {
 	fn.LabelTarget = nil

pass/cleanup_test.go

@@ -42,3 +42,45 @@ func TestPruneSelfMoves(t *testing.T) {
 		t.Fatal("unexpected result from self-move pruning")
 	}
 }
+
+func TestPruneJumpToFollowingLabel(t *testing.T) {
+	// Construct a function containing a jump to following.
+	ctx := build.NewContext()
+	ctx.Function("add")
+	ctx.XORQ(reg.RAX, reg.RAX)
+	ctx.JMP(operand.LabelRef("next"))
+	ctx.Label("next")
+	ctx.XORQ(reg.RAX, reg.RAX)
+
+	// Build the function with the PruneJumpToFollowingLabel pass.
+	fn := BuildFunction(t, ctx, pass.PruneJumpToFollowingLabel)
+
+	// Confirm no JMP instruction remains.
+	for _, i := range fn.Instructions() {
+		if i.Opcode == "JMP" {
+			t.Fatal("JMP instruction not removed")
+		}
+	}
+}
+
+func TestPruneDanglingLabels(t *testing.T) {
+	// Construct a function containing an unreferenced label.
+	ctx := build.NewContext()
+	ctx.Function("add")
+	ctx.XORQ(reg.RAX, reg.RAX)
+	ctx.JMP(operand.LabelRef("referenced"))
+	ctx.XORQ(reg.RAX, reg.RAX)
+	ctx.Label("dangling")
+	ctx.XORQ(reg.RAX, reg.RAX)
+	ctx.Label("referenced")
+	ctx.XORQ(reg.RAX, reg.RAX)
+
+	// Build the function with the PruneDanglingLabels pass.
+	fn := BuildFunction(t, ctx, pass.PruneDanglingLabels)
+
+	// Confirm the only label remaining is "referenced".
+	expect := []ir.Label{"referenced"}
+	if !reflect.DeepEqual(expect, fn.Labels()) {
+		t.Fatal("expected dangling label to be removed")
+	}
+}

pass/pass.go

@@ -11,6 +11,8 @@ import (
 // Compile pass compiles an avo file. Upon successful completion the avo file
 // may be printed to Go assembly.
 var Compile = Concat(
+	FunctionPass(PruneJumpToFollowingLabel),
+	FunctionPass(PruneDanglingLabels),
 	FunctionPass(LabelTarget),
 	FunctionPass(CFG),
 	FunctionPass(Liveness),

tests/fmt/asm.go

@@ -4,6 +4,7 @@ package main
 
 import (
 	. "github.com/mmcloughlin/avo/build"
+	. "github.com/mmcloughlin/avo/operand"
 	. "github.com/mmcloughlin/avo/reg"
 )
 
@@ -19,6 +20,7 @@ func main() {
 	Label("label")
 	Comment("Comment after label.")
 	ADDQ(R8, R8)
+	JMP(LabelRef("label"))
 
 	RET()
 

tests/fmt/fmt.s

@@ -13,4 +13,5 @@ TEXT ·Formatting(SB), NOSPLIT, $0
 label:
 	// Comment after label.
 	ADDQ R8, R8
+	JMP  label
 	RET

tests/labels/asm.go

@@ -0,0 +1,28 @@
+// +build ignore
+
+package main
+
+import (
+	. "github.com/mmcloughlin/avo/build"
+	. "github.com/mmcloughlin/avo/operand"
+	. "github.com/mmcloughlin/avo/reg"
+)
+
+func main() {
+	TEXT("Labels", NOSPLIT, "func() uint64")
+	XORQ(RAX, RAX)
+	INCQ(RAX)
+	Label("neverused")
+	INCQ(RAX)
+	INCQ(RAX)
+	INCQ(RAX)
+	INCQ(RAX)
+	JMP(LabelRef("next"))
+	Label("next")
+	INCQ(RAX)
+	INCQ(RAX)
+	Store(RAX, ReturnIndex(0))
+	RET()
+
+	Generate()
+}

tests/labels/doc.go

@@ -0,0 +1,2 @@
+// Package labels tests for cleanup of redundant labels.
+package labels

tests/labels/labels.s

@@ -0,0 +1,16 @@
+// Code generated by command: go run asm.go -out labels.s -stubs stub.go. DO NOT EDIT.
+
+#include "textflag.h"
+
+// func Labels() uint64
+TEXT ·Labels(SB), NOSPLIT, $0-8
+	XORQ AX, AX
+	INCQ AX
+	INCQ AX
+	INCQ AX
+	INCQ AX
+	INCQ AX
+	INCQ AX
+	INCQ AX
+	MOVQ AX, ret+0(FP)
+	RET

tests/labels/labels_test.go

@@ -0,0 +1,15 @@
+package labels
+
+import (
+	"testing"
+	"testing/quick"
+)
+
+//go:generate go run asm.go -out labels.s -stubs stub.go
+
+func TestLabels(t *testing.T) {
+	expect := func() uint64 { return 7 }
+	if err := quick.CheckEqual(Labels, expect, nil); err != nil {
+		t.Fatal(err)
+	}
+}

tests/labels/stub.go

@@ -0,0 +1,5 @@
+// Code generated by command: go run asm.go -out labels.s -stubs stub.go. DO NOT EDIT.
+
+package labels
+
+func Labels() uint64