diff --git a/go/tools/asthelpergen/asthelpergen.go b/go/tools/asthelpergen/asthelpergen.go index 549469071eb..fa2b2b78c22 100644 --- a/go/tools/asthelpergen/asthelpergen.go +++ b/go/tools/asthelpergen/asthelpergen.go @@ -14,13 +14,34 @@ See the License for the specific language governing permissions and limitations under the License. */ +// Package asthelpergen provides code generation for AST (Abstract Syntax Tree) helper methods. +// +// This package automatically generates helper methods for AST nodes including: +// - Deep cloning (Clone methods) +// - Equality comparison (Equals methods) +// - Visitor pattern support (Visit methods) +// - AST rewriting/transformation (Rewrite methods) +// - Path enumeration for navigation +// - Copy-on-write functionality +// +// The generator works by discovering all types that implement a root interface and +// then generating the appropriate helper methods for each type using a plugin architecture. +// +// Usage: +// +// result, err := asthelpergen.GenerateASTHelpers(&asthelpergen.Options{ +// Packages: []string{"./mypackage"}, +// RootInterface: "github.com/example/mypackage.MyASTInterface", +// }) +// +// The generated code follows Go conventions and includes proper error handling, +// nil checks, and type safety. package asthelpergen import ( "bytes" "fmt" "go/types" - "log" "os" "path" "strings" @@ -32,7 +53,13 @@ import ( "vitess.io/vitess/go/tools/codegen" ) -const licenseFileHeader = `Copyright 2025 The Vitess Authors. +const ( + // Common constants used across generators + visitableName = "Visitable" + anyTypeName = "any" + + // License header for generated files + licenseFileHeader = `Copyright 2025 The Vitess Authors. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -45,35 +72,64 @@ distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.` +) type ( + // generatorSPI provides services to individual generators during code generation. + // It acts as a service provider interface, giving generators access to type discovery + // and scope information needed for generating helper methods. generatorSPI interface { + // addType adds a newly discovered type to the processing queue addType(t types.Type) + // scope returns the type scope for finding implementations scope() *types.Scope + // findImplementations finds all types that implement the given interface findImplementations(iff *types.Interface, impl func(types.Type) error) error - iface() *types.Interface // the root interface that all nodes are expected to implement + // iface returns the root interface that all nodes are expected to implement + iface() *types.Interface } + + // generator defines the interface that all specialized generators must implement. + // Each generator handles specific types of Go constructs (structs, interfaces, etc.) + // and produces the appropriate helper methods for those types. generator interface { + // genFile generates the final output file for this generator genFile(generatorSPI) (string, *jen.File) + // interfaceMethod handles interface types with type switching logic interfaceMethod(t types.Type, iface *types.Interface, spi generatorSPI) error + // structMethod handles struct types with field iteration structMethod(t types.Type, strct *types.Struct, spi generatorSPI) error + // ptrToStructMethod handles pointer-to-struct types ptrToStructMethod(t types.Type, strct *types.Struct, spi generatorSPI) error + // ptrToBasicMethod handles pointer-to-basic types (e.g., *int, *string) ptrToBasicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error + // sliceMethod handles slice types with element processing sliceMethod(t types.Type, slice *types.Slice, spi generatorSPI) error + // basicMethod handles basic types (int, string, bool, etc.) basicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error } - // astHelperGen finds implementations of the given interface, - // and uses the supplied `generator`s to produce the output code + + // astHelperGen is the main orchestrator that coordinates the code generation process. + // It discovers implementations of a root interface and uses multiple specialized + // generators to produce helper methods for all discovered types. astHelperGen struct { + // DebugTypes enables debug output for type processing DebugTypes bool - mod *packages.Module - sizes types.Sizes + // mod is the Go module information for path resolution + mod *packages.Module + // sizes provides platform-specific type size information + sizes types.Sizes + // namedIface is the root interface type for which helpers are generated namedIface *types.Named - _iface *types.Interface - gens []generator + // _iface is the underlying interface type + _iface *types.Interface + // gens is the list of specialized generators (clone, equals, visit, etc.) + gens []generator + // _scope is the type scope for finding implementations _scope *types.Scope - todo []types.Type + // todo is the queue of types that need to be processed + todo []types.Type } ) @@ -81,6 +137,15 @@ func (gen *astHelperGen) iface() *types.Interface { return gen._iface } +// newGenerator creates a new AST helper generator with the specified configuration. +// +// Parameters: +// - mod: Go module information for path resolution +// - sizes: Platform-specific type size information +// - named: The root interface type for which helpers will be generated +// - generators: Specialized generators for different helper types (clone, equals, etc.) +// +// Returns a configured astHelperGen ready to process types and generate code. func newGenerator(mod *packages.Module, sizes types.Sizes, named *types.Named, generators ...generator) *astHelperGen { return &astHelperGen{ DebugTypes: true, @@ -107,7 +172,7 @@ func findImplementations(scope *types.Scope, iff *types.Interface, impl func(typ case *types.Interface: // This is OK; interfaces are references default: - panic(fmt.Errorf("interface %s implemented by %s (%s as %T) without ptr", iff.String(), baseType, tt.String(), tt)) + return fmt.Errorf("interface %s implemented by %s (%s as %T) without ptr", iff.String(), baseType, tt.String(), tt) } } if types.TypeString(baseType, noQualifier) == visitableName { @@ -140,7 +205,10 @@ func (gen *astHelperGen) GenerateCode() (map[string]*jen.File, error) { gen._scope = pkg.Scope() gen.todo = append(gen.todo, gen.namedIface) - jenFiles := gen.createFiles() + jenFiles, err := gen.createFiles() + if err != nil { + return nil, err + } result := map[string]*jen.File{} for fName, genFile := range jenFiles { @@ -175,16 +243,34 @@ func VerifyFilesOnDisk(result map[string]*jen.File) (errors []error) { return errors } +// Options configures the AST helper generation process. type Options struct { - Packages []string + // Packages specifies the Go packages to analyze for AST types. + // Can be package paths like "./mypackage" or import paths like "github.com/example/ast". + Packages []string + + // RootInterface is the fully qualified name of the root interface that all AST nodes implement. + // Format: "package.path.InterfaceName" (e.g., "github.com/example/ast.Node") RootInterface string - Clone CloneOptions + // Clone configures the clone generator options + Clone CloneOptions + + // Equals configures the equality comparison generator options Equals EqualsOptions } -// GenerateASTHelpers loads the input code, constructs the necessary generators, -// and generates the rewriter and clone methods for the AST +// GenerateASTHelpers is the main entry point for generating AST helper methods. +// +// It loads the specified packages, analyzes the types that implement the root interface, +// and generates comprehensive helper methods including clone, equals, visit, rewrite, +// path enumeration, and copy-on-write functionality. +// +// The function returns a map where keys are file paths and values are the generated +// Go source files. The caller is responsible for writing these files to disk. +// +// Returns an error if package loading fails, the root interface cannot be found, +// or code generation encounters any issues. func GenerateASTHelpers(options *Options) (map[string]*jen.File, error) { loaded, err := packages.Load(&packages.Config{ Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesSizes | packages.NeedTypesInfo | packages.NeedDeps | packages.NeedImports | packages.NeedModule, @@ -260,48 +346,76 @@ func (gen *astHelperGen) addType(t types.Type) { gen.todo = append(gen.todo, t) } -func (gen *astHelperGen) createFiles() map[string]*jen.File { +func (gen *astHelperGen) createFiles() (map[string]*jen.File, error) { + if err := gen.processTypeQueue(); err != nil { + return nil, err + } + return gen.generateOutputFiles(), nil +} + +// processTypeQueue processes all types in the todo queue with all generators +func (gen *astHelperGen) processTypeQueue() error { alreadyDone := map[string]bool{} for len(gen.todo) > 0 { t := gen.todo[0] - underlying := t.Underlying() typeName := printableTypeName(t) gen.todo = gen.todo[1:] if alreadyDone[typeName] { continue } - var err error - for _, g := range gen.gens { - switch underlying := underlying.(type) { - case *types.Interface: - err = g.interfaceMethod(t, underlying, gen) - case *types.Slice: - err = g.sliceMethod(t, underlying, gen) - case *types.Struct: - err = g.structMethod(t, underlying, gen) - case *types.Pointer: - ptrToType := underlying.Elem().Underlying() - switch ptrToType := ptrToType.(type) { - case *types.Struct: - err = g.ptrToStructMethod(t, ptrToType, gen) - case *types.Basic: - err = g.ptrToBasicMethod(t, ptrToType, gen) - default: - panic(fmt.Sprintf("%T", ptrToType)) - } - case *types.Basic: - err = g.basicMethod(t, underlying, gen) - default: - log.Fatalf("don't know how to handle %s %T", typeName, underlying) - } - if err != nil { - log.Fatal(err) - } + + if err := gen.processTypeWithGenerators(t); err != nil { + return fmt.Errorf("failed to process type %s: %w", typeName, err) } alreadyDone[typeName] = true } + return nil +} + +// processTypeWithGenerators dispatches a type to all generators based on its underlying type +func (gen *astHelperGen) processTypeWithGenerators(t types.Type) error { + underlying := t.Underlying() + typeName := printableTypeName(t) + + for _, g := range gen.gens { + var err error + switch underlying := underlying.(type) { + case *types.Interface: + err = g.interfaceMethod(t, underlying, gen) + case *types.Slice: + err = g.sliceMethod(t, underlying, gen) + case *types.Struct: + err = g.structMethod(t, underlying, gen) + case *types.Pointer: + err = gen.handlePointerType(t, underlying, g) + case *types.Basic: + err = g.basicMethod(t, underlying, gen) + default: + return fmt.Errorf("don't know how to handle type %s %T", typeName, underlying) + } + if err != nil { + return fmt.Errorf("generator failed for type %s: %w", typeName, err) + } + } + return nil +} +// handlePointerType handles pointer types by dispatching to the appropriate method +func (gen *astHelperGen) handlePointerType(t types.Type, ptr *types.Pointer, g generator) error { + ptrToType := ptr.Elem().Underlying() + switch ptrToType := ptrToType.(type) { + case *types.Struct: + return g.ptrToStructMethod(t, ptrToType, gen) + case *types.Basic: + return g.ptrToBasicMethod(t, ptrToType, gen) + default: + return fmt.Errorf("unsupported pointer type %T", ptrToType) + } +} + +// generateOutputFiles collects the generated files from all generators +func (gen *astHelperGen) generateOutputFiles() map[string]*jen.File { result := map[string]*jen.File{} for _, g := range gen.gens { fName, jenFile := g.genFile(gen) @@ -310,6 +424,9 @@ func (gen *astHelperGen) createFiles() map[string]*jen.File { return result } +// noQualifier is used to print types without package qualifiers +var noQualifier = func(*types.Package) string { return "" } + // printableTypeName returns a string that can be used as a valid golang identifier func printableTypeName(t types.Type) string { switch t := t.(type) { diff --git a/go/tools/asthelpergen/clone_gen.go b/go/tools/asthelpergen/clone_gen.go index eed3ff685ee..28a2607cee4 100644 --- a/go/tools/asthelpergen/clone_gen.go +++ b/go/tools/asthelpergen/clone_gen.go @@ -26,7 +26,11 @@ import ( "github.com/dave/jennifer/jen" ) +// CloneOptions configures the clone generator behavior. type CloneOptions struct { + // Exclude specifies type patterns that should not be deep cloned. + // Types matching these patterns will be returned as-is instead of being cloned. + // Patterns use glob-style matching (e.g., "*NoCloneType"). Exclude []string } @@ -68,7 +72,7 @@ func (c *cloneGen) readValueOfType(t types.Type, expr jen.Code, spi generatorSPI case *types.Basic: return expr case *types.Interface: - if types.TypeString(t, noQualifier) == "any" { + if types.TypeString(t, noQualifier) == anyTypeName { // these fields have to be taken care of manually return expr } diff --git a/go/tools/asthelpergen/design.md b/go/tools/asthelpergen/design.md new file mode 100644 index 00000000000..5199c4b5609 --- /dev/null +++ b/go/tools/asthelpergen/design.md @@ -0,0 +1,182 @@ +# AST Helper Generator Design + +## Overview + +The `asthelpergen` package is a code generation tool that automatically creates helper methods for Abstract Syntax Tree (AST) implementations. Given a root interface, it generates various AST manipulation utilities including cloning, equality comparison, visiting, and rewriting functions. + +## Architecture + +### Core Components + +#### 1. Main Generator (`astHelperGen`) +- **Purpose**: Central orchestrator that coordinates the code generation process +- **Key Fields**: + - `namedIface`: The root interface for which helpers are generated + - `gens`: Collection of specialized generators (clone, equals, visit, rewrite, etc.) + - `todo`: Queue of types discovered during generation that need processing + - `_scope`: Type scope for finding implementations + +#### 2. Generator Interface System +The system uses a plugin-like architecture with two key interfaces: + +**`generatorSPI`** (Service Provider Interface): +- Provides services to individual generators +- `addType()`: Adds newly discovered types to the processing queue +- `findImplementations()`: Discovers all types implementing a given interface +- `scope()`: Access to the type system scope + +**`generator`** Interface: +- Implemented by each specialized generator +- Handles different type categories (struct, pointer, slice, basic, interface) +- Generates the final output file + +### Specialized Generators + +#### 1. Clone Generator (`cloneGen`) +- **Purpose**: Creates deep clone methods for AST nodes +- **Output**: `ast_clone.go` +- **Key Features**: + - Recursive cloning with type dispatch + - Exclusion list support for types that shouldn't be cloned + - Handles pointers, slices, structs, and interfaces + +#### 2. Equals Generator (`equalsGen`) +- **Purpose**: Creates deep equality comparison methods +- **Output**: `ast_equals.go` +- **Key Features**: + - Custom comparator support via `Comparator` struct + - Field-by-field comparison for structs + - Type-safe equality checks with nil handling + +#### 3. Visit Generator (`visitGen`) +- **Purpose**: Creates visitor pattern methods for AST traversal +- **Output**: `ast_visit.go` +- **Key Features**: + - Pre-order traversal of AST nodes + - Error propagation during traversal + - Support for `Visitable` interface with custom visit logic + +#### 4. Rewrite Generator (`rewriteGen`) +- **Purpose**: Creates AST rewriting/transformation methods +- **Output**: `ast_rewrite.go` +- **Key Features**: + - Pre and post-order hooks for transformations + - Path collection for tracking node locations + - Safe node replacement with type checking + +#### 5. Path Generator (`pathGen`) +- **Purpose**: Creates path enumeration for AST nodes +- **Output**: `ast_path.go` +- **Key Features**: + - Generates constants for field paths + - Supports path-based navigation + +#### 6. Copy-on-Write Generator (`cowGen`) +- **Purpose**: Creates copy-on-write functionality +- **Output**: `ast_copy_on_rewrite.go` +- **Key Features**: + - Efficient AST modification with minimal copying + - Shared immutable structures where possible + +## Code Generation Process + +### 1. Discovery Phase +1. Load packages using `go/packages` +2. Find the root interface type +3. Initialize generators with configuration options +4. Create the main `astHelperGen` instance + +### 2. Type Analysis Phase +1. Start with the root interface in the `todo` queue +2. For each type in the queue: + - Determine its underlying type (struct, pointer, slice, etc.) + - Call appropriate method on each generator + - Generators may add new types to the queue +3. Continue until all discovered types are processed + +### 3. Code Generation Phase +1. Each generator creates its output file +2. Files are formatted and optimized +3. License headers and generation comments are added + +### 4. Output and Verification +1. Generated files are returned as a map of file paths to content +2. Optional verification against existing files on disk +3. Files can be written to disk or verified for CI/CD + +## Type System Integration + +### Type Discovery +The generator uses Go's type system to: +- Find all types implementing the root interface +- Discover field types that need helper methods +- Handle complex type relationships (pointers, slices, nested structs) + +### Type Dispatch +Each generator implements methods for different type categories: +- `interfaceMethod`: Handles interface types with type switching +- `structMethod`: Handles struct types with field iteration +- `ptrToStructMethod`: Handles pointer-to-struct types +- `sliceMethod`: Handles slice types with element processing +- `basicMethod`: Handles basic types (int, string, etc.) + +## Configuration and Options + +### Clone Options +- `Exclude`: List of type patterns to exclude from cloning + +### Equals Options +- `AllowCustom`: List of types that can have custom comparators + +### Main Options +- `Packages`: Go packages to analyze +- `RootInterface`: Fully qualified name of the root interface + +## Integration Points + +### Command Line Interface (`main/main.go`) +- Uses `spf13/pflag` for command-line argument parsing +- Supports verification mode for CI/CD pipelines +- Integrates with Vitess codegen utilities + +### Testing Integration +- Integration tests verify generated code matches expectations +- Test helpers in `integration/` package provide real-world examples +- Comprehensive type coverage including edge cases + +## Key Design Patterns + +### 1. Plugin Architecture +Multiple generators implement the same interface, allowing easy extension with new generator types. + +### 2. Type-Driven Generation +The generation process is driven by Go's type system, ensuring type safety and correctness. + +### 3. Incremental Discovery +Types are discovered incrementally as generators process fields and references, ensuring complete coverage. + +### 4. Template-Free Generation +Uses `dave/jennifer` for programmatic Go code generation instead of text templates, providing better type safety and refactoring support. + +## Performance Considerations + +- **Lazy Type Discovery**: Types are only processed when referenced +- **Deduplication**: Each type is processed only once using the `alreadyDone` map +- **Efficient Code Generation**: Jennifer library provides efficient Go code generation +- **Minimal Dependencies**: Generated code has minimal runtime dependencies + +## Error Handling + +- **Package Loading Errors**: Detailed error reporting for missing packages +- **Type Resolution Errors**: Clear messages when interfaces/types cannot be found +- **Generation Errors**: Panic on unexpected type scenarios with helpful context +- **Verification Errors**: File comparison errors for CI/CD integration + +## Future Extensibility + +The architecture supports easy addition of new generators: +1. Implement the `generator` interface +2. Add to the generator list in `GenerateASTHelpers()` +3. Follow established patterns for type dispatch and code generation + +This design provides a robust, extensible foundation for AST helper code generation while maintaining type safety and performance. \ No newline at end of file diff --git a/go/tools/asthelpergen/equals_gen.go b/go/tools/asthelpergen/equals_gen.go index 2e34eb5a36c..d3c6f473472 100644 --- a/go/tools/asthelpergen/equals_gen.go +++ b/go/tools/asthelpergen/equals_gen.go @@ -26,7 +26,11 @@ import ( const Comparator = "Comparator" +// EqualsOptions configures the equals generator behavior. type EqualsOptions struct { + // AllowCustom specifies types that can have custom equality comparators. + // For these types, the generated Comparator struct will include function fields + // that allow custom comparison logic to be injected at runtime. AllowCustom []string } @@ -167,7 +171,7 @@ func compareAllStructFields(strct *types.Struct, spi generatorSPI) jen.Code { var others []*jen.Statement for i := 0; i < strct.NumFields(); i++ { field := strct.Field(i) - if field.Type().Underlying().String() == "any" || strings.HasPrefix(field.Name(), "_") { + if field.Type().Underlying().String() == anyTypeName || strings.HasPrefix(field.Name(), "_") { // we can safely ignore this, we do not want ast to contain `any` types. continue } diff --git a/go/tools/asthelpergen/rewrite_gen.go b/go/tools/asthelpergen/rewrite_gen.go index 6fa5a296103..4478fb90157 100644 --- a/go/tools/asthelpergen/rewrite_gen.go +++ b/go/tools/asthelpergen/rewrite_gen.go @@ -24,8 +24,7 @@ import ( ) const ( - rewriteName = "rewrite" - visitableName = "Visitable" + rewriteName = "rewrite" ) type rewriteGen struct { @@ -550,7 +549,3 @@ func returnTrue() jen.Code { func returnFalse() jen.Code { return jen.Return(jen.False()) } - -var noQualifier = func(p *types.Package) string { - return "" -}