/internal/cmd/units-codegen/template.go

package main

import (
	"fmt"
	"regexp"
	"strings"
	"sync"
	"text/template"
)

type pkgConfig struct {
	sync.RWMutex
	compiled bool
	PkgName  string
	Types    map[string]typeDef
}

// Compile realizes generator fields so that the struct can be used for template rendering.
func (p *pkgConfig) Compile() error {
	if p.compiled {
		return nil
	}
	wgs := make(map[string]*sync.WaitGroup)
	for k := range p.Types {
		wgs[k] = &sync.WaitGroup{}
		wgs[k].Add(1)
	}

	p.RLock()
	for k, t := range p.Types {
		go t.compileRoutine(k, p, wgs)
	}
	p.RUnlock()

	for k, wg := range wgs {
		wg.Wait()
		p.RLock()
		err := p.Types[k].compileError
		p.RUnlock()
		if err != nil {
			return err
		}
	}

	p.compiled = true
	return nil
}

type typeDef struct {
	compileError    error
	PkgName         string
	Import          []string
	Self            string
	Name            string `json:"-"`
	DivDurationType string
	MulDurationType string
	Units           []unitDef
	UnitGenerators  []unitGenerator
}

// compileRoutine should be started concurrently with go, and all typeDefs should be passed the same WaitGroup lookup.
func (t typeDef) compileRoutine(name string, p *pkgConfig, wgs map[string]*sync.WaitGroup) {
	defer wgs[name].Done()
	defer func() {
		p.Lock()
		p.Types[name] = t
		p.Unlock()
	}()

	var hasTime bool

	t.Name = name
	t.PkgName = p.PkgName

	// Append import time if needed.
	if len(t.DivDurationType) > 0 || len(t.MulDurationType) > 0 {
		hasTime = false
		for _, s := range t.Import {
			if s == "time" {
				hasTime = true
				break
			}
		}
		if !hasTime {
			t.Import = append(t.Import, "time")
		}
	}

	// Generate units -- wait for dependent types to complete.
	for i, g := range t.UnitGenerators {
		wg, ok := wgs[g.From]
		if !ok {
			t.compileError = fmt.Errorf("types.%s.unitGenerators[%d].from: '%s' not found in type lookup",
				name, i, g.From,
			)
			return
		}
		wg.Wait()

		p.RLock()
		units, err := g.units(name, p.Types)
		p.RUnlock()
		if err != nil {
			t.compileError = fmt.Errorf("types.%s.unitGenerators[%d]%s", name, i, err.Error())
			return
		}
		t.Units = append(t.Units, units...)
	}
	if len(t.UnitGenerators) > 0 {
		t.UnitGenerators = nil
	}
}

type unitDef struct {
	Name       string
	NamePlural string
	Value      string
}

type unitGenerator struct {
	From        string
	Pattern     string
	NameSuffix  string
	ValueSuffix string
}

func (g unitGenerator) units(name string, types map[string]typeDef) ([]unitDef, error) {
	var r *regexp.Regexp
	var err error

	t, ok := types[g.From]
	if !ok {
		return nil, fmt.Errorf(".from: '%s' not found in type lookup", g.From)
	}
	if g.Pattern != "" {
		r, err = regexp.Compile(g.Pattern)
		if err != nil {
			return nil, fmt.Errorf(".pattern: %s", err.Error())
		}
	}
	units := make([]unitDef, 0, len(t.Units))
	for _, from := range t.Units {
		if r != nil && !r.MatchString(from.Name) {
			continue
		}
		units = append(units, unitDef{
			Name:       from.Name + g.NameSuffix,
			NamePlural: from.NamePlural + g.NameSuffix,
			Value:      fmt.Sprintf("%s(%s)%s", name, from.Name, g.ValueSuffix),
		})
	}
	return units, nil
}

var tmplFunctions = template.FuncMap{
	"lower": strings.ToLower,
}

const tmpl = `// This code is generated by units-codegen; do not manualy edit this file.
{{$self := .Self}}{{$name := .Name}}
package {{.PkgName}}

{{if .Import}}import ({{range .Import}}
	"{{.}}"{{end}}
){{end}}

// Units for {{.Name}} values. Always multiply with a unit when setting the initial value like you would for
// time.Time. This prevents you from having to worry about the internal storage format.
const ({{range .Units}}
	{{.Name}} {{$name}} = {{.Value}}{{end}}
)
{{range .Units}}
// {{.NamePlural}} returns {{$self}} as a floating point number of {{.NamePlural|lower}}.
func ({{$self}} {{$name}}) {{.NamePlural}}() float64 {
	return float64({{$self}} / {{.Name}})
}
{{end}}
// Abs returns the absolute value of {{.Self}} as a copy.
func ({{.Self}} {{.Name}}) Abs() {{.Name}} {
	if {{.Self}} < 0 {
		return -{{.Self}}
	}
	return {{.Self}}
}

// Mul returns the product of {{.Self}} * x as a new {{.Name}}.
func ({{.Self}} {{.Name}}) Mul(x float64) {{.Name}} {
	return {{.Self}} * {{.Name}}(x)
}

// Div returns the quotient of {{.Self}} / x as a new {{.Name}}.
func ({{.Self}} {{.Name}}) Div(x float64) {{.Name}} {
	return {{.Self}} / {{.Name}}(x)
}

// Div{{.Name}} returns the quotient of {{.Self}} / x as a floating point number.
func ({{.Self}} {{.Name}}) Div{{.Name}}(x {{.Name}}) float64 {
	return float64({{.Self}} / x)
}
{{if .DivDurationType}}
// DivDuration returns the quotient of {{.Self}} / t as a {{.DivDurationType}}.
func ({{.Self}} {{.Name}}) DivDuration(t time.Duration) {{.DivDurationType}} {
	return {{.DivDurationType}}(float64({{.Self}}) / float64(t))
}

// Div{{.DivDurationType}} returns the quotient of {{.Self}} / x as a time.Duration.
func ({{.Self}} {{.Name}}) Div{{.DivDurationType}}(x {{.DivDurationType}}) time.Duration {
	return time.Duration(float64({{.Self}}) / float64(x))
}{{end}}
{{if .MulDurationType}}
// MulDuration returns the product of {{.Self}} * t as a {{.MulDurationType}}.
func ({{.Self}} {{.Name}}) MulDuration(t time.Duration) {{.MulDurationType}} {
	return {{.MulDurationType}}(float64({{.Self}}) * float64(t))
}{{end}}
`