main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strings"
	"unicode"
)

type field struct {
	name string
	typ  string
	zero string
}

type builder struct {
	fields map[string]field
	fns    []string
}

type namedImport struct {
	name string
	path string
}

func main() {
	flag.Parse()
	args := flag.Args()
	filename := args[0]
	inTypes := args[1]
	types := strings.Split(inTypes, ",")
	typeSet := make(map[string]struct{})
	for _, typ := range types {
		typeSet[typ] = struct{}{}
	}

	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, filename, nil, 0)
	if err != nil {
		log.Fatal(err)
	}

	imports := make(map[string]string)
	for _, impt := range file.Imports {
		name := packageName(impt)
		imports[name] = impt.Path.Value
	}

	builders := make(map[string]builder)
	builderImports := make([]namedImport, 0)

	ast.Inspect(file, func(n ast.Node) bool {
		ts, ok := n.(*ast.TypeSpec)
		if !ok {
			return true
		}

		st, ok := ts.Type.(*ast.StructType)
		if !ok {
			return true
		}

		if len(st.Fields.List) == 0 {
			return true
		}

		// TODO only generate for provided types
		if _, ok := typeSet[ts.Name.Name]; !ok {
			return true
		}

		bldr := builder{fields: make(map[string]field)}

		for _, fld := range st.Fields.List {
			typ := getType(fld)
			name := getName(fld, typ)
			lowerName := toLower(name)
			zero := getZero(fld)

			parts := strings.Split(typ, ".")
			if len(parts) > 1 {
				pkgName := parts[0]
				builderImports = append(builderImports, namedImport{name: pkgName, path: imports[pkgName]})
			}

			bldr.fields[lowerName] = field{name: name, typ: typ, zero: zero}
			fn := fmt.Sprintf("func (b *%sBuilder) %s(%s %s) *%sBuilder {\n", ts.Name.Name, name, lowerName, typ, ts.Name)
			fn = fn + fmt.Sprintf("\tb.%s = %s\n", lowerName, lowerName)
			fn = fn + "\treturn b\n"
			fn = fn + "}\n"

			bldr.fns = append(bldr.fns, fn)
		}

		newFn := fmt.Sprintf("func New%sBuilder() *%sBuilder {\n", ts.Name.Name, ts.Name.Name)
		newFn = newFn + fmt.Sprintf("\treturn &%sBuilder{", ts.Name.Name)

		buildFn := fmt.Sprintf("func (b *%sBuilder) Build() %s {\n", ts.Name.Name, ts.Name.Name)
		buildFn = buildFn + fmt.Sprintf("\treturn %s{\n", ts.Name.Name)

		hasZero := false
		for fieldName, fld := range bldr.fields {
			if fld.zero != "" {
				newFn = newFn + fmt.Sprintf("\n\t\t%s: %s,", fieldName, fld.zero)
				hasZero = true
			}
			buildFn = buildFn + fmt.Sprintf("\t\t%s: b.%s,\n", fld.name, fieldName)
		}

		if hasZero {
			newFn = newFn + "\n\t}\n"
		} else {
			newFn = newFn + "}\n"
		}

		newFn = newFn + "}\n"
		buildFn = buildFn + "\t}\n"
		buildFn = buildFn + "}\n"

		fmt.Println(newFn)
		bldr.fns = append(bldr.fns, newFn, buildFn)
		builders[ts.Name.String()] = bldr

		return false
	})

	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "// Code generated by \"buildify %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
	fmt.Fprintln(buf)
	fmt.Fprintln(buf, "package main") // TODO package name from output directory name
	fmt.Fprintln(buf)

	if len(builderImports) > 0 {
		fmt.Fprintln(buf, "import (")
		for _, namedImport := range builderImports {
			if namedImport.name != "" {
				fmt.Fprintf(buf, "\t%s %s\n", namedImport.name, namedImport.path)
			} else {
				fmt.Fprintf(buf, "\t%s\n", namedImport.path)
			}
		}
		fmt.Fprintln(buf, ")")
		fmt.Fprintln(buf)
	}

	for typ, b := range builders {
		fmt.Fprintf(buf, "type %sBuilder struct {\n", typ)
		for fieldName, fld := range b.fields {
			fmt.Fprintf(buf, "\t%s  %s\n", fieldName, fld.typ)
		}
		fmt.Fprintln(buf, "}")
		fmt.Fprintln(buf)

		for _, fn := range b.fns {
			fmt.Fprintln(buf, fn)
			fmt.Fprintln(buf)
		}

		fmt.Fprintln(buf)
	}

	src, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(err)
	}

	err = os.WriteFile("buildify.out.go", src, 0644) // TODO output file from flags/args
	if err != nil {
		log.Fatal(err)
	}
}

func getType(field *ast.Field) string {
	switch typ := field.Type.(type) {
	case *ast.Ident:
		return typ.Name
	case *ast.SelectorExpr:
		pkg, ok := typ.X.(*ast.Ident)
		if !ok {
			log.Fatalf("typ.X is a %T", typ.X)
			return ""
		}

		return fmt.Sprintf("%s.%s", pkg.Name, typ.Sel.Name)
	default:
		log.Fatalf("unknown field type %T", field.Type)
		return ""
	}
}

func getName(field *ast.Field, typ string) string {
	if field.Names == nil {
		parts := strings.Split(typ, ".")
		if len(parts) == 1 {
			return parts[0]
		}

		return parts[1]
	}

	return field.Names[0].Name
}

func getZero(field *ast.Field) string {
	if field.Tag == nil {
		return ""
	}
	rawTag := strings.Trim(field.Tag.Value, "`")
	tags := strings.Split(rawTag, " ")
	for _, tag := range tags {
		parts := strings.Split(tag, ":")
		tagName, value := parts[0], parts[1]
		if tagName == "zero" {
			return value
		}
	}

	return ""
}

func packageName(impt *ast.ImportSpec) string {
	if impt.Name != nil {
		return impt.Name.Name
	}

	parts := strings.Split(impt.Path.Value, "/")
	return strings.Trim(parts[len(parts)-1], "\"")
}

func toLower(s string) string {
	for _, r := range s {
		return string(unicode.ToLower(r)) + s[1:]
	}

	return s
}