gen.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"runtime/debug"
	"strings"

	. "github.com/dave/jennifer/jen"
	"github.com/samber/lo"
)

var moduleName string

func init() {
	info, ok := debug.ReadBuildInfo()
	if ok {
		moduleName = info.Main.Path
	}
	flag.StringVar(&pkg, "pkg", moduleName, "generated file package name")
	flag.StringVar(&filename, "file", "./middle.go", "name of the file to write generated code in")
	flag.IntVar(&n, "n", 27, "maximum generated number of chains")
	flag.BoolVar(&noHeader, "no-header", false, "do not generate GENERATED header comment")
}

var alphabets = []string{
	"A", "B", "C", "D", "E", "F",
	"G", "H", "I", "J", "K", "L", "M",
	"N", "O", "P", "Q", "R", "S", "T",
	"U", "V", "W", "X", "Y", "Z",
}

func chainStructName(i int) string {
	return fmt.Sprintf("ChainHandler%d", i)
}

func factoryFuncName(i int) string {
	return fmt.Sprintf("Chain%d", i)
}

func genericTypes(i int) []Code {
	return lo.Times(i-1, func(j int) Code { return Id(alphabets[j]).Any() })
}

func genericTypeParamName(i int) string {
	return strings.ToLower(alphabets[i])
}

func parameterGenericTypes(i int) []Code {
	return lo.Times(i-1, func(j int) Code { return Id(alphabets[j]) })
}

func fnParams(i int) []Code {
	return append(
		lo.Times(i-1, func(j int) Code {
			return Id(fnName(j + 1)).
				Func().
				Params(
					append(
						[]Code{
							Qual("net/http", "ResponseWriter"),
							Add(Op("*")).Qual("net/http", "Request"),
						},
						lo.Times(j, func(j int) Code { return Id(alphabets[j]) })...,
					)...,
				).
				Parens(
					List(
						Id(alphabets[j]),
						Error(),
					),
				)
		}),
		Id(fnName(i)).
			Func().
			Params(
				append(
					[]Code{
						Qual("net/http", "ResponseWriter"),
						Add(Op("*")).Qual("net/http", "Request"),
					},
					lo.Times(i-1, func(j int) Code { return Id(alphabets[j]) })...,
				)...,
			).
			Error(),
	)
}

func fnName(n int) string {
	return fmt.Sprintf("f%d", n)
}

var (
	pkg      string
	filename string
	n        int
	noHeader bool
)

func validateFlags() error {
	if moduleName == "" {
		return fmt.Errorf("pkg option is required")
	}
	if n < 1 || n > 27 {
		return fmt.Errorf("n cannot be < 1 or > 27")
	}
	return nil
}

