ir,build: pragma support (#97)

Adds support for arbitrary compiler directives.

Fixes #15

README.md

@@ -174,6 +174,7 @@ For demonstrations of `avo` features:
 * **[complex](examples/complex):** Working with `complex{64,128}` types.
 * **[data](examples/data):** Defining `DATA` sections.
 * **[ext](examples/ext):** Interacting with types from external packages.
+* **[pragma](examples/pragma):** Apply compiler directives to generated functions.
 
 ### Real Examples
 

build/context.go

@@ -92,6 +92,11 @@ func (c *Context) Doc(lines ...string) {
 	c.activefunc().Doc = lines
 }
 
+// Pragma adds a compiler directive to the currently active function.
+func (c *Context) Pragma(directive string, args ...string) {
+	c.activefunc().AddPragma(directive, args...)
+}
+
 // Attributes sets function attributes for the currently active function.
 func (c *Context) Attributes(a attr.Attribute) {
 	c.activefunc().Attributes = a

build/global.go

@@ -118,6 +118,9 @@ func Dereference(ptr gotypes.Component) gotypes.Component { return ctx.Dereferen
 // Doc sets documentation comment lines for the currently active function.
 func Doc(lines ...string) { ctx.Doc(lines...) }
 
+// Pragma adds a compiler directive to the currently active function.
+func Pragma(directive string, args ...string) { ctx.Pragma(directive, args...) }
+
 // Attributes sets function attributes for the currently active function.
 func Attributes(a attr.Attribute) { ctx.Attributes(a) }
 

examples/README.md

@@ -12,7 +12,7 @@ Features:
 * **[complex](complex):** Working with `complex{64,128}` types.
 * **[data](data):** Defining `DATA` sections.
 * **[ext](ext):** Interacting with types from external packages.
-
+* **[pragma](pragma):** Apply compiler directives to generated functions.
 
 "Real" examples:
 

examples/pragma/README.md

@@ -0,0 +1,33 @@
+# pragma
+
+Apply [compiler directives](https://golang.org/pkg/cmd/compile/#hdr-Compiler_Directives) to `avo` functions.
+
+The [code generator](asm.go) uses the `Pragma` function to apply the `//go:noescape` directive to the `Add` function:
+
+[embedmd]:# (asm.go go /func main/ /^}/)
+```go
+func main() {
+	TEXT("Add", NOSPLIT, "func(z, x, y *uint64)")
+	Pragma("noescape")
+	Doc("Add adds the values at x and y and writes the result to z.")
+	zptr := Mem{Base: Load(Param("z"), GP64())}
+	xptr := Mem{Base: Load(Param("x"), GP64())}
+	yptr := Mem{Base: Load(Param("y"), GP64())}
+	x, y := GP64(), GP64()
+	MOVQ(xptr, x)
+	MOVQ(yptr, y)
+	ADDQ(x, y)
+	MOVQ(y, zptr)
+	RET()
+	Generate()
+}
+```
+
+Note the directive is applied in the generated stub file:
+
+[embedmd]:# (stub.go go /\/\/ Add/ /func/)
+```go
+// Add adds the values at x and y and writes the result to z.
+//go:noescape
+func
+```

examples/pragma/asm.go

@@ -0,0 +1,24 @@
+// +build ignore
+
+package main
+
+import (
+	. "github.com/mmcloughlin/avo/build"
+	. "github.com/mmcloughlin/avo/operand"
+)
+
+func main() {
+	TEXT("Add", NOSPLIT, "func(z, x, y *uint64)")
+	Pragma("noescape")
+	Doc("Add adds the values at x and y and writes the result to z.")
+	zptr := Mem{Base: Load(Param("z"), GP64())}
+	xptr := Mem{Base: Load(Param("x"), GP64())}
+	yptr := Mem{Base: Load(Param("y"), GP64())}
+	x, y := GP64(), GP64()
+	MOVQ(xptr, x)
+	MOVQ(yptr, y)
+	ADDQ(x, y)
+	MOVQ(y, zptr)
+	RET()
+	Generate()
+}

examples/pragma/doc.go

@@ -0,0 +1,2 @@
+// Package pragma demonstrates the use of compiler directives.
+package pragma

examples/pragma/pragma.s

@@ -0,0 +1,14 @@
+// Code generated by command: go run asm.go -out pragma.s -stubs stub.go. DO NOT EDIT.
+
+#include "textflag.h"
+
+// func Add(z *uint64, x *uint64, y *uint64)
+TEXT ·Add(SB), NOSPLIT, $0-24
+	MOVQ z+0(FP), AX
+	MOVQ x+8(FP), CX
+	MOVQ y+16(FP), DX
+	MOVQ (CX), CX
+	MOVQ (DX), DX
+	ADDQ CX, DX
+	MOVQ DX, (AX)
+	RET

examples/pragma/pragma_test.go

@@ -0,0 +1,16 @@
+package pragma
+
+import (
+	"testing"
+	"testing/quick"
+)
+
+//go:generate go run asm.go -out pragma.s -stubs stub.go
+
+func TestAdd(t *testing.T) {
+	got := func(x, y uint64) (z uint64) { Add(&z, &x, &y); return }
+	expect := func(x, y uint64) uint64 { return x + y }
+	if err := quick.CheckEqual(got, expect, nil); err != nil {
+		t.Fatal(err)
+	}
+}

examples/pragma/stub.go

@@ -0,0 +1,7 @@
+// Code generated by command: go run asm.go -out pragma.s -stubs stub.go. DO NOT EDIT.
+
+package pragma
+
+// Add adds the values at x and y and writes the result to z.
+//go:noescape
+func Add(z *uint64, x *uint64, y *uint64)

ir/ir.go

@@ -144,10 +144,17 @@ func (f *File) Functions() []*Function {
 	return fns
 }
 
+// Pragma represents a function compiler directive.
+type Pragma struct {
+	Directive string
+	Arguments []string
+}
+
 // Function represents an assembly function.
 type Function struct {
 	Name       string
 	Attributes attr.Attribute
+	Pragmas    []Pragma
 	Doc        []string
 	Signature  *gotypes.Signature
 	LocalSize  int
@@ -171,6 +178,14 @@ func NewFunction(name string) *Function {
 	}
 }
 
+// AddPragma adds a pragma to this function.
+func (f *Function) AddPragma(directive string, args ...string) {
+	f.Pragmas = append(f.Pragmas, Pragma{
+		Directive: directive,
+		Arguments: args,
+	})
+}
+
 // SetSignature sets the function signature.
 func (f *Function) SetSignature(s *gotypes.Signature) {
 	f.Signature = s

printer/stubs.go

@@ -28,7 +28,18 @@ func (s *stubs) Print(f *ir.File) ([]byte, error) {
 	for _, fn := range f.Functions() {
 		s.NL()
 		s.Comment(fn.Doc...)
+		for _, pragma := range fn.Pragmas {
+			s.pragma(pragma)
+		}
 		s.Printf("%s\n", fn.Stub())
 	}
 	return s.Result()
 }
+
+func (s *stubs) pragma(p ir.Pragma) {
+	s.Printf("//go:%s", p.Directive)
+	for _, arg := range p.Arguments {
+		s.Printf(" %s", arg)
+	}
+	s.NL()
+}

printer/stubs_test.go

@@ -0,0 +1,28 @@
+package printer_test
+
+import (
+	"testing"
+
+	"github.com/mmcloughlin/avo/build"
+	"github.com/mmcloughlin/avo/printer"
+)
+
+func TestStubsPragmas(t *testing.T) {
+	ctx := build.NewContext()
+	ctx.Function("f")
+	ctx.Pragma("noescape")
+	ctx.Pragma("linkname f remote.f")
+	ctx.SignatureExpr("func(x *uint64)")
+	ctx.RET()
+
+	AssertPrintsLines(t, ctx, printer.NewStubs, []string{
+		"// Code generated by avo. DO NOT EDIT.",
+		"",
+		"package printer",
+		"",
+		"//go:noescape",
+		"//go:linkname f remote.f",
+		"func f(x *uint64)",
+		"",
+	})
+}

printer/util_test.go

@@ -15,6 +15,7 @@ func AssertPrintsLines(t *testing.T, ctx *build.Context, pb printer.Builder, exp
 	lines := strings.Split(output, "\n")
 
 	if len(expect) != len(lines) {
+		t.Logf("output:\n%s", output)
 		t.Fatalf("have %d lines of output; expected %d", len(lines), len(expect))
 	}