add offset tracking when visiting slices

Signed-off-by: Andres Taylor <[email protected]>
vitessio · Jan 30, 2025 · f5fc6df · f5fc6df
1 parent f2f00da
commit f5fc6df
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 14 deletions.
diff --git a/go/tools/asthelpergen/integration/ast_rewrite.go b/go/tools/asthelpergen/integration/ast_rewrite.go
diff --git a/go/tools/asthelpergen/rewrite_gen.go b/go/tools/asthelpergen/rewrite_gen.go
@@ -222,7 +222,7 @@ func (r *rewriteGen) sliceMethod(t types.Type, slice *types.Slice, spi generator
 			),
 		}
 		stmts = append(stmts, forBlock...)
-		rewriteChild := r.rewriteChildSlice(t, slice.Elem(), "notUsed", jen.Id("el"), jen.Index(jen.Id("idx")), false)
+		rewriteChild := r.rewriteChildSlice(t, slice.Elem(), "Offset", jen.Id("el"), jen.Index(jen.Id("idx")), false)
 
 		stmts = append(stmts,
 			jen.For(jen.Id("x, el").Op(":=").Id("range node")).
@@ -336,12 +336,8 @@ func (r *rewriteGen) rewriteAllStructFields(t types.Type, strct *types.Struct, s
 			}
 
 			spi.addType(slice.Elem())
-			id := jen.Id("x")
-			if fail {
-				id = jen.Id("_")
-			}
 			output = append(output,
-				jen.For(jen.List(id, jen.Id("el")).Op(":=").Id("range node."+field.Name())).
+				jen.For(jen.List(jen.Id("x"), jen.Id("el")).Op(":=").Id("range node."+field.Name())).
 					Block(r.rewriteChildSlice(t, slice.Elem(), field.Name(), jen.Id("el"), jen.Dot(field.Name()).Index(jen.Id("idx")), fail)...))
 			fieldNumber++
 		}
@@ -444,7 +440,15 @@ func (r *rewriteGen) rewriteChildSlice(t, field types.Type, fieldName string, pa
 			param,
 			funcBlock).Block(returnFalse()))
 
+	savePath := jen.If(jen.Id("a.collectPaths")).Block(
+		jen.Id("a.cur.current").Op("=").Id("AddStepWithSliceIndex").Call(
+			jen.Id("path"),
+			jen.Id(printableTypeName(t)+fieldName+"8"),
+			jen.Id("x")),
+	)
+
 	return []jen.Code{
+		savePath,
 		rewriteField,
 	}
 }