package main import ( "bufio" "bytes" _ "embed" "fmt" "go/format" "io" "maps" "os" "reflect" "slices" "strconv" "strings" "text/template" "github.com/Masterminds/sprig/v3" "github.com/evcc-io/evcc/api" combinations "github.com/mxschmitt/golang-combinations" "github.com/samber/lo" "github.com/spf13/pflag" "golang.org/x/tools/imports" ) //go:generate go tool decorate //evcc:function decorateTest //evcc:basetype api.Charger //evcc:types api.MeterEnergy,api.PhaseSwitcher,api.PhaseGetter //go:embed decorate.tpl var srcTmpl string //go:embed header.tpl var header string type funcStruct struct { Signature, Function, VarName, ReturnTypes string Params []string } type typeStruct struct { Type, ShortType string Functions []funcStruct } var interfaces = make(map[string]reflect.Type) var dependents = make(map[string][]string) func init() { reflectTypes := map[reflect.Type][]reflect.Type{ reflect.TypeFor[api.Meter](): {reflect.TypeFor[api.MeterEnergy](), reflect.TypeFor[api.PhaseCurrents](), reflect.TypeFor[api.PhaseVoltages](), reflect.TypeFor[api.MaxACPowerGetter]()}, reflect.TypeFor[api.PhaseCurrents](): {reflect.TypeFor[api.PhasePowers]()}, // phase powers are only used to determine currents sign reflect.TypeFor[api.PhaseSwitcher](): {reflect.TypeFor[api.PhaseGetter]()}, reflect.TypeFor[api.Battery](): {reflect.TypeFor[api.BatteryCapacity](), reflect.TypeFor[api.SocLimiter](), reflect.TypeFor[api.BatteryController](), reflect.TypeFor[api.BatterySocLimiter](), reflect.TypeFor[api.BatteryPowerLimiter]()}, reflect.TypeFor[api.ChargeState](): {reflect.TypeFor[api.ChargeController](), reflect.TypeFor[api.CurrentController]()}, reflect.TypeFor[api.CurrentController](): {reflect.TypeFor[api.CurrentGetter]()}, } for typ, types := range reflectTypes { interfaces[typ.String()] = typ for _, t := range types { interfaces[t.String()] = t } dependents[typ.String()] = lo.Map(types, func(typ reflect.Type, _ int) string { return typ.String() }) } for _, typ := range []reflect.Type{ reflect.TypeFor[api.Curtailer](), reflect.TypeFor[api.Resurrector](), reflect.TypeFor[api.VehicleOdometer](), reflect.TypeFor[api.VehicleRange](), reflect.TypeFor[api.VehicleClimater](), reflect.TypeFor[api.VehicleFinishTimer](), reflect.TypeFor[api.Identifier](), reflect.TypeFor[api.ChargerEx](), reflect.TypeFor[api.ChargeRater](), reflect.TypeFor[api.StatusReasoner](), } { interfaces[typ.String()] = typ } } // hasIntersection returns if the slices intersect func hasIntersection[T comparable](a, b []T) bool { for _, el := range a { if slices.Contains(b, el) { return true } } return false } func getCombinations(combos []string) [][]string { validCombos := make([][]string, 0) sortedDependents := slices.Sorted(maps.Keys(dependents)) COMBO: for _, c := range combinations.All(combos) { // order the cases for generation for _, master := range sortedDependents { details := dependents[master] // prune combinations where ... // - master is part of the decorators // - master is not part of the currently evaluated combination // - details are part of the currently evaluated combination // ... and remove details from the combination if slices.Contains(combos, master) && !slices.Contains(c, master) && hasIntersection(c, details) { c = lo.Without(c, details...) if len(c) == 0 { continue COMBO } } } // prune duplicates for _, v := range validCombos { if slices.Equal(v, c) { continue COMBO } } validCombos = append(validCombos, c) } return validCombos } func getTemplate(dtypes []reflect.Type, types map[string]typeStruct, combos []string) *template.Template { tmpl, err := template.New("gen").Funcs(sprig.FuncMap()).Funcs(template.FuncMap{ // contains checks if slice contains string "contains": slices.Contains[[]string, string], // ordered returns a slice of funcStruct ordered by dynamicType "ordered": func() []funcStruct { ordered := make([]funcStruct, 0) for _, t := range dtypes { for _, f := range types[getTypeImport(t)].Functions { ordered = append(ordered, f) } } return ordered }, "requiredType": func(c []string, typ string) bool { for master, details := range dependents { // exclude combinations where ... // - master is part of the decorators // - master is not part of the currently evaluated combination // - details are part of the currently evaluated combination if slices.Contains(combos, master) && !slices.Contains(c, master) && slices.Contains(details, typ) { return false } } return true }, "empty": func() []string { return nil }, }).Parse(srcTmpl) if err != nil { fmt.Printf("invalid template: %s", err) os.Exit(2) } return tmpl } func getTypeImport(t reflect.Type) string { n := t.Name() if p := t.PkgPath(); p != "" { if s := strings.Split(p, "github.com/evcc-io/evcc/"); len(s) == 2 { return fmt.Sprintf("%s.%s", s[1], n) } else { return fmt.Sprintf("%s.%s", p, n) } } return n } func generate(out io.Writer, functionName, baseType string, dtypes []reflect.Type) error { var combos []string types := make(map[string]typeStruct) for _, t := range dtypes { lastPart := t.Name() var funcs []funcStruct for i := 0; i < t.NumMethod(); i++ { m := t.Method(i) varName := strings.ToLower(lastPart[:1]) + lastPart[1:] if t.NumMethod() > 1 { varName += strconv.Itoa(i) } var params []string for input := range m.Type.Ins() { params = append(params, getTypeImport(input)) } var returns []string for output := range m.Type.Outs() { returns = append(returns, getTypeImport(output)) } funcs = append(funcs, funcStruct{ VarName: varName, Signature: fmt.Sprintf("func(%s) (%s)", strings.Join(params, ", "), strings.Join(returns, ", ")), Function: m.Name, Params: params, ReturnTypes: fmt.Sprintf("(%s)", strings.Join(returns, ",")), }) } types[getTypeImport(t)] = typeStruct{ Type: t.Name(), ShortType: lastPart, Functions: funcs, } combos = append(combos, getTypeImport(t)) } returnType := *ret if returnType == "" { returnType = baseType } shortBase := strings.TrimLeft(baseType, "*") if baseTypeParts := strings.SplitN(baseType, ".", 2); len(baseTypeParts) > 1 { shortBase = baseTypeParts[1] } vars := struct { Function string BaseType, ShortBase string ReturnType string Types map[string]typeStruct Combinations [][]string }{ Function: functionName, BaseType: baseType, ShortBase: shortBase, ReturnType: returnType, Types: types, Combinations: getCombinations(combos), } return getTemplate(dtypes, types, combos).Execute(out, vars) } type decorationSet struct { function, base, ret, types string } var ( target = pflag.StringP("out", "o", "", "output file") pkg = pflag.StringP("package", "p", "", "package name") funcname = pflag.StringP("function", "f", "", "function name") base = pflag.StringP("base", "b", "", "base type") ret = pflag.StringP("return", "r", "", "return type") types = pflag.StringP("type", "t", "", "comma-separated list of type definitions") ) // Usage prints flags usage func Usage() { fmt.Fprintf(os.Stderr, "Usage of decorate:\n") fmt.Fprintf(os.Stderr, "\ndecorate [flags] -type interface,interface function,function signature\n") fmt.Fprintf(os.Stderr, "\nFlags:\n") pflag.PrintDefaults() } func parseFile(file string) ([]decorationSet, error) { var res []decorationSet f, err := os.Open(file) if err != nil { return nil, err } defer f.Close() var current decorationSet scanner := bufio.NewScanner(f) for scanner.Scan() { line := scanner.Text() if s, ok := strings.CutPrefix(line, "//evcc:"); ok { segs := strings.SplitN(s, " ", 2) if len(segs) != 2 { panic("invalid segments: " + s) } switch segs[0] { case "function": // must be first if current.function != "" { res = append(res, current) current = decorationSet{} } current.function = segs[1] case "basetype": current.base = segs[1] case "returntype": current.ret = segs[1] case "types": current.types = segs[1] default: panic("invalid directive //evcc:" + segs[0]) } } } if current.function != "" { res = append(res, current) } return res, scanner.Err() } func main() { pflag.Usage = Usage pflag.Parse() // read target from go:generate gofile, ok := os.LookupEnv("GOFILE") if *target == "" && ok { gofile := strings.TrimSuffix(gofile, ".go") + "_decorators.go" target = &gofile } // read target from go:generate if gopkg, ok := os.LookupEnv("GOPACKAGE"); *pkg == "" && ok { pkg = &gopkg } sets := []decorationSet{{*funcname, *base, *ret, *types}} if *funcname == "" { all, err := parseFile(gofile) if err != nil { fmt.Println(err) os.Exit(2) } sets = all } if *pkg == "" || len(sets) == 0 || sets[0].base == "" || len(sets[0].types) == 0 { Usage() os.Exit(2) } var out io.Writer = os.Stdout var name string if target != nil { name = *target if !strings.HasSuffix(name, ".go") { name += ".go" } dst, err := os.Create(name) if err != nil { fmt.Println(err) os.Exit(2) } defer dst.Close() out = dst } generated := new(bytes.Buffer) fmt.Fprintln(generated, strings.ReplaceAll(header, "{{.Package}}", *pkg)) for _, set := range sets { var types []reflect.Type for t := range strings.SplitSeq(set.types, ",") { typ, ok := interfaces[t] if !ok { fmt.Printf("don't know interface %s\n", t) os.Exit(2) } types = append(types, typ) } var buf bytes.Buffer if err := generate(&buf, set.function, set.base, types); err != nil { fmt.Println(err) os.Exit(2) } fmt.Fprintln(generated, buf.String()) } formatted, err := format.Source(generated.Bytes()) if err != nil { formatted = generated.Bytes() } formatted, err = imports.Process(name, formatted, nil) if err != nil { fmt.Println(err) os.Exit(3) } if _, err := out.Write(formatted); err != nil { fmt.Println(err) os.Exit(2) } }