go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/projects/nodes/cmd/nodegen/main.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package main

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"slices"
	"strings"
	"text/template"
	"unicode"

	"golang.org/x/tools/imports"

	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/sdk/iter"
)

func init() {
	flag.Parse()
	inputTypeToOutputTypes = generateOutputTypesForInputType()
}

func main() {
	var inputTypesFunc func() []string
	var outputTypesFunc any
	var templateBody string
	switch *flagNodeMode {
	case "both": // emit instances for both input and output types
		templateBody = templateBoth
		outputTypesFunc = getOutputTypesForInput
		inputTypesFunc = getInputTypes
	case "input": // emit instances just for input types
		templateBody = templateInput
		outputTypesFunc = func() []string { return nil }
		inputTypesFunc = getInputTypes
	case "output": // emit instances just for output types
		templateBody = templateOutput
		outputTypesFunc = getOutputTypes
		inputTypesFunc = getInputTypes
	case "sort": // a special case for the sort node type
		templateBody = templateSort
		outputTypesFunc = func() []string { return nil }
		inputTypesFunc = getSortInputTypes
	case "none":
		templateBody = templateNone
		outputTypesFunc = func() []string { return nil }
		inputTypesFunc = func() []string { return nil }
	default:
		fmt.Fprintf(os.Stderr, "invalid node mode %v\n", *flagNodeMode)
		os.Exit(1)
	}

	// if *flagDebug {
	// 	fmt.Printf("using input-types: %s\n", strings.Join(inputTypesFunc(), ","))
	// 	if typed, ok := outputTypesFunc.(func() []string); ok {
	// 		fmt.Printf("using output-types: %s\n", strings.Join(typed(), ","))
	// 	} else if typed, ok := outputTypesFunc.(func(string) []string); ok {
	// 		fullTypes := make(map[string][]string)
	// 		for _, inputType := range inputTypesFunc() {
	// 			fullTypes[inputType] = typed(inputType)
	// 		}
	// 		fmt.Printf("using output-types: %#v\n", fullTypes)
	// 	}
	// 	// fmt.Printf("using template:\n%s\n", templateBody)
	// }

	t, err := template.New("").Funcs(template.FuncMap{
		"toExported":  toExported,
		"toLower":     toLower,
		"toScalar":    toScalar,
		"toArray":     toArray,
		"apiType":     incrutil.APIValueTypeForGoType,
		"incrFn":      formatIncrFn,
		"incrManyFn":  formatIncrManyFn,
		"incrType":    formatIncrType,
		"isArray":     isArray,
		"outputTypes": outputTypesFunc,
	}).Parse(templateBody)
	if err != nil {
		fmt.Fprintf(os.Stderr, "template parse error: %+v\n", err)
		os.Exit(1)
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, map[string]any{
		"inputTypes": inputTypesFunc(),
		"nodeType":   *flagNodeType,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "template execution error: %+v\n", err)
		os.Exit(1)
	}
	formatted, err := imports.Process(flag.Args()[0], buf.Bytes(), nil)
	if err != nil {
		// fmt.Fprintln(os.Stdout, buf.String())
		fmt.Fprintf(os.Stderr, "format source error: %+v\n", err)
		os.Exit(1)
	}
	os.WriteFile(flag.Args()[0], formatted, 0644)
}

var flagDebug = flag.Bool("debug", false, "If debug output should be shown.")

var flagNodeType = flag.String("node-type", "merge", "The node type (will be a component of the exported function name)")
var flagNodeMode = flag.String("node-mode", "both", "The node mode (e.g. input|output|both)")

var flagIncrType = flag.String("incr-type", "", "The incremental node type")
var flagIncrFn = flag.String("incr-fn", "", "The incremental function value (can be a template)")
var flagIncrManyFn = flag.String("incr-many-fn", "", "The incremental function name for array types")

var flagInputTypes = flag.String("i-types", "", "specific input types to include")
var flagInputTypesScalar = flag.Bool("i-scalar", false, "if we should include scalar input types")
var flagInputTypesArray = flag.Bool("i-array", false, "if we should include array input types")
var flagInputTypesTable = flag.Bool("i-table", false, "if we should include the special table input type")
var flagInputTypesSVG = flag.Bool("i-svg", false, "if we should include the special svg input type")
var flagInputTypesAny = flag.Bool("i-any", false, "if we should include the special any input type")
var flagInputTypesAnyArray = flag.Bool("i-any-array", false, "if we should include the special []any input type")

var flagInputTypesFilterMath = flag.Bool("i-filter-math", false, "if input types should be math only")

