feat: gethclone patch set for stateful precompiles

x/gethclone/astpatch/astpatch.go

@@ -15,6 +15,12 @@ type (
 	// A non-nil error is equivalent to returning false and will also abort all
 	// further calls to other patches.
 	Patch func(*astutil.Cursor) error
+	// A TypedPatch functions identically to a `Patch` except that it also
+	// receives the `ast.Node` as its concrete type.
+	//
+	// Invariant: `c.Node()` and `n` MUST be of the same concrete type and point
+	// to the same memory; i.e. `c.Node().(N) == n` doesn't panic and is true.
+	TypedPatch[N ast.Node] func(c *astutil.Cursor, n N) error
 	// A PatchRegistry maps [Go package path] -> [ast.Node concrete types] ->
 	// [all `Patch` functions that must be applied to said node types in said
 	// package].
@@ -23,15 +29,35 @@ type (
 	PatchRegistry map[string]map[reflect.Type][]Patch
 )
 
-// Add is a convenience wrapper for registering a new `Patch` in the registry.
-// The `zeroNode` can be any type (including nil pointers) that implements
-// `ast.Node`.
+// Apply is equivalent to `astutil.ApplyFunc()` except that it accepts
+// `Patch`es. See `Patch` comment for error-handling semantics.
+func Apply(root ast.Node, pre, post Patch) (ast.Node, error) {
+	var err error
+	x := func(p Patch) astutil.ApplyFunc {
+		return func(c *astutil.Cursor) bool {
+			if err != nil {
+				return false
+			}
+			if p == nil {
+				return true
+			}
+			err = p(c)
+			return err == nil
+		}
+	}
+	n := astutil.Apply(root, x(pre), x(post))
+	return n, err
+}
+
+// AddForType is a convenience wrapper for registering a new `Patch` in the
+// registry. The `zeroNode` can be any type (including nil pointers) that
+// implements `ast.Node`.
 //
 // The special `pkgPath` value "*" will match all package paths. While there is
 // no specific requirement for `pkgPath` other than it matching the equivalent
-// argument passed to `Apply()`, it is typically sourced from
-// `golang.org/x/tools/go/packages.Package.PkgPath`.
-func (r PatchRegistry) Add(pkgPath string, zeroNode ast.Node, fn Patch) {
+// argument passed to `Apply()`, it is typically the import path of the package
+// being patched.
+func (r PatchRegistry) AddForType(pkgPath string, zeroNode ast.Node, fn Patch) {
 	pkg, ok := r[pkgPath]
 	if !ok {
 		pkg = make(map[reflect.Type][]Patch)
@@ -42,6 +68,30 @@ func (r PatchRegistry) Add(pkgPath string, zeroNode ast.Node, fn Patch) {
 	pkg[t] = append(pkg[t], fn)
 }
 
+// A Patcher couples a `Patch` with the specific `ast.Node` type to which it
+// applies. It is useful when `PatchRegistry.AddForType()` MUST receive a
+// specific `Node` type for a particular `Patch`, in which case
+// `PatchRegistry.Add()` SHOULD be used instead.
+type Patcher interface {
+	Type() ast.Node
+	Patch(*astutil.Cursor) error
+}
+
+// Add is a synonym of `AddForType()`, instead accepting an argument that
+// provides the `Node` type and the `Patch`.
+func (r PatchRegistry) Add(pkgPath string, tp Patcher) {
+	r.AddForType(pkgPath, tp.Type(), tp.Patch)
+}
+
+// patcher implements the `Patcher` interface.
+type patcher struct {
+	typ   ast.Node
+	patch Patch
+}
+
+func (p patcher) Type() ast.Node                { return p.typ }
+func (p patcher) Patch(c *astutil.Cursor) error { return p.patch(c) }
+
 // Apply calls `astutil.Apply()` on `node`, calling the appropriate `Patch`
 // functions as the syntax tree is traversed. Patches are applied as the `pre`
 // argument to `astutil.Apply()`.
@@ -51,19 +101,13 @@ func (r PatchRegistry) Add(pkgPath string, zeroNode ast.Node, fn Patch) {
 //
 // If any `Patch` returns an error then no further patches will be called, and
 // the error will be returned by `Apply()`.
-func (r PatchRegistry) Apply(pkgPath string, node ast.Node) error {
-	var err error
-	astutil.Apply(node, func(c *astutil.Cursor) bool {
-		if err != nil {
-			return false
-		}
-		if err = r.applyToCursor("*", c); err != nil {
-			return false
+func (r PatchRegistry) Apply(pkgPath string, node ast.Node) (ast.Node, error) {
+	return Apply(node, func(c *astutil.Cursor) error {
+		if err := r.applyToCursor("*", c); err != nil {
+			return err
 		}
-		err = r.applyToCursor(pkgPath, c)
-		return err == nil
+		return r.applyToCursor(pkgPath, c)
 	}, nil)
-	return err
 }
 
 func (r PatchRegistry) applyToCursor(pkgPath string, c *astutil.Cursor) error {