mukan-ignite/ignite/pkg/xast/global.go

package xast

import (
	"bytes"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"

	"github.com/ignite/cli/v29/ignite/pkg/errors"
)

type (
	// globalOpts represent the options for globals.
	globalOpts struct {
		globals []global
	}

	// GlobalOptions configures code generation.
	GlobalOptions func(*globalOpts)

	global struct {
		name, varType, value string
	}

	// GlobalType represents the global type.
	GlobalType string
)

const (
	GlobalTypeVar   GlobalType = "var"
	GlobalTypeConst GlobalType = "const"
)

// WithGlobal add a new global.
func WithGlobal(name, varType, value string) GlobalOptions {
	return func(c *globalOpts) {
		c.globals = append(c.globals, global{
			name:    name,
			varType: varType,
			value:   value,
		})
	}
}

func newGlobalOptions() globalOpts {
	return globalOpts{
		globals: make([]global, 0),
	}
}

// InsertGlobal inserts global variables or constants into the provided Go source code content after the import section.
// The function parses the provided content, locates the import section, and inserts the global declarations immediately after it.
// The type of globals (variables or constants) is specified by the globalType parameter.
// Each global declaration is defined by calling WithGlobal function with appropriate arguments.
// The function returns the modified content with the inserted global declarations.
func InsertGlobal(fileContent string, globalType GlobalType, globals ...GlobalOptions) (modifiedContent string, err error) {
	// apply global options.
	opts := newGlobalOptions()
	for _, o := range globals {
		o(&opts)
	}
	if len(opts.globals) == 0 {
		return fileContent, nil
	}

	tok, err := globalTypeToken(globalType)
	if err != nil {
		return "", err
	}

	fileSet := token.NewFileSet()

	// Parse the Go source code content.
	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	// Find the index of the import declaration or package declaration if no imports.
	var insertIndex int
	for i, decl := range f.Decls {
		if genDecl, ok := decl.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
			insertIndex = i + 1
			break
		} else if funcDecl, ok := decl.(*ast.FuncDecl); ok {
			insertIndex = i
			if funcDecl.Doc == nil {
				insertIndex++
			}
			break
		}
	}

	for _, global := range opts.globals {
		spec, err := newGlobalValueSpec(fileSet, global)
		if err != nil {
			return "", err
		}

		f.Decls = append(
			f.Decls[:insertIndex],
			append([]ast.Decl{
				&ast.GenDecl{
					TokPos: 1,
					Tok:    tok,
					Specs:  []ast.Spec{spec},
				},
			}, f.Decls[insertIndex:]...)...,
		)
		insertIndex++
	}

	f.Comments = cmap.Filter(f).Comments()

	// Format the modified AST.
	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	// Return the modified content.
	return buf.String(), nil
}

// AppendFunction appends a new function to the end of the Go source code content.
func AppendFunction(fileContent string, function string) (modifiedContent string, err error) {
	fileSet := token.NewFileSet()

	// Parse the function body as a separate file.
	funcFile, err := parser.ParseFile(fileSet, "", "package main\n"+function, parser.AllErrors)
	if err != nil {
		return "", err
	}

	// Extract the first declaration, assuming it's a function declaration.
	var funcDecl *ast.FuncDecl
	for _, decl := range funcFile.Decls {
		if fDecl, ok := decl.(*ast.FuncDecl); ok {
			funcDecl = fDecl
			break
		}
	}
	if funcDecl == nil {
		return "", errors.Errorf("no function declaration found in the provided function body")
	}

	// Parse the Go source code content.
	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	// Append the function declaration to the file's declarations.
	f.Decls = append(f.Decls, funcDecl)

	f.Comments = cmap.Filter(f).Comments()

	// Format the modified AST.
	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	return buf.String(), nil
}

type (
	// structOpts represent the options for structs.
	structOpts struct {
		values []structValue
	}

	// StructOpts configures struct changes.
	StructOpts func(*structOpts)

	structValue struct {
		value     string
		valueType string
	}
)

// AppendStructValue add a new value inside struct. For instances,
// the struct have only one field 'test struct{ test1 string }' and we want to add
// the `test2 int` the result will be 'test struct{ test1 string, test int }'.
func AppendStructValue(value, valueType string) StructOpts {
	return func(c *structOpts) {
		c.values = append(c.values, structValue{
			value:     value,
			valueType: valueType,
		})
	}
}

func newStructOptions() structOpts {
	return structOpts{
		values: make([]structValue, 0),
	}
}