var flagOutputTypes = flag.String("o-types", "", "specific output types to include")
var flagOutputTypesScalar = flag.Bool("o-scalar", false, "if we should include scalar output types")
var flagOutputTypesArray = flag.Bool("o-array", false, "if we should include array output types")
var flagOutputTypesTable = flag.Bool("o-table", false, "if we should include the special table output type")
var flagOutputTypesSVG = flag.Bool("o-svg", false, "if we should include the special svg output type")
var flagOutputTypesAny = flag.Bool("o-any", false, "if we should include the special any output type")
var flagOutputTypesAnyArray = flag.Bool("o-any-array", false, "if we should include the special []any output type")

var flagOutputTypesFilterMath = flag.Bool("o-filter-math", false, "if output types should be math only")

var flagOutputTypesMatchInput = flag.Bool("o-match-input", false, "if the output type should match the include type")
var flagOutputTypesMatchInputAsScalar = flag.Bool("o-match-input-scalar", false, "if the output type should match the include type as a scalar type")
var flagOutputTypesMatchInputAsArray = flag.Bool("o-match-input-array", false, "if the output type should match the include type as an array type")

var inputTypeToOutputTypes map[string][]string

func generateOutputTypesForInputType() map[string][]string {
	var inputToOutputTypes = make(map[string][]string)

	if *flagInputTypes != "" {
		for _, t := range iter.Apply(strings.Split(trimQuotes(*flagInputTypes), ","), trimQuotes) {
			inputToOutputTypes[t] = []string{}
		}
	} else {
		if *flagInputTypesScalar {
			for _, t := range incrutil.ValueScalarTypes {
				if !*flagInputTypesFilterMath || (*flagInputTypesFilterMath && incrutil.ValueTypeIsMath(t)) {
					inputToOutputTypes[t] = []string{}
				}
			}
		}
		if *flagInputTypesArray {
			for _, t := range incrutil.ValueArrayTypes {
				if !*flagInputTypesFilterMath || (*flagInputTypesFilterMath && incrutil.ValueTypeIsMath(t)) {
					inputToOutputTypes[t] = []string{}
				}
			}
		}
		if *flagInputTypesTable {
			inputToOutputTypes[incrutil.ValueTypeTable] = []string{}
		}
		if *flagInputTypesSVG {
			inputToOutputTypes[incrutil.ValueTypeSVG] = []string{}
		}
		if *flagInputTypesAny {
			inputToOutputTypes[incrutil.ValueTypeAny] = []string{}
		}
		if *flagInputTypesAnyArray {
			inputToOutputTypes[incrutil.ValueTypeAnyArray] = []string{}
		}
	}

	for t := range inputToOutputTypes {
		if *flagOutputTypes != "" {
			inputToOutputTypes[t] = iter.Apply(strings.Split(trimQuotes(*flagOutputTypes), ","), trimQuotes)
		} else {
			if *flagOutputTypesMatchInput {
				if !*flagOutputTypesFilterMath || (*flagOutputTypesFilterMath && incrutil.ValueTypeIsMath(t)) {
					inputToOutputTypes[t] = []string{t}
				}
			} else if *flagOutputTypesMatchInputAsArray {
				if !*flagOutputTypesFilterMath || (*flagOutputTypesFilterMath && incrutil.ValueTypeIsMath(t)) {
					if isArray(t) {
						inputToOutputTypes[t] = []string{t}
					} else {
						inputToOutputTypes[t] = []string{"[]" + t}
					}
				}
			} else if *flagOutputTypesMatchInputAsScalar {
				if !*flagOutputTypesFilterMath || (*flagOutputTypesFilterMath && incrutil.ValueTypeIsMath(t)) {
					if isArray(t) {
						inputToOutputTypes[t] = []string{toScalar(t)}
					} else {
						inputToOutputTypes[t] = []string{t}
					}
				}
			} else {
				if *flagOutputTypesScalar {
					if *flagOutputTypesFilterMath {
						inputToOutputTypes[t] = iter.Filter(incrutil.ValueScalarTypes, incrutil.ValueTypeIsMath)
					} else {
						inputToOutputTypes[t] = incrutil.ValueScalarTypes
					}
				}
				if *flagOutputTypesArray {
					if *flagOutputTypesFilterMath {
						inputToOutputTypes[t] = append(inputToOutputTypes[t], iter.Filter(incrutil.ValueArrayTypes, incrutil.ValueTypeIsMath)...)
					} else {
						inputToOutputTypes[t] = append(inputToOutputTypes[t], incrutil.ValueArrayTypes...)
					}
				}
				if *flagOutputTypesTable {
					inputToOutputTypes[t] = append(inputToOutputTypes[t], incrutil.ValueTypeTable)
				}
				if *flagOutputTypesSVG {
					inputToOutputTypes[t] = append(inputToOutputTypes[t], incrutil.ValueTypeSVG)
				}
				if *flagOutputTypesAny {
					inputToOutputTypes[t] = append(inputToOutputTypes[t], incrutil.ValueTypeAny)
				}
				if *flagOutputTypesAnyArray {
					inputToOutputTypes[t] = append(inputToOutputTypes[t], incrutil.ValueTypeAnyArray)
				}
			}
		}
	}
	return inputToOutputTypes
}

func getOutputTypesForInput(inputType string) []string {
	outputTypes, _ := inputTypeToOutputTypes[inputType]
	slices.Sort(outputTypes)
	return outputTypes
}

func getInputTypes() []string {
	var inputTypes []string
	for k := range inputTypeToOutputTypes {
		inputTypes = append(inputTypes, k)
	}
	slices.Sort(inputTypes)
	return inputTypes
}

func getOutputTypes() []string {
	var outputTypes []string
	if *flagOutputTypes != "" {
		outputTypes = iter.Apply(strings.Split(trimQuotes(*flagOutputTypes), ","), trimQuotes)
	} else {
		if *flagOutputTypesScalar {
			outputTypes = append(outputTypes, incrutil.ValueScalarTypes...)
		}
		if *flagOutputTypesArray {
			outputTypes = append(outputTypes, incrutil.ValueArrayTypes...)
		}
		if *flagOutputTypesTable {
			outputTypes = append(outputTypes, incrutil.ValueTypeTable)
		}
		if *flagOutputTypesSVG {
			outputTypes = append(outputTypes, incrutil.ValueTypeSVG)
		}
		if *flagOutputTypesAny {
			outputTypes = append(outputTypes, incrutil.ValueTypeAny)
		}
		if *flagOutputTypesFilterMath {
			outputTypes = iter.Filter(outputTypes, incrutil.ValueTypeIsMath)
		}
	}
	slices.Sort(outputTypes)
	return outputTypes
}

func renderFunctionSnippet(tmpl, inputType, outputType string) (string, error) {
	t, err := template.New("").Funcs(template.FuncMap{
		"toScalar":   toScalar,
		"toArray":    toArray,
		"toLower":    toLower,
		"toExported": toExported,
	}).Parse(trimQuotes(tmpl))
	if err != nil {
		return "", err
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, map[string]any{
		"inputType":  inputType,
		"outputType": outputType,
	})
	if err != nil {
		return "", err
	}
	return buf.String(), nil
}

func formatIncrFn(inputType, outputType string) (string, error) {
	return renderFunctionSnippet(*flagIncrFn, inputType, outputType)
}

func formatIncrManyFn(inputType, outputType string) (string, error) {
	var fn string
	if flagIncrManyFn != nil && *flagIncrManyFn != "" {
		fn = *flagIncrManyFn
	} else {
		fn = *flagIncrFn
	}
	return renderFunctionSnippet(fn, inputType, outputType)
}

func formatIncrType(inputType, outputType string) (string, error) {
	return renderFunctionSnippet(*flagIncrType, inputType, outputType)
}

func isArray(typeName string) bool {
	return strings.HasPrefix(strings.TrimSpace(typeName), "[]")
}

func getSortInputTypes() []string {
	return iter.Filter(incrutil.ValueTypes, func(v string) bool { return v != "[]bool" && v != "bool" })
}

const templateBoth = `
// File generated by nodegen. DO NOT EDIT.
package gen

import (
	"fmt"
	"time"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/projects/nodes/pkg/types"
	"go.charczuk.com/projects/nodes/pkg/funcs"
)

{{ $vm := . }}
func {{ .nodeType | toExported }}ForNode(graph *incr.Graph, n *types.Node) (output incr.INode, err error) {
	switch n.Metadata.InputType {
	{{- range $inputIndex, $inputType := $vm.inputTypes }}
	case "{{ $inputType | apiType }}":
		switch n.Metadata.OutputType {
		{{- range $outputIndex, $outputType := outputTypes $inputType }}
		case "{{ $outputType | apiType }}":
		{{- if isArray $inputType }}
			output = {{ incrType $inputType $outputType }}(graph, {{ incrManyFn $inputType $outputType }})
			return
		{{- else }}
			output = {{ incrType $inputType $outputType }}(graph, {{ incrFn $inputType $outputType }})
			return
		{{- end }}
		{{- end }}
		default:
			err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} output type %v for input type %v", n.Metadata.OutputType, n.Metadata.InputType)
			return
		}
	{{ end }}
	default:
		err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} input type %v", n.Metadata.InputType)
		return
	}
}
`

const templateSort = `
// File generated by nodegen. DO NOT EDIT.
package gen

import (
	"fmt"
	"time"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/sdk/iter"
	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/projects/nodes/pkg/types"
	"go.charczuk.com/projects/nodes/pkg/funcs"
)

{{ $vm := . }}
func {{ .nodeType | toExported }}ForNode(graph *incr.Graph, n *types.Node) (output incr.INode, err error) {
	switch n.Metadata.InputType {
	{{- range $inputIndex, $inputType := $vm.inputTypes }}
	case "{{ $inputType | apiType }}":
		switch n.Metadata.OutputType {
		case "{{ $inputType | apiType }}":
		{{- if $inputType | eq "[]time.Time" }}
			output = incrutil.MapN[{{ $inputType }}](graph, funcs.SortComparerMany[{{ $inputType | toScalar }}](iter.SorterComparerFunc[{{ $inputType | toScalar }}](funcs.TimestampAsc)))
		{{- else if $inputType | eq "time.Time" }}
			output = incrutil.MapN[{{ $inputType }}](graph, funcs.SortComparer[{{ $inputType | toScalar }}](iter.SorterComparerFunc[{{ $inputType | toScalar }}](funcs.TimestampAsc)))
		{{- else }}
		{{- if isArray $inputType }}
			output = incrutil.MapN[{{ $inputType }}](graph, funcs.SortMany[{{ $inputType | toScalar }}])
		{{- else }}
			output = incrutil.MapN[{{ $inputType }}](graph, funcs.Sort[{{ $inputType | toScalar }}])
		{{- end }}
		{{- end }}
			return
		default:
			err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} output type %v for input type %v", n.Metadata.OutputType, n.Metadata.InputType)
			return
		}
	{{ end }}
	default:
		err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} input type %v", n.Metadata.InputType)
		return
	}
}
`

const templateInput = `
// File generated by nodegen. DO NOT EDIT.
package gen

import (
	"fmt"
	"time"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/projects/nodes/pkg/types"
	"go.charczuk.com/projects/nodes/pkg/funcs"
)

{{ $vm := . }}
func {{ .nodeType | toExported }}ForNode(graph *incr.Graph, n *types.Node) (output incr.INode, err error) {
	switch n.Metadata.InputType {
	{{- range $inputIndex, $inputType := .inputTypes }}
	case "{{ $inputType | apiType }}":
		output = {{ incrType $inputType "" }}(graph, {{ incrFn $inputType "" }})
		return
	{{- end }}
	default:
		err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} input type %v", n.Metadata.InputType)
		return
	}
}`
const templateOutput = `
// File generated by nodegen. DO NOT EDIT.
package gen

import (
	"fmt"
	"time"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/projects/nodes/pkg/types"
	"go.charczuk.com/projects/nodes/pkg/funcs"
)

{{ $vm := . }}
func {{ .nodeType | toExported }}ForNode(graph *incr.Graph, n *types.Node) (output incr.INode, err error) {
	switch n.Metadata.OutputType {
	{{ range $outputIndex, $outputType := outputTypes }}
	case "{{ $outputType | apiType }}":
	{{- if isArray $outputType }}
		output = {{ incrType "" $outputType }}(graph, {{ incrManyFn "" $outputType }})
		return
	{{- else }}
		output = {{ incrType "" $outputType }}(graph, {{ incrFn "" $outputType }})
		return
	{{- end }}
	{{- end }}
	default:
		err = fmt.Errorf("invalid {{ $vm.nodeType | toLower }} output type %v", n.Metadata.OutputType)
		return
	}
}`

const templateNone = `
// File generated by nodegen. DO NOT EDIT.
package gen

import (
	"fmt"
	"time"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/projects/nodes/pkg/incrutil"
	"go.charczuk.com/projects/nodes/pkg/types"
	"go.charczuk.com/projects/nodes/pkg/funcs"
)

{{ $vm := . }}
func {{ .nodeType | toExported }}ForNode(graph *incr.Graph, n *types.Node) (output incr.INode, err error) {
	output = {{ incrType "" "" }}(graph, {{ incrFn "" "" }})
	return
}`

func toExported(input string) string {
	if input == "" {
		return ""
	}
	inputRunes := []rune(input)
	inputRunes[0] = unicode.ToUpper(inputRunes[0])
	return string(inputRunes)
}

func toLower(input string) string {
	return strings.ToLower(input)
}

func toScalar(typeName string) string {
	return strings.TrimPrefix(strings.TrimSpace(typeName), "[]")
}

func toArray(typeName string) string {
	if strings.HasPrefix(typeName, "[]") {
		return typeName
	}
	return "[]" + typeName
}

func trimQuotes(str string) string {
	str = strings.TrimSpace(str)
	str = strings.TrimPrefix(str, `'`)
	str = strings.TrimPrefix(str, `"`)
	str = strings.TrimSuffix(str, `'`)
	str = strings.TrimSuffix(str, `"`)
	str = strings.TrimSpace(str)
	return str
}