Support pre-processor macros and includes

build/context.go

@@ -56,6 +56,12 @@ func (c *Context) Package(path string) {
 	c.pkg = pkg
 }
 
+// IncludeHeader adds an include pre-processor directive for the provided path, if it
+// is not already present.
+func (c *Context) IncludeHeader(path string) {
+	c.file.IncludeHeader(path)
+}
+
 // Constraints sets build constraints for the file.
 func (c *Context) Constraints(t buildtags.ConstraintsConvertable) {
 	cs := t.ToConstraints()
@@ -160,6 +166,26 @@ func (c *Context) Label(name string) {
 	c.activefunc().AddLabel(ir.Label(name))
 }
 
+// PreprocIfdef adds an #ifdef macro to the active function.
+func (c *Context) PreprocIfdef(condition string) {
+	c.activefunc().AddPreprocIfdef(condition)
+}
+
+// PreprocIfndef adds an #ifndef macro to the active function.
+func (c *Context) PreprocIfndef(condition string) {
+	c.activefunc().AddPreprocIfndef(condition)
+}
+
+// PreprocElse adds an #else macro to the active function.
+func (c *Context) PreprocElse() {
+	c.activefunc().AddPreprocElse()
+}
+
+// PreprocEndif adds an #endif macro to the active function.
+func (c *Context) PreprocEndif() {
+	c.activefunc().AddPreprocEndif()
+}
+
 // Comment adds comment lines to the active function.
 func (c *Context) Comment(lines ...string) {
 	c.activefunc().AddComment(lines...)

build/global.go

@@ -15,6 +15,11 @@ import (
 // ctx provides a global build context.
 var ctx = NewContext()
 
+// Returns the global context.
+func GlobalContext() *Context {
+	return ctx
+}
+
 // TEXT starts building a new function called name, with attributes a, and sets its signature (see SignatureExpr).
 func TEXT(name string, a attr.Attribute, signature string) {
 	ctx.Function(name)
@@ -55,6 +60,10 @@ func Generate() {
 	}
 }
 
+// IncludeHeader adds an include pre-processor directive for the provided path
+// to the current file, if it is not already present.
+func IncludeHeader(path string) { ctx.IncludeHeader(path) }
+
 // Package sets the package the generated file will belong to. Required to be able to reference types in the package.
 func Package(path string) { ctx.Package(path) }
 

ir/ir.go

@@ -15,6 +15,48 @@ type Node interface {
 	node()
 }
 
+// PreprocIfDef is the #ifdef preprocessor macro node, represented as its
+// condition expression.
+type PreprocIfdef string
+
+func (p PreprocIfdef) node() {}
+
+// NewPreprocIfdef builds a PreprocIfdef node from the provided condition.
+func NewPreprocIfdef(condition string) PreprocIfdef {
+	return PreprocIfdef(condition)
+}
+
+// PreprocIfndef is the #ifndef preprocessor macro node, represented as its
+// condition expression.
+type PreprocIfndef string
+
+func (p PreprocIfndef) node() {}
+
+// NewPreprocIfndef builds a PreprocIfndef node from the provided condition.
+func NewPreprocIfndef(condition string) PreprocIfndef {
+	return PreprocIfndef(condition)
+}
+
+// PreprocElse is the #else preprocessor macro node.
+type PreprocElse struct{}
+
+func (p PreprocElse) node() {}
+
+// NewPreprocElse builds a PreprocElse node.
+func NewPreprocElse() PreprocElse {
+	return PreprocElse{}
+}
+
+// PreprocEndif is the #endif preprocessor macro node.
+type PreprocEndif struct{}
+
+func (p PreprocEndif) node() {}
+
+// NewPreprocEnd builds a PreprocEnd node.
+func NewPreprocEndif() PreprocEndif {
+	return PreprocEndif{}
+}
+
 // Label within a function.
 type Label string
 
@@ -162,6 +204,17 @@ func (f *File) Functions() []*Function {
 	return fns
 }
 
+// IncludeHeader adds an include pre-processor directive for the provided path,
+// if it is not already present.
+func (f *File) IncludeHeader(path string) {
+	for _, p := range f.Includes {
+		if p == path {
+			return
+		}
+	}
+	f.Includes = append(f.Includes, path)
+}
+
 // Pragma represents a function compiler directive.
 type Pragma struct {
 	Directive string
@@ -230,6 +283,26 @@ func (f *Function) AddLabel(l Label) {
 	f.AddNode(l)
 }
 
+// AddPreprocIfdef adds an #ifdef preprocessor macro to f.
+func (f *Function) AddPreprocIfdef(condition string) {
+	f.AddNode(NewPreprocIfdef(condition))
+}
+
+// AddPreprocIfndef adds an #ifndef preprocessor macro to f.
+func (f *Function) AddPreprocIfndef(condition string) {
+	f.AddNode(NewPreprocIfndef(condition))
+}
+
+// AddPreprocElse adds an #else preprocessor macro to f.
+func (f *Function) AddPreprocElse() {
+	f.AddNode(NewPreprocElse())
+}
+
+// AddPreprocEndif adds an #endif preprocessor macro to f.
+func (f *Function) AddPreprocEndif() {
+	f.AddNode(NewPreprocEndif())
+}
+
 // AddComment adds comment lines to f.
 func (f *Function) AddComment(lines ...string) {
 	f.AddNode(NewComment(lines...))

preproc/cond.go

@@ -0,0 +1,27 @@
+package preproc
+
+import "github.com/mmcloughlin/avo/build"
+
+// Ifdef inserts an #ifdef preprocessor macro. Condition should be a valid
+// preprocessor expression, though this is not checked.
+func Ifdef(condition string) {
+	build.GlobalContext().PreprocIfdef(condition)
+}
+
+// Ifndef inserts an #ifndef preprocessor macro. Condition should be a valid
+// preprocessor expression, though this is not checked.
+func Ifndef(condition string) {
+	build.GlobalContext().PreprocIfndef(condition)
+}
+
+// Else inserts an #else preprocessor macro. It needs to be preceded by an
+// #ifdef or #ifndef macro.
+func Else() {
+	build.GlobalContext().PreprocElse()
+}
+
+// End inserts an #endif preprocessor macro. It needs to be preceded by an
+// #ifdef, #ifndef, or #else macro.
+func Endif() {
+	build.GlobalContext().PreprocEndif()
+}

printer/goasm.go

@@ -111,6 +111,18 @@ func (p *goasm) function(f *ir.Function) {
 			for _, line := range n.Lines {
 				p.Printf("\t// %s\n", line)
 			}
+		case ir.PreprocIfdef:
+			p.flush()
+			p.Printf("#ifdef %s\n", string(n))
+		case ir.PreprocIfndef:
+			p.flush()
+			p.Printf("#ifndef %s\n", string(n))
+		case ir.PreprocElse:
+			p.flush()
+			p.Printf("#else\n")
+		case ir.PreprocEndif:
+			p.flush()
+			p.Printf("#endif\n")
 		default:
 			panic("unexpected node type")
 		}

printer/goasm_test.go

@@ -135,3 +135,51 @@ func TestOpcodeSuffixes(t *testing.T) {
 		"",
 	})
 }
+
+func TestPreprocessor(t *testing.T) {
+	ctx := build.NewContext()
+	ctx.Function("preprocessor")
+	ctx.SignatureExpr("func()")
+	ctx.PreprocIfndef("hasAVX2")
+	ctx.ADDQ(reg.RBX, reg.RAX)
+	ctx.VZEROUPPER()
+	ctx.PreprocElse()
+	ctx.ADDQ(reg.RAX, reg.RBX)
+	ctx.PreprocEndif()
+
+	AssertPrintsLines(t, ctx, printer.NewGoAsm, []string{
+		"// Code generated by avo. DO NOT EDIT.",
+		"",
+		"// func preprocessor()",
+		"TEXT ·preprocessor(SB), $0",
+		"#ifndef hasAVX2",
+		"\tADDQ BX, AX",
+		"\tVZEROUPPER",
+		"#else",
+		"\tADDQ AX, BX",
+		"#endif",
+		"",
+	})
+}
+
+func TestInclude(t *testing.T) {
+	ctx := build.NewContext()
+	ctx.IncludeHeader("before.h")
+	ctx.Function("preprocessor")
+	ctx.SignatureExpr("func()")
+	ctx.IncludeHeader("after.h")
+	ctx.VZEROUPPER()
+
+	AssertPrintsLines(t, ctx, printer.NewGoAsm, []string{
+		"// Code generated by avo. DO NOT EDIT.",
+		"",
+		"#include \"before.h\"",
+		"#include \"after.h\"",
+		"",
+		"// func preprocessor()",
+		"TEXT ·preprocessor(SB), $0",
+		"\tVZEROUPPER",
+		"",
+	})
+
+}