gen4: Faster analysis

go/tools/asthelpergen/integration/integration_rewriter_test.go

@@ -164,6 +164,46 @@ func TestRewriteVisitInterfaceSlice(t *testing.T) {
 	})
 }
 
+func TestRewriteAndRevisitInterfaceSlice(t *testing.T) {
+	leaf1 := &Leaf{2}
+	leaf2 := &Leaf{3}
+	ast := InterfaceSlice{
+		leaf1,
+		leaf2,
+	}
+	ast2 := InterfaceSlice{
+		leaf2,
+		leaf1,
+	}
+
+	tv := &rewriteTestVisitor{}
+
+	a := false
+	_ = Rewrite(ast, func(cursor *Cursor) bool {
+		tv.pre(cursor)
+		switch cursor.node.(type) {
+		case InterfaceSlice:
+			if a {
+				break
+			}
+			a = true
+			cursor.ReplaceAndRevisit(ast2)
+		}
+		return true
+	}, tv.post)
+
+	tv.assertEquals(t, []step{
+		Pre{ast}, // when we visit ast, we want to replace and revisit,
+		// which means that we don't do a post on this node, or visit the children
+		Pre{ast2},
+		Pre{leaf2},
+		Post{leaf2},
+		Pre{leaf1},
+		Post{leaf1},
+		Post{ast2},
+	})
+}
+
 func TestRewriteVisitRefContainerReplace(t *testing.T) {
 	ast := &RefContainer{
 		ASTType:               &RefContainer{NotASTType: 12},

go/tools/asthelpergen/integration/test_helpers.go

@@ -47,6 +47,8 @@ type Cursor struct {
 	parent   AST
 	replacer replacerFunc
 	node     AST
+	// marks that the node has been replaced, and the new node should be visited
+	revisit bool
 }
 
 // Node returns the current Node.
@@ -55,13 +57,30 @@ func (c *Cursor) Node() AST { return c.node }
 // Parent returns the parent of the current Node.
 func (c *Cursor) Parent() AST { return c.parent }
 
-// Replace replaces the current node in the parent field with this new object. The use needs to make sure to not
+// Replace replaces the current node in the parent field with this new object. The user needs to make sure to not
 // replace the object with something of the wrong type, or the visitor will panic.
 func (c *Cursor) Replace(newNode AST) {
 	c.replacer(newNode, c.parent)
 	c.node = newNode
 }
 
+// ReplaceAndRevisit replaces the current node in the parent field with this new object.
+// When used, this will abort the visitation of the current node - no post or children visited,
+// and the new node visited.
+func (c *Cursor) ReplaceAndRevisit(newNode AST) {
+	switch newNode.(type) {
+	case InterfaceSlice:
+	default:
+		// We need to add support to the generated code for when to look at the revisit flag. At the moment it is only
+		// there for slices of AST implementations
+		panic("no support added for this type yet")
+	}
+
+	c.replacer(newNode, c.parent)
+	c.node = newNode
+	c.revisit = true
+}
+
 type replacerFunc func(newNode, parent AST)
 
 // Rewrite is the api.

go/tools/asthelpergen/rewrite_gen.go

@@ -178,7 +178,21 @@ func (r *rewriteGen) sliceMethod(t types.Type, slice *types.Slice, spi generator
 	stmts := []jen.Code{
 		jen.If(jen.Id("node == nil").Block(returnTrue())),
 	}
-	stmts = append(stmts, executePre())
+
+	typeString := types.TypeString(t, noQualifier)
+
+	preStmts := setupCursor()
+	preStmts = append(preStmts,
+		jen.Id("kontinue").Op(":=").Id("!a.pre(&a.cur)"),
+		jen.If(jen.Id("a.cur.revisit").Block(
+			jen.Id("node").Op("=").Id("a.cur.node.("+typeString+")"),
+			jen.Id("a.cur.revisit").Op("=").False(),
+			jen.Return(jen.Id("a.rewrite"+typeString+"(parent, node, replacer)")),
+		)),
+		jen.If(jen.Id("kontinue").Block(jen.Return(jen.True()))),
+	)
+
+	stmts = append(stmts, jen.If(jen.Id("a.pre!= nil").Block(preStmts...)))
 
 	haveChildren := false
 	if shouldAdd(slice.Elem(), spi.iface()) {