From 4aaf6bc7ed8e103e97b5bb5fe5744831035db7f1 Mon Sep 17 00:00:00 2001
From: Michael McLoughlin <mmcloughlin@gmail.com>
Date: Mon, 31 Dec 2018 00:08:48 -0800
Subject: [PATCH] ast,printer: add Includes field to File

Updates #12
---
 ast.go           |  5 ++++-
 printer/goasm.go | 13 +++++++++----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/ast.go b/ast.go
index e852dfe..a332726 100644
--- a/ast.go
+++ b/ast.go
@@ -94,11 +94,14 @@ type Section interface {
 
 // File represents an assembly file.
 type File struct {
+	Includes []string
 	Sections []Section
 }
 
 func NewFile() *File {
-	return &File{}
+	return &File{
+		Includes: []string{"textflag.h"},
+	}
 }
 
 func (f *File) AddSection(s Section) {
diff --git a/printer/goasm.go b/printer/goasm.go
index adc7217..2685062 100644
--- a/printer/goasm.go
+++ b/printer/goasm.go
@@ -23,6 +23,7 @@ func NewGoAsm(cfg Config) Printer {
 
 func (p *goasm) Print(f *avo.File) ([]byte, error) {
 	p.header()
+	p.includes(f.Includes)
 	for _, s := range f.Sections {
 		switch s := s.(type) {
 		case *avo.Function:
@@ -38,12 +39,16 @@ func (p *goasm) Print(f *avo.File) ([]byte, error) {
 
 func (p *goasm) header() {
 	p.Comment(p.cfg.GeneratedWarning())
-	p.NL()
-	p.include("textflag.h")
 }
 
-func (p *goasm) include(path string) {
-	p.Printf("#include \"%s\"\n", path)
+func (p *goasm) includes(paths []string) {
+	if len(paths) == 0 {
+		return
+	}
+	p.NL()
+	for _, path := range paths {
+		p.Printf("#include \"%s\"\n", path)
+	}
 }
 
 func (p *goasm) function(f *avo.Function) {