allow per field flags

Author: gobwas

cmd/gtrace/main.go

@@ -258,8 +258,8 @@ func main() {
 				for i, c := range v.List {
 					logf("#%d comment %q", i, c.Text)
 
-					text := strings.TrimPrefix(c.Text, "//gtrace:")
-					if c.Text != text {
+					text, ok := TrimConfigComment(c.Text)
+					if ok {
 						if item == nil {
 							item = &GenItem{
 								File: pkgFiles[i],
@@ -284,6 +284,7 @@ func main() {
 				}
 				return false
 			}
+
 			return true
 		})
 	}
@@ -298,20 +299,39 @@ func main() {
 		}
 		for _, field := range item.StructType.Fields.List {
 			name := field.Names[0].Name
-			if fn, ok := field.Type.(*ast.FuncType); ok {
-				f, err := buildFunc(info, fn)
-				if err != nil {
-					log.Printf(
-						"skipping hook %s due to error: %v",
-						name, err,
-					)
-					continue
+			fn, ok := field.Type.(*ast.FuncType)
+			if !ok {
+				continue
+			}
+			f, err := buildFunc(info, fn)
+			if err != nil {
+				log.Printf(
+					"skipping hook %s due to error: %v",
+					name, err,
+				)
+				continue
+			}
+			var config GenConfig
+			if doc := field.Doc; doc != nil {
+				for _, line := range doc.List {
+					text, ok := TrimConfigComment(line.Text)
+					if !ok {
+						continue
+					}
+					err := config.ParseComment(text)
+					if err != nil {
+						log.Fatalf(
+							"malformed comment string: %q: %v",
+							text, err,
+						)
+					}
 				}
-				t.Hooks = append(t.Hooks, Hook{
-					Name: name,
-					Func: f,
-				})
 			}
+			t.Hooks = append(t.Hooks, Hook{
+				Name: name,
+				Func: f,
+				Flag: item.GenConfig.Flag | config.Flag,
+			})
 		}
 		p.Traces = append(p.Traces, t)
 	}
@@ -412,6 +432,7 @@ type Trace struct {
 type Hook struct {
 	Name string
 	Func Func
+	Flag GenFlag
 }
 
 type Param struct {
@@ -442,44 +463,55 @@ const (
 	GenAll = ^GenFlag(0)
 )
 
-type GenItem struct {
-	File       *os.File
-	Ident      *ast.Ident
-	TypeSpec   *ast.TypeSpec
-	StructType *ast.StructType
-
+type GenConfig struct {
 	Flag GenFlag
 }
 
-func (x *GenItem) ParseComment(text string) (err error) {
+func TrimConfigComment(text string) (string, bool) {
+	s := strings.TrimPrefix(text, "//gtrace:")
+	if text != s {
+		return s, true
+	}
+	return "", false
+}
+
+func (g *GenConfig) ParseComment(text string) (err error) {
 	prefix, text := split(text, ' ')
 	switch prefix {
 	case "gen":
 	case "set":
-		return x.ParseParameter(text)
+		return g.ParseParameter(text)
 	default:
 		return fmt.Errorf("unknown prefix: %q", prefix)
 	}
 	return nil
 }
 
-func (x *GenItem) ParseParameter(text string) (err error) {
+func (g *GenConfig) ParseParameter(text string) (err error) {
 	text = strings.TrimSpace(text)
 	param, _ := split(text, '=')
 	if param == "" {
 		return nil
 	}
 	switch param {
 	case "shortcut":
-		x.Flag |= GenShortcut
+		g.Flag |= GenShortcut
 	case "context":
-		x.Flag |= GenContext
+		g.Flag |= GenContext
 	default:
 		return fmt.Errorf("unexpected parameter: %q", param)
 	}
 	return nil
 }
 
+type GenItem struct {
+	GenConfig
+	File       *os.File
+	Ident      *ast.Ident
+	TypeSpec   *ast.TypeSpec
+	StructType *ast.StructType
+}
+
 func split(s string, c byte) (s1, s2 string) {
 	i := strings.IndexByte(s, c)
 	if i == -1 {

cmd/gtrace/writer.go

@@ -87,10 +87,10 @@ func (w *Writer) Write(p Package) error {
 			}
 		}
 		for _, trace := range p.Traces {
-			if !trace.Flag.Has(GenShortcut) {
-				continue
-			}
 			for _, hook := range trace.Hooks {
+				if !hook.Flag.Has(GenShortcut) {
+					continue
+				}
 				if w.Stub {
 					w.stubHookShortcut(trace, hook)
 				} else {
@@ -203,10 +203,10 @@ func unwrapStruct(t types.Type) (n *types.Named, s *types.Struct) {
 	return
 }
 
-func (w *Writer) traceFuncImports(dst []dep, trace Trace, fn Func) []dep {
+func (w *Writer) funcImports(dst []dep, flag GenFlag, fn Func) []dep {
 	for _, p := range fn.Params {
 		dst = w.typeImports(dst, p.Type)
-		if !trace.Flag.Has(GenShortcut) {
+		if !flag.Has(GenShortcut) {
 			continue
 		}
 		if _, s := unwrapStruct(p.Type); s != nil {
@@ -218,7 +218,7 @@ func (w *Writer) traceFuncImports(dst []dep, trace Trace, fn Func) []dep {
 		}
 	}
 	for _, fn := range fn.Result {
-		dst = w.traceFuncImports(dst, trace, fn)
+		dst = w.funcImports(dst, flag, fn)
 	}
 	return dst
 }
@@ -232,7 +232,7 @@ func (w *Writer) traceImports(dst []dep, t Trace) []dep {
 		})
 	}
 	for _, h := range t.Hooks {
-		dst = w.traceFuncImports(dst, t, h.Func)
+		dst = w.funcImports(dst, h.Flag, h.Func)
 	}
 	return dst
 }

test/test_shortcut.go

@@ -0,0 +1,10 @@
+package test
+
+//go:generate gtrace -v
+
+//gtrace:gen
+type ShortcutPerFieldTrace struct {
+	//gtrace:set shortcut
+	OnFoo func()
+	OnBar func()
+}

test/test_shortcut_gtrace.go

@@ -186,6 +186,19 @@ func TestCompose(t *testing.T) {
 	}
 }
 
+func TestShortcutPerFieldTrace(t *testing.T) {
+	var called bool
+	t0 := ShortcutPerFieldTrace{
+		OnFoo: func() {
+			called = true
+		},
+	}
+	shortcutPerFieldTraceOnFoo(t0)
+	if !called {
+		t.Fatalf("hook wasn't called")
+	}
+}
+
 func TestBuildTagTrace(t *testing.T) {
 	t0 := BuildTagTrace{
 		OnSomethingA: func() func() {