Merge pull request #7583 from planetscale/deep-clone

Optimise AST rewriting
vitessio · Mar 2, 2021 · d2f939e · d2f939e
2 parents f28f6f1 + 0b963ec
commit d2f939e
Show file tree

Hide file tree

Showing 3 changed files with 162 additions and 99 deletions.
diff --git a/go/tools/asthelpergen/integration/rewriter.go b/go/tools/asthelpergen/integration/rewriter.go
diff --git a/go/tools/asthelpergen/rewriter_gen.go b/go/tools/asthelpergen/rewriter_gen.go
@@ -133,13 +133,14 @@ func (r *rewriterGen) createReplacementMethod(container, elem types.Type, x jen.
 			return func(newnode, container AST) {
 				container.(InterfaceSlice)[idx] = newnode.(AST)
 			}
-		}
-
+		}(x)
 	*/
-	return jen.Func().Params(jen.List(jen.Id("newNode"), jen.Id("container")).Id(r.ifaceName)).Block(
-		jen.Id("container").Assert(jen.Id(types.TypeString(container, noQualifier))).Add(x).Index(jen.Id("x")).Op("=").
-			Id("newNode").Assert(jen.Id(types.TypeString(elem, noQualifier))),
-	)
+	return jen.Func().Params(jen.Id("idx").Int()).Func().Params(jen.List(jen.Id(r.ifaceName), jen.Id(r.ifaceName))).Block(
+		jen.Return(jen.Func().Params(jen.List(jen.Id("newNode"), jen.Id("container")).Id(r.ifaceName))).Block(
+			jen.Id("container").Assert(jen.Id(types.TypeString(container, noQualifier))).Add(x).Index(jen.Id("idx")).Op("=").
+				Id("newNode").Assert(jen.Id(types.TypeString(elem, noQualifier))),
+		),
+	).Call(jen.Id("x"))
 }
 
 func (r *rewriterGen) createFile(pkgName string) (string, *jen.File) {