support inline directive at anonymous function assignments

gcassert.go

@@ -128,14 +128,25 @@ func (v *assertVisitor) Visit(node ast.Node) ast.Visitor {
 					continue
 				}
 				if directive == inline {
+					// If the node is a FuncDecl or an assignment of an
+					// anonymous function, add the ident's Object to our map of
+					// must-inline functions.
+					var obj types.Object
 					switch n := node.(type) {
 					case *ast.FuncDecl:
-						// Add the Object that this FuncDecl's ident is connected
-						// to our map of must-inline functions.
-						obj := v.p.TypesInfo.Defs[n.Name]
-						if obj != nil {
-							v.mustInlineFuncs[obj] = struct{}{}
+						obj = v.p.TypesInfo.Defs[n.Name]
+					case *ast.AssignStmt:
+						if _, ok := n.Rhs[0].(*ast.FuncLit); !ok {
+							break
 						}
+						id, ok := n.Lhs[0].(*ast.Ident)
+						if !ok {
+							break
+						}
+						obj = v.p.TypesInfo.Defs[id]
+					}
+					if obj != nil {
+						v.mustInlineFuncs[obj] = struct{}{}
 						continue
 					}
 				}

gcassert_test.go

@@ -68,12 +68,17 @@ func badDirective3() {
 			23: {directives: []assertDirective{bce}},
 		},
 		"testdata/inline.go": {
-			45: {inlinableCallsites: []passInfo{{colNo: 15}}},
-			49: {directives: []assertDirective{inline}},
-			51: {directives: []assertDirective{inline}},
-			55: {directives: []assertDirective{inline}},
-			57: {inlinableCallsites: []passInfo{{colNo: 36}}},
-			58: {inlinableCallsites: []passInfo{{colNo: 35}}},
+			45:  {inlinableCallsites: []passInfo{{colNo: 15}}},
+			49:  {directives: []assertDirective{inline}},
+			51:  {directives: []assertDirective{inline}},
+			55:  {directives: []assertDirective{inline}},
+			57:  {inlinableCallsites: []passInfo{{colNo: 36}}},
+			58:  {inlinableCallsites: []passInfo{{colNo: 35}}},
+			85:  {inlinableCallsites: []passInfo{{colNo: 21}}},
+			86:  {inlinableCallsites: []passInfo{{colNo: 26}}},
+			88:  {directives: []assertDirective{inline}},
+			90:  {directives: []assertDirective{inline}},
+			102: {inlinableCallsites: []passInfo{{colNo: 12}}},
 		},
 		"testdata/noescape.go": {
 			11: {directives: []assertDirective{noescape}},
@@ -120,10 +125,12 @@ testdata/inline.go:55:	sum += 1: call was not inlined
 testdata/inline.go:58:	test(0).neverInlinedMethod(10): call was not inlined
 testdata/inline.go:60:	otherpkg.A{}.NeverInlined(sum): call was not inlined
 testdata/inline.go:62:	otherpkg.NeverInlinedFunc(sum): call was not inlined
+testdata/inline.go:86:	anonNeverInlined(i): call was not inlined
+testdata/inline.go:90:	sum += anonNeverInlinedNoAssert(i): call was not inlined
 testdata/issue5.go:4:	Gen().Layout(): call was not inlined
 `
 
-	testCases := []struct{
+	testCases := []struct {
 		name     string
 		pkgs     []string
 		cwd      string

testdata/inline.go

@@ -60,4 +60,59 @@ func caller() {
 	otherpkg.A{}.NeverInlined(sum)
 
 	otherpkg.NeverInlinedFunc(sum)
+
+	//gcassert:inline
+	anonInlined := func(a int) int {
+		return a + 10
+	}
+
+	anonInlinedNoAssert := func(a int) int {
+		return a + 10
+	}
+
+	//gcassert:inline
+	anonNeverInlined := func(a int) int {
+		fmt.Println(a)
+		return a + 10
+	}
+
+	anonNeverInlinedNoAssert := func(a int) int {
+		fmt.Println(a)
+		return a + 10
+	}
+
+	for i := 0; i < 10; i++ {
+		sum += anonInlined(i)
+		sum += anonNeverInlined(i)
+		//gcassert:inline
+		sum += anonInlinedNoAssert(i)
+		//gcassert:inline
+		sum += anonNeverInlinedNoAssert(i)
+	}
+
+	// The assertion for the anonymous function assigned to baz below does not
+	// apply to this call of the named function baz, which cannot be inlined and
+	// has no inline assertion.
+	sum += baz(1)
+
+	//gcassert:inline
+	baz := func(a int) int {
+		return a + 10
+	}
+	sum += baz(1)
+
+	{
+		// The assertion for the baz in the parent scope does not apply to this
+		// baz, which cannot be inlined and has no assertion.
+		baz := func(a int) int {
+			fmt.Println(a)
+			return a + 10
+		}
+		sum += baz(1)
+	}
+}
+
+//go:noinline
+func baz(a int) int {
+	return a + 10
 }