func main() {
	flag.Parse()
	if err := validateFlags(); nil != err {
		log.Fatalf("provided flags are invalid: %v", err)
	}
	filePathPkg := pkg
	if matches, _ := regexp.MatchString(`/v\d$`, pkg); matches {
		filePathPkg = pkg[:strings.LastIndex(pkg, "/")]
	}
	f := NewFilePath(filePathPkg)
	if !noHeader {
		f.HeaderComment(fmt.Sprintf("Code generated by %s. DO NOT EDIT.", pkg))
	}
	f.Var().Defs(
		Commentf("ErrAbort can be used to stop the middleware chain execution.").
			Line().
			Id("ErrAbort").Op("=").Qual("errors", "New").Call(Lit("chain execution stopped")),
	)
	for i := 1; i <= n; i++ {
		structName := chainStructName(i)
		f.
			Commentf("%s provides capability of processing chain functions in order by satisfying [net/http.Handler], or with an optional chain error handler via [%s.Finally] by satisfying [net/http.HandlerFunc]", structName, structName).
			Line().
			Type().
			Id(structName).
			Types(genericTypes(i)...).
			Struct(fnParams(i)...)

		f.Line()

		if i < 2 {
			f.Comment("ServeHTTP satisfies [net/http.Handler]. It executes the handler function, passing request, and response to it.")
		} else {
			f.Commentf("ServeHTTP satisfies [net/http.Handler]. It executes functions in the chain in order, passing results of all previous function calls to it. If any of the functions in the chain returns a non-nil error, the execution stops.")
		}
		f.
			Func().
			Params(Id("chain").Id(structName).Types(parameterGenericTypes(i)...)).
			Id("ServeHTTP").
			Params(
				Id("response").Qual("net/http", "ResponseWriter"),
				Id("request").Add(Op("*")).Qual("net/http", "Request"),
			).
			Block(
				append(
					lo.Flatten(
						lo.Times(i-1, func(j int) []Code {
							return []Code{
								List(
									Id(genericTypeParamName(j)),
									Err(),
								).
									Op(":=").
									Id("chain").
									Dot(fnName(j + 1)).
									Call(
										append(
											[]Code{
												Id("response"),
												Id("request"),
											},
											lo.Times(j, func(k int) Code { return Id(genericTypeParamName(k)) })...,
										)...,
									),
								If(
									Nil().
										Op("!=").
										Err(),
								).
									Block(
										Return(),
									),
							}
						}),
					),
					Id("_").
						Op("=").
						Id("chain").
						Dot(fnName(i)).
						Call(
							append(
								[]Code{
									Id("response"),
									Id("request"),
								},
								lo.Times(i-1, func(j int) Code { return Id(genericTypeParamName(j)) })...,
							)...,
						),
				)...,
			)

		f.Line()

		if i < 2 {
			f.Commentf("Finally executes middleware function registered via [%s], passing request, and response to it.", factoryFuncName(i))
		} else {
			f.Commentf("Finally executes middleware functions registered via [%s] in order, passing results of all previous function calls to it. If any of the functions in the chain returns a non-nil error, the execution stops, and executes catch with that error. If the error is [ErrAbort] according to [errors.Is] semantics, it is ignored, and catch will not be called, although the chain execution stops.", factoryFuncName(i))
		}
		f.
			Func().
			Params(Id("chain").Id(structName).Types(parameterGenericTypes(i)...)).
			Id("Finally").
			Params(
				Id("catch").
					Func().
					Params(
						Qual("net/http", "ResponseWriter"),
						Add(Op("*")).Qual("net/http", "Request"),
						Error(),
					),
			).
			Qual("net/http", "HandlerFunc").
			Block(
				Return(
					Func().
						Params(
							Id("response").Qual("net/http", "ResponseWriter"),
							Id("request").Add(Op("*")).Qual("net/http", "Request"),
						).
						Block(
							append(
								lo.Flatten(
									lo.Times(i-1, func(j int) []Code {
										return []Code{
											List(
												Id(genericTypeParamName(j)),
												Err(),
											).
												Op(":=").
												Id("chain").
												Dot(fnName(j + 1)).
												Call(
													append(
														[]Code{
															Id("response"),
															Id("request"),
														},
														lo.Times(j, func(k int) Code { return Id(genericTypeParamName(k)) })...,
													)...,
												),
											If(
												Nil().
													Op("!=").
													Err(),
											).
												Block(
													If(
														Op("!").
															Add().
															Qual("errors", "Is").
															Call(Err(), Id("ErrAbort")),
													).
														Block(
															Id("catch").Call(Id("response"), Id("request"), Err()),
														),
													Return(),
												),
										}
									}),
								),
								If(
									Err().
										Op(":=").
										Id("chain").
										Dot(fnName(i)).
										Call(
											append(
												[]Code{
													Id("response"),
													Id("request"),
												},
												lo.Times(i-1, func(j int) Code { return Id(genericTypeParamName(j)) })...,
											)...,
										),
									Nil().Op("!=").Err(),
								).
									Block(
										If(
											Op("!").
												Add().
												Qual("errors", "Is").
												Call(Err(), Id("ErrAbort")),
										).
											Block(
												Id("catch").Call(Id("response"), Id("request"), Err()),
											),
									),
							)...,
						),
				),
			)

		f.Line()

		f.Commentf("%s creates a chain of exactly %d function%s that will be executed in order.", factoryFuncName(i), i, lo.Ternary(i > 1, "s", ""))
		f.Func().
			Id(factoryFuncName(i)).
			Types(genericTypes(i)...).
			Params(fnParams(i)...).
			Id(structName).Types(parameterGenericTypes(i)...).
			Block(
				Return(
					Id(structName).
						Types(parameterGenericTypes(i)...).
						Values(lo.Times(i, func(j int) Code { return Id(fnName(j + 1)) })...),
				),
			)
	}

	var buf bytes.Buffer
	if err := f.Render(&buf); nil != err {
		log.Fatalf("failed to generate code: %v\n", err)
	}
	if err := os.WriteFile(filename, buf.Bytes(), 0644); nil != err {
		log.Fatalf("failed to write generated code to %q: %v\n", filename, err)
	}
}