// ModifyStruct modifies a struct in the provided Go source code.
func ModifyStruct(fileContent, structName string, options ...StructOpts) (string, error) {
	// Apply struct options.
	opts := newStructOptions()
	for _, o := range options {
		o(&opts)
	}

	fileSet := token.NewFileSet()

	// Parse the Go source code content.
	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	// Locate and modify the struct declaration.
	var (
		found      bool
		structType *ast.StructType
	)
	for _, decl := range f.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok || genDecl.Tok != token.TYPE {
			continue
		}

		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok || typeSpec.Name.Name != structName {
				continue
			}

			structType, ok = typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}

			found = true
			break
		}
		if found {
			break
		}
	}
	if !found {
		return "", errors.Errorf("struct %q not found in file content", structName)
	}
	for _, v := range opts.values {
		structType.Fields.List = append(structType.Fields.List, &ast.Field{
			Names: []*ast.Ident{ast.NewIdent(v.value)},
			Type:  ast.NewIdent(v.valueType),
		})
	}

	f.Comments = cmap.Filter(f).Comments()

	// Format the modified AST.
	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	// Return the modified content.
	return buf.String(), nil
}

type (
	// globalArrayOpts represent the options for globar array variables.
	globalArrayOpts struct {
		values []string
	}

	// GlobalArrayOpts configures global array variable changes.
	GlobalArrayOpts func(*globalArrayOpts)
)

// AppendGlobalArrayValue add a new value inside a global array variable. For instances,
// the variable have only one field '[]]' and we want to add
// the `test2 int` the result will be 'test struct{ test1 string, test int }'.
func AppendGlobalArrayValue(value string) GlobalArrayOpts {
	return func(c *globalArrayOpts) {
		c.values = append(c.values, value)
	}
}

func newGlobalArrayOptions() globalArrayOpts {
	return globalArrayOpts{
		values: make([]string, 0),
	}
}

// ModifyGlobalArrayVar modifies an array global array variable in the provided Go source code by appending new values.
func ModifyGlobalArrayVar(fileContent, globalName string, options ...GlobalArrayOpts) (string, error) {
	opts := newGlobalArrayOptions()
	for _, o := range options {
		o(&opts)
	}
	if len(opts.values) == 0 {
		return fileContent, nil
	}

	fileSet := token.NewFileSet()

	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	var (
		found   bool
		compLit *ast.CompositeLit
	)
	for _, decl := range f.Decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok || genDecl.Tok != token.VAR {
			continue
		}

		for _, spec := range genDecl.Specs {
			valueSpec, ok := spec.(*ast.ValueSpec)
			if !ok || len(valueSpec.Names) == 0 || valueSpec.Names[0].Name != globalName || len(valueSpec.Values) == 0 {
				continue
			}

			compLit, ok = valueSpec.Values[0].(*ast.CompositeLit)
			if !ok {
				continue
			}

			found = true
			break
		}
		if found {
			break
		}
	}

	if !found {
		return "", errors.Errorf("global array %q not found in file content", globalName)
	}
	appendCompositeLiteralValues(fileSet, compLit, opts.values)

	f.Comments = cmap.Filter(f).Comments()

	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	return buf.String(), nil
}

func globalTypeToken(globalType GlobalType) (token.Token, error) {
	switch globalType {
	case GlobalTypeVar:
		return token.VAR, nil
	case GlobalTypeConst:
		return token.CONST, nil
	default:
		return token.ILLEGAL, errors.Errorf("unsupported global type: %s", string(globalType))
	}
}

func newGlobalValueSpec(fileSet *token.FileSet, global global) (*ast.ValueSpec, error) {
	spec := &ast.ValueSpec{
		Names: []*ast.Ident{ast.NewIdent(global.name)},
	}
	if global.varType != "" {
		spec.Type = ast.NewIdent(global.varType)
	}
	if global.value == "" {
		return spec, nil
	}

	valueExpr, err := parser.ParseExprFrom(fileSet, "", []byte(global.value), parser.ParseComments)
	if err != nil {
		return nil, err
	}
	spec.Values = []ast.Expr{valueExpr}
	return spec, nil
}

func appendCompositeLiteralValues(fileSet *token.FileSet, compLit *ast.CompositeLit, values []string) {
	file := fileSet.File(compLit.Pos())
	maxOffset := file.Offset(compLit.Rbrace)
	for _, elt := range compLit.Elts {
		if pos := elt.End(); pos.IsValid() {
			if offset := file.Offset(pos); offset > maxOffset {
				maxOffset = offset
			}
		}
	}

	for i, valueName := range values {
		insertPos := file.Pos(maxOffset + i)
		value := ast.NewIdent(valueName)
		value.NamePos = insertPos

		compLit.Elts = append(compLit.Elts, value)
		compLit.Rbrace += token.Pos(i + 1)
	}

	if len(compLit.Elts) == 0 {
		return
	}

	last := compLit.Elts[len(compLit.Elts)-1]
	if file.Line(compLit.Rbrace) == file.Line(last.End())-1 {
		file.AddLine(file.Offset(compLit.Rbrace))
		compLit.Rbrace += token.Pos(1)
	}
}