don't evict global declarations

Signed-off-by: Achille Roussel <[email protected]>

compiler/Makefile

@@ -1,6 +1,8 @@
 GO ?= go
 
-TARGET = testdata/coroutine_durable.go testdata/coroutine_functypes.go
+TARGET = \
+	testdata/coroutine_generated.go \
+	testdata/testdata_generated.go
 
 test: clean generate
 	$(GO) test ./...

compiler/cmd/coroc/main.go

@@ -17,8 +17,6 @@ USAGE:
 OPTIONS:
       --output <FILENAME>  Name of the Go file to generate in each package
 
-      --tags   <TAGS...>   Build tags to set on generated files
-
   -h, --help               Show this help information
 `
 
@@ -30,8 +28,6 @@ func main() {
 }
 
 func run() error {
-	buildTags := flag.String("tags", "", "")
-
 	flag.Usage = func() { println(usage[1:]) }
 	flag.Parse()
 
@@ -49,10 +45,5 @@ func run() error {
 		}
 	}
 
-	var options []compiler.Option
-	if *buildTags != "" {
-		options = append(options, compiler.WithBuildTags(*buildTags))
-	}
-
-	return compiler.Compile(path, options...)
+	return compiler.Compile(path)
 }

compiler/comment.go

@@ -1,6 +1,11 @@
 package compiler
 
-import "go/ast"
+import (
+	"go/ast"
+	"go/build/constraint"
+	"slices"
+	"strings"
+)
 
 func appendCommentGroup(comments []*ast.Comment, group *ast.CommentGroup) []*ast.Comment {
 	if group != nil && len(group.List) > 0 {
@@ -19,3 +24,35 @@ func appendComment(comments []*ast.Comment, text string) []*ast.Comment {
 		Text: text,
 	})
 }
+
+func commentGroupsOf(file *ast.File) []*ast.CommentGroup {
+	groups := make([]*ast.CommentGroup, 0, 1+len(file.Comments))
+	groups = append(groups, file.Comments...)
+	if file.Doc != nil {
+		groups = append(groups, file.Doc)
+	}
+	return groups
+}
+
+const gobuild = "//go:build"
+
+func parseBuildTags(file *ast.File) (constraint.Expr, error) {
+	for _, group := range commentGroupsOf(file) {
+		for _, comment := range group.List {
+			if strings.HasPrefix(comment.Text, gobuild) {
+				return constraint.Parse(comment.Text)
+			}
+		}
+	}
+	return nil, nil
+}
+
+func stripBuildTagsOf(file *ast.File) {
+	for _, group := range commentGroupsOf(file) {
+		for i, comment := range group.List {
+			if strings.HasPrefix(comment.Text, gobuild) {
+				group.List = slices.Delete(group.List, i, i+1)
+			}
+		}
+	}
+}

compiler/compile.go

@@ -3,6 +3,7 @@ package compiler
 import (
 	"fmt"
 	"go/ast"
+	"go/build/constraint"
 	"go/format"
 	"go/token"
 	"go/types"
@@ -44,15 +45,7 @@ func Compile(path string, options ...Option) error {
 // Option configures the compiler.
 type Option func(*compiler)
 
-// WithBuildTags instructs the compiler to attach the specified build
-// tags to generated files.
-func WithBuildTags(buildTags string) Option {
-	return func(c *compiler) { c.buildTags = buildTags }
-}
-
 type compiler struct {
-	buildTags string
-
 	fset *token.FileSet
 }
 
@@ -192,20 +185,26 @@ func (c *compiler) compile(path string) error {
 	return nil
 }
 
-func (c *compiler) writeFile(path string, file *ast.File) error {
+func (c *compiler) writeFile(path string, file *ast.File, changeBuildTags func(constraint.Expr) constraint.Expr) error {
 	f, err := os.Create(path)
 	if err != nil {
 		return err
 	}
 	defer f.Close()
+
+	buildTags, err := parseBuildTags(file)
+	if err != nil {
+		return err
+	}
+	buildTags = changeBuildTags(buildTags)
+	stripBuildTagsOf(file)
+
 	// Comments are awkward to attach to the tree (they rely on token.Pos, which
 	// is coupled to a token.FileSet). Instead, just write out the raw strings.
 	var b strings.Builder
-	b.WriteString(`// Code generated by coroc. DO NOT EDIT`)
-	b.WriteString("\n\n")
-	if c.buildTags != "" {
+	if buildTags != nil {
 		b.WriteString(`//go:build `)
-		b.WriteString(c.buildTags)
+		b.WriteString(buildTags.String())
 		b.WriteString("\n\n")
 	}
 	if _, err := f.WriteString(b.String()); err != nil {
@@ -233,37 +232,50 @@ func (c *compiler) compilePackage(p *packages.Package, colors functionColors) er
 		colorsByFunc[decl] = color
 	}
 
+	buildTag := &constraint.TagExpr{
+		Tag: "durable",
+	}
+
 	for i, f := range p.Syntax {
+		if err := c.writeFile(p.GoFiles[i], f, func(expr constraint.Expr) constraint.Expr {
+			return withoutBuildTag(expr, buildTag)
+		}); err != nil {
+			return err
+		}
+
 		// Generate the coroutine AST.
 		gen := &ast.File{
 			Name: ast.NewIdent(p.Name),
 		}
 
 		for _, anydecl := range f.Decls {
-			decl, ok := anydecl.(*ast.FuncDecl)
-			if !ok {
-				continue
-			}
-			color, ok := colorsByFunc[decl]
-			if !ok {
-				gen.Decls = append(gen.Decls, decl)
-				continue
-			}
-			// Reject certain language features for now.
-			if err := unsupported(decl, p.TypesInfo); err != nil {
-				return err
-			}
+			switch decl := anydecl.(type) {
+			case *ast.GenDecl:
+				// Imports get re-added by addImports below, so no need to carry
+				// them from declarations in the input file.
+				if decl.Tok != token.IMPORT {
+					gen.Decls = append(gen.Decls, decl)
+					continue
+				}
 
-			scope := &scope{colors: colorsByFunc}
-			// If the function has a single expression it does not contain a
-			// deferred closure; it won't be added to the list of colored
-			// functions so generateFunctypes does not mistakenly increment the
-			// local symbol counter when generating closure names.
-			gen.Decls = append(gen.Decls, scope.compileFuncDecl(p, decl, color))
-		}
+			case *ast.FuncDecl:
+				color, ok := colorsByFunc[decl]
+				if !ok {
+					gen.Decls = append(gen.Decls, decl)
+					continue
+				}
+				// Reject certain language features for now.
+				if err := unsupported(decl, p.TypesInfo); err != nil {
+					return err
+				}
 
-		if len(gen.Decls) == 0 {
-			continue
+				scope := &scope{colors: colorsByFunc}
+				// If the function has a single expression it does not contain a
+				// deferred closure; it won't be added to the list of colored
+				// functions so generateFunctypes does not mistakenly increment the
+				// local symbol counter when generating closure names.
+				gen.Decls = append(gen.Decls, scope.compileFuncDecl(p, decl, color))
+			}
 		}
 
 		clearPos(gen)
@@ -273,13 +285,12 @@ func (c *compiler) compilePackage(p *packages.Package, colors functionColors) er
 		// Find all the required imports for this file.
 		gen = addImports(p, gen)
 
-		outputPath, _ := strings.CutSuffix(p.GoFiles[i], ".go")
-		if c.buildTags != "" {
-			outputPath += "_" + strings.ReplaceAll(c.buildTags, ",", "_") + ".go"
-		} else {
-			outputPath += "_generated.go"
-		}
-		if err := c.writeFile(outputPath, gen); err != nil {
+		outputPath := strings.TrimSuffix(p.GoFiles[i], ".go")
+		outputPath += "_generated.go"
+
+		if err := c.writeFile(outputPath, gen, func(expr constraint.Expr) constraint.Expr {
+			return withBuildTag(expr, buildTag)
+		}); err != nil {
 			return err
 		}
 	}

compiler/constraint.go

@@ -0,0 +1,36 @@
+package compiler
+
+import (
+	"go/build/constraint"
+	"reflect"
+)
+
+func containsExpr(expr, contains constraint.Expr) bool {
+	switch x := expr.(type) {
+	case *constraint.AndExpr:
+		return containsExpr(x.X, contains) || containsExpr(x.Y, contains)
+	default:
+		return reflect.DeepEqual(expr, contains)
+	}
+}
+
+func withBuildTag(expr constraint.Expr, buildTag *constraint.TagExpr) constraint.Expr {
+	if buildTag == nil || containsExpr(expr, buildTag) {
+		return expr
+	} else if expr == nil {
+		return buildTag
+	} else {
+		return &constraint.AndExpr{X: expr, Y: buildTag}
+	}
+}
+
+func withoutBuildTag(expr constraint.Expr, buildTag *constraint.TagExpr) constraint.Expr {
+	notBuildTag := &constraint.NotExpr{X: buildTag}
+	if buildTag == nil || containsExpr(expr, notBuildTag) {
+		return expr
+	} else if expr == nil {
+		return notBuildTag
+	} else {
+		return &constraint.AndExpr{X: expr, Y: notBuildTag}
+	}
+}

compiler/function.go

@@ -240,14 +240,16 @@ func generateFunctypes(p *packages.Package, f *ast.File, colors map[ast.Node]*ty
 		})
 	}
 
-	astutil.AddNamedImport(nil, f, "_types", "github.com/stealthrocket/coroutine/types")
-
-	f.Decls = append(f.Decls,
-		&ast.FuncDecl{
-			Name: ast.NewIdent("init"),
-			Type: &ast.FuncType{Params: new(ast.FieldList)},
-			Body: init,
-		})
+	if len(init.List) > 0 {
+		astutil.AddNamedImport(nil, f, "_types", "github.com/stealthrocket/coroutine/types")
+
+		f.Decls = append(f.Decls,
+			&ast.FuncDecl{
+				Name: ast.NewIdent("init"),
+				Type: &ast.FuncType{Params: new(ast.FieldList)},
+				Body: init,
+			})
+	}
 }
 
 // This function computes the name that the linker gives to anonymous functions,

compiler/testdata/coroutine.go

@@ -9,7 +9,7 @@ import (
 	"github.com/stealthrocket/coroutine"
 )
 
-//go:generate coroc --tags durable
+//go:generate coroc
 
 func SomeFunctionThatShouldExistInTheCompiledFile() {
 }