add simple 1-stmt level normalization

Signed-off-by: Iskander Sharipov <[email protected]>

astnorm.go

@@ -12,3 +12,7 @@ type Config struct {
 func Expr(cfg *Config, x ast.Expr) ast.Expr {
 	return newNormalizer(cfg).normalizeExpr(x)
 }
+
+func Stmt(cfg *Config, x ast.Stmt) ast.Stmt {
+	return newNormalizer(cfg).normalizeStmt(x)
+}

normalizer.go

@@ -2,7 +2,10 @@ package astnorm
 
 import (
 	"go/ast"
+	"go/token"
 
+	"github.com/go-toolsmith/astcast"
+	"github.com/go-toolsmith/astequal"
 	"github.com/go-toolsmith/typep"
 )
 
@@ -35,3 +38,42 @@ func (n *normalizer) normalizeExprList(xs []ast.Expr) []ast.Expr {
 	}
 	return xs
 }
+
+func (n *normalizer) normalizeStmt(x ast.Stmt) ast.Stmt {
+	switch x := x.(type) {
+	case *ast.AssignStmt:
+		return n.normalizeAssignStmt(x)
+	case *ast.BlockStmt:
+		return n.normalizeBlockStmt(x)
+	default:
+		return x
+	}
+}
+
+func (n *normalizer) normalizeBlockStmt(b *ast.BlockStmt) *ast.BlockStmt {
+	for i, x := range b.List {
+		b.List[i] = n.normalizeStmt(x)
+	}
+	return b
+}
+
+func (n *normalizer) normalizeAssignStmt(assign *ast.AssignStmt) ast.Stmt {
+	assign = n.normalizeAssignOp(assign)
+	return assign
+}
+
+func (n *normalizer) normalizeAssignOp(assign *ast.AssignStmt) *ast.AssignStmt {
+	if assign.Tok != token.ASSIGN || len(assign.Lhs) != 1 {
+		return assign
+	}
+	rhs := astcast.ToBinaryExpr(assign.Rhs[0])
+	if !astequal.Expr(assign.Lhs[0], rhs.X) {
+		return assign
+	}
+	switch rhs.Op {
+	case token.ADD:
+		assign.Tok = token.ADD_ASSIGN
+		assign.Rhs[0] = n.normalizeExpr(rhs.Y)
+	}
+	return assign
+}

normalizer_test.go

@@ -11,6 +11,13 @@ import (
 	"golang.org/x/tools/go/packages"
 )
 
+func isTestCase(assign *ast.AssignStmt) bool {
+	return len(assign.Lhs) == 2 &&
+		len(assign.Rhs) == 2 &&
+		astcast.ToIdent(assign.Lhs[0]).Name == "_" &&
+		astcast.ToIdent(assign.Lhs[1]).Name == "_"
+}
+
 func TestNormalizeExpr(t *testing.T) {
 	pkg := loadPackage(t, "./testdata/normalize_expr.go")
 	funcs := collectFuncDecls(pkg)
@@ -19,13 +26,7 @@ func TestNormalizeExpr(t *testing.T) {
 	for _, fn := range funcs {
 		for _, stmt := range fn.Body.List {
 			assign, ok := stmt.(*ast.AssignStmt)
-			if !ok || len(assign.Lhs) != 2 || len(assign.Rhs) != 2 {
-				continue
-			}
-			if astcast.ToIdent(assign.Lhs[0]).Name != "_" {
-				continue
-			}
-			if astcast.ToIdent(assign.Lhs[1]).Name != "_" {
+			if !ok || !isTestCase(assign) {
 				continue
 			}
 			input := assign.Rhs[0]
@@ -40,6 +41,29 @@ func TestNormalizeExpr(t *testing.T) {
 	}
 }
 
+func TestNormalizeStmt(t *testing.T) {
+	pkg := loadPackage(t, "./testdata/normalize_stmt.go")
+	funcs := collectFuncDecls(pkg)
+	cfg := &Config{Info: pkg.TypesInfo}
+
+	for _, fn := range funcs {
+		for _, stmt := range fn.Body.List {
+			assign, ok := stmt.(*ast.AssignStmt)
+			if !ok || !isTestCase(assign) {
+				continue
+			}
+			input := assign.Rhs[0].(*ast.FuncLit).Body
+			want := assign.Rhs[1].(*ast.FuncLit).Body
+			have := Stmt(cfg, input)
+			if !astequal.Stmt(have, want) {
+				pos := pkg.Fset.Position(assign.Pos())
+				t.Errorf("%s:\nhave: %s\nwant: %s",
+					pos, astfmt.Sprint(have), astfmt.Sprint(want))
+			}
+		}
+	}
+}
+
 func collectFuncDecls(pkg *packages.Package) []*ast.FuncDecl {
 	var funcs []*ast.FuncDecl
 	for _, f := range pkg.Syntax {

testdata/normalize_stmt.go

@@ -0,0 +1,31 @@
+package normalize_stmt
+
+func identityTest() {
+	var x int
+
+	_, _ = func() {
+		x += 1
+	}, func() {
+		x += 1
+	}
+}
+
+func assignOpTest() {
+	var x int
+
+	_, _ = func() {
+		x = x + 1
+	}, func() {
+		x += 1
+	}
+}
+
+func combinedTest() {
+	var x int
+
+	_, _ = func() {
+		x = x + (1)
+	}, func() {
+		x += 1
+	}
+}