gen_dtypes.go

//go:build ignore

// gen_dtypes writes zz_dtypes_gen.go: the type switches that turn a dtype or
// an untyped storage slice into a call on the concrete element type. Run it
// with `go generate`.
package main

import (
	"bytes"
	"fmt"
	"go/format"
	"os"
	"strings"
)

type typ struct {
	Go, DType string
	kind      byte // b(ool) i(nt) u(int) f(loat) c(omplex)
	bits      int
}

var types = []typ{
	{"bool", "Bool", 'b', 8},
	{"int8", "Int8", 'i', 8}, {"int16", "Int16", 'i', 16}, {"int32", "Int32", 'i', 32}, {"int64", "Int64", 'i', 64},
	{"uint8", "Uint8", 'u', 8}, {"uint16", "Uint16", 'u', 16}, {"uint32", "Uint32", 'u', 32}, {"uint64", "Uint64", 'u', 64},
	{"float32", "Float32", 'f', 32}, {"float64", "Float64", 'f', 64},
	{"complex64", "Complex64", 'c', 64}, {"complex128", "Complex128", 'c', 128},
}

func (t typ) isInt() bool { return t.kind == 'i' || t.kind == 'u' }

// fbits is the precision of a float or complex component, for formatting.
func (t typ) fbits() int {
	if t.kind == 'c' {
		return t.bits / 2
	}
	return t.bits
}

// narrow lists the types sums and products never run in: they accumulate in
// int64, uint64 or the input's own inexact type.
const narrow = "int8,int16,int32,uint8,uint16,uint32,float64"

var b bytes.Buffer

func p(format string, args ...any) { fmt.Fprintf(&b, format+"\n", args...) }

// conv returns the Go expression converting the value expression x of type
// s to type d.
func conv(s, d typ, x string) string {
	switch {
	case s == d:
		return x
	case d.kind == 'b':
		return x + " != 0"
	case s.kind == 'b':
		return "boolTo[" + d.Go + "](" + x + ")"
	case s.kind == 'c' && d.kind == 'c':
		return d.Go + "(" + x + ")"
	case s.kind == 'c':
		return conv(typ{"float64", "Float64", 'f', 64}, d, "float64(real("+x+"))")
	case d.kind == 'c':
		return d.Go + "(complex(float64(" + x + "), 0))"
	case s.kind == 'f' && d.isInt():
		return "satInt[" + d.Go + "](float64(" + x + "))"
	}
	return d.Go + "(" + x + ")"
}

// each emits one case per type whose kind letter is in set, substituting $T
// for the Go type and $D for the dtype constant. A set may end in
// "-name,name" to leave those Go types out: the Float64 cases of the
// operations whose Float64 path never reaches the generic code, and the
// narrow integers of the accumulations, which always run in 64 bits.
func each(set string, body string) {
	set, skip, _ := strings.Cut(set, "-")
	for _, t := range types {
		if strings.IndexByte(set, t.kind) < 0 || strings.Contains(","+skip+",", ","+t.Go+",") {
			continue
		}
		p(strings.ReplaceAll(strings.ReplaceAll(body, "$T", t.Go), "$D", t.DType))
	}
}

func main() {
	p("// Code generated by gen_dtypes.go; DO NOT EDIT.")
	p("")
	p("package ndarray")
	p("")
	p(`import "strconv"`)
	p("")

	p("// dtypeOfStore returns the dtype of a storage slice.")
	p("func dtypeOfStore(s any) DType {")
	p("switch s.(type) {")
	each("biucf-float64", "case []$T: return $D") // fromStore handles []float64 itself
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// makeStore returns a zeroed storage slice of n elements of dtype dt.")
	p("func makeStore(dt DType, n int) any {")
	p("switch dt {")
	each("biucf", "case $D: return make([]$T, n)")
	p("}")
	p(`panic("ndarray: unknown dtype " + dt.String())`)
	p("}")
	p("")

	p("// storeLen returns the length of a storage slice.")
	p("func storeLen(s any) int {")
	p("switch s := s.(type) {")
	each("biucf", "case []$T: return len(s)")
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// cloneStore returns a copy of a storage slice.")
	p("func cloneStore(s any) any {")
	p("switch s := s.(type) {")
	each("biucf", "case []$T: return append([]$T(nil), s...)")
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// fillStore sets every element of dst to the single element of one, a")
	p("// storage slice of the same dtype.")
	p("func fillStore(dst, one any) {")
	p("switch d := dst.(type) {")
	each("biucf", "case []$T: v := one.([]$T)[0]; for i := range d { d[i] = v }")
	p("}")
	p("}")
	p("")

	p("// copyStore copies n elements from src[so:] to dst[do:], two storage slices")
	p("// of the same dtype.")
	p("func copyStore(dst any, do int, src any, so, n int) {")
	p("switch d := dst.(type) {")
	each("biucf", "case []$T: copy(d[do:do+n], src.([]$T)[so:so+n])")
	p("}")
	p("}")
	p("")

	p("// takeStore returns the elements of a storage slice at the given positions.")
	p("func takeStore(s any, idx []int) any {")
	p("switch s := s.(type) {")
	each("biucf", "case []$T: out := make([]$T, len(idx)); for k, i := range idx { out[k] = s[i] }; return out")
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// gatherStore returns a contiguous copy of the elements of the view a.")
	p("func gatherStore(a *Array) any {")
	p("switch s := a.store().(type) {")
	each("biucf-float64", "case []$T: return gather(s, a.shape, a.strides, a.offset)") // Float64 uses materialize
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// broadcastStore returns a contiguous copy of a's elements broadcast to shape.")
	p("func broadcastStore(a *Array, shape []int) any {")
	p("es := a.broadcastStrides(shape)")
	p("switch s := a.store().(type) {")
	each("biucf", "case []$T: return gather(s, shape, es, a.offset)")
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// convertStore converts every element of src into dst, two storage slices")
	p("// of the same length, as AsType documents.")
	p("func convertStore(dst, src any) {")
	p("switch s := src.(type) {")
	for _, st := range types {
		p("case []%s:", st.Go)
		p("switch d := dst.(type) {")
		for _, dt := range types {
			if st == dt {
				p("case []%s: copy(d, s)", dt.Go)
				continue
			}
			p("case []%s: for i, v := range s { d[i] = %s }", dt.Go, conv(st, dt, "v"))
		}
		p("}")
	}
	p("}")
	p("}")
	p("")

	c128 := typ{"complex128", "Complex128", 'c', 128}
	f64 := typ{"float64", "Float64", 'f', 64}
	p("// elemComplex returns element i of a storage slice as a complex128.")
	p("func elemComplex(s any, i int) complex128 {")
	p("switch s := s.(type) {")
	for _, t := range types {
		p("case []%s: return %s", t.Go, conv(t, c128, "s[i]"))
	}
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// setComplex stores v, converted, as element i of a storage slice.")
	p("func setComplex(s any, i int, v complex128) {")
	p("switch s := s.(type) {")
	for _, t := range types {
		p("case []%s: s[i] = %s", t.Go, conv(c128, t, "v"))
	}
	p("}")
	p("}")
	p("")

	p("// elemFloat returns element i of a storage slice of a real dtype as a float64.")
	p("func elemFloat(s any, i int) float64 {")
	p("switch s := s.(type) {")
	for _, t := range types {
		if t.kind != 'c' && t != f64 { // At reads Float64 storage itself
			p("case []%s: return %s", t.Go, conv(t, f64, "s[i]"))
		}
	}
	p("}")
	p(`panic("ndarray: complex element read as float64; use AtComplex")`)
	p("}")
	p("")

	p("// setFloat stores v, converted, as element i of a storage slice.")
	p("func setFloat(s any, i int, v float64) {")
	p("switch s := s.(type) {")
	for _, t := range types {
		if t != f64 { // Set writes Float64 storage itself
			p("case []%s: s[i] = %s", t.Go, conv(f64, t, "v"))
		}
	}
	p("}")
	p("}")
	p("")

	p("// elemString formats element i of a storage slice as NumPy prints it.")
	p("func elemString(s any, i int) string {")
	p("switch s := s.(type) {")
	for _, t := range types {
		switch t.kind {
		case 'b':
			p(`case []bool: if s[i] { return "true" }; return "false"`)
		case 'i':
			p("case []%s: return strconv.FormatInt(int64(s[i]), 10)", t.Go)
		case 'u':
			p("case []%s: return strconv.FormatUint(uint64(s[i]), 10)", t.Go)
		case 'f':
			p("case []%s: return floatString(float64(s[i]), %d)", t.Go, t.fbits())
		case 'c':
			p("case []%s: return complexString(complex128(s[i]), %d)", t.Go, t.fbits())
		}
	}
	p("}")
	p(`panic("ndarray: not a storage slice")`)
	p("}")
	p("")

	p("// arithStore applies a binary operation to storage slices of one dtype.")
	p("func arithStore(op binCode, d, x, y any) {")
	p("switch d := d.(type) {")
	each("iu", "case []$T: binInt(op, d, x.([]$T), y.([]$T))")
	each("f-float64", "case []$T: binFloat(op, d, x.([]$T), y.([]$T))")
	each("c", "case []$T: binComplex(op, d, x.([]$T), y.([]$T))")
	each("b", "case []$T: binBool(op, d, x.([]$T), y.([]$T))")
	p("}")
	p("}")
	p("")

	p("// cmpStore compares storage slices of one dtype into d.")
	p("func cmpStore(op cmpCode, d []bool, x, y any) {")
	p("switch x := x.(type) {")
	each("iuf", "case []$T: cmpOrdered(op, d, x, y.([]$T))")
	each("c", "case []$T: cmpComplex(op, d, x, y.([]$T))")
	each("b", "case []$T: cmpBool(op, d, x, y.([]$T))")
	p("}")
	p("}")
	p("")

	p("// accumulateStore sums or multiplies src along the middle axis of")
	p("// [outer][axisLen][inner] into dst; both have an accumulation dtype.")
	p("func accumulateStore(op reduceCode, dst, src any, outer, axisLen, inner int) {")
	p("switch d := dst.(type) {")
	each("iufc-"+narrow, "case []$T: reduceNum(op, d, src.([]$T), outer, axisLen, inner)")
	p("}")
	p("}")
	p("")

	p("// extremeStore takes the maximum or minimum of src along the middle axis.")
	p("func extremeStore(op reduceCode, dst, src any, outer, axisLen, inner int) {")
	p("switch d := dst.(type) {")
	each("iuf-float64", "case []$T: reduceOrdered(op, d, src.([]$T), outer, axisLen, inner)")
	each("c", "case []$T: reduceComplex(op, d, src.([]$T), outer, axisLen, inner)")
	each("b", "case []$T: reduceBool(op, d, src.([]$T), outer, axisLen, inner)")
	p("}")
	p("}")
	p("")

	p("// argStore writes the arg-max (or arg-min) along the middle axis into dst.")
	p("func argStore(larger bool, dst []int64, src any, outer, axisLen, inner int) {")
	p("switch s := src.(type) {")
	each("iuf", "case []$T: argOrdered(larger, dst, s, outer, axisLen, inner)")
	each("c", "case []$T: argComplex(larger, dst, s, outer, axisLen, inner)")
	each("b", "case []$T: argBool(larger, dst, s, outer, axisLen, inner)")
	p("}")
	p("}")
	p("")

	p("// scanStore is the cumulative sum or product along the middle axis.")
	p("func scanStore(op reduceCode, dst, src any, outer, axisLen, inner int) {")
	p("switch d := dst.(type) {")
	each("iufc-"+narrow, "case []$T: scanNum(op, d, src.([]$T), outer, axisLen, inner)")
	p("}")
	p("}")
	p("")

	p("// whereStore picks from t where c is true and from f elsewhere.")
	p("func whereStore(d any, c []bool, t, f any) {")
	p("switch d := d.(type) {")
	each("biucf", "case []$T: whereT(d, c, t.([]$T), f.([]$T))")
	p("}")
	p("}")
	p("")

	p("// intUnaryStore negates, takes the absolute value of, or squares integers.")
	p("func intUnaryStore(code unaryCode, d, x any) {")
	p("switch d := d.(type) {")
	each("iu", "case []$T: s := x.([]$T); switch code { case uNeg: negInt(d, s); case uAbs: absInt(d, s); default: squareInt(d, s) }")
	p("}")
	p("}")
	p("")

	p("// matmulStore multiplies (m x k) by (k x n) storage slices of one dtype into d.")
	p("func matmulStore(d, x, y any, m, k, n int) {")
	p("switch d := d.(type) {")
	each("iufc-float64", "case []$T: matmulNum(d, x.([]$T), y.([]$T), m, k, n)")
	each("b", "case []$T: matmulBool(d, x.([]$T), y.([]$T), m, k, n)")
	p("}")
	p("}")
	p("")

	src, err := format.Source(b.Bytes())
	if err != nil {
		os.Stdout.Write(b.Bytes())
		panic(err)
	}
	if err := os.WriteFile("zz_dtypes_gen.go", src, 0o644); err != nil {
		panic(err)
	}
}