// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build ignore

package main

// this generates type-instantiated boilerplate code for
// slice operations and tests

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io"
	"os"
	"strings"
	"text/template"
)

// shapes describes a combination of vector widths and various element types
type shapes struct {
	vecs   []int // Vector bit width for this shape.
	ints   []int // Int element bit width(s) for this shape
	uints  []int // Unsigned int element bit width(s) for this shape
	floats []int // Float element bit width(s) for this shape
}

// shapeAndTemplate is a template and the set of shapes on which it will be expanded
type shapeAndTemplate struct {
	s *shapes
	t *template.Template
}

var allShapes = &shapes{
	vecs:   []int{128, 256, 512},
	ints:   []int{8, 16, 32, 64},
	uints:  []int{8, 16, 32, 64},
	floats: []int{32, 64},
}

// these are the shapes that are currently converted to int32
// (not all conversions are available, yet)
var toInt32Shapes = &shapes{
	vecs:   []int{128, 256, 512},
	floats: []int{32},
}

func oneTemplate(t *template.Template, baseType string, width, count int, out io.Writer) {
	b := width * count
	if b < 128 || b > 512 {
		return
	}
	BaseType := strings.ToUpper(baseType[:1]) + baseType[1:]
	eType := fmt.Sprintf("%s%d", baseType, width)
	wxc := fmt.Sprintf("%dx%d", width, count)
	vType := fmt.Sprintf("%s%s", BaseType, wxc)
	aOrAn := "a"
	if strings.Contains("aeiou", baseType[:1]) {
		aOrAn = "an"
	}
	t.Execute(out, struct {
		Vec   string // the type of the vector, e.g. Float32x4
		AOrAn string // for documentation, the article "a" or "an"
		Width int    // the bit width of the element type, e.g. 32
		Count int    // the number of elements, e.g. 4
		WxC   string // the width-by-type string, e.g., "32x4"
		Type  string // the element type, e.g. "float32"
	}{
		Vec:   vType,
		AOrAn: aOrAn,
		Width: width,
		Count: count,
		WxC:   wxc,
		Type:  eType,
	})
}

// forTemplates expands the template sat.t for each shape
// in sat.s, writing to out.
func (sat shapeAndTemplate) forTemplates(out io.Writer) {
	t, s := sat.t, sat.s
	vecs := s.vecs
	ints := s.ints
	uints := s.uints
	floats := s.floats
	for _, v := range vecs {
		for _, w := range ints {
			c := v / w
			oneTemplate(t, "int", w, c, out)
		}
		for _, w := range uints {
			c := v / w
			oneTemplate(t, "uint", w, c, out)
		}
		for _, w := range floats {
			c := v / w
			oneTemplate(t, "float", w, c, out)
		}
	}
}

func prologue(s string, out io.Writer) {
	fmt.Fprintf(out,
		`// Code generated by '%s'; DO NOT EDIT.

//go:build goexperiment.simd

package simd

`, s)
}

func testPrologue(t, s string, out io.Writer) {
	fmt.Fprintf(out,
		`// Code generated by '%s'; DO NOT EDIT.

//go:build goexperiment.simd

// This file contains functions testing %s.
// Each function in this file is specialized for a
// particular simd type <BaseType><Width>x<Count>.

package simd_test

import (
	"simd"
	"testing"
)

`, s, t)
}

func curryTestPrologue(t string) func(s string, out io.Writer) {
	return func(s string, out io.Writer) {
		testPrologue(t, s, out)
	}
}

// //go:noescape
// func LoadUint8x16Slice(s []uint8) Uint8x16 {
// 	return LoadUint8x16((*[16]uint8)(s[:16]))
// }

// //go:noescape
// func (x Uint8x16) StoreSlice(s []uint8) {
//    x.Store((*[16]uint8)(s[:16]))
// }

func templateOf(name, temp string) shapeAndTemplate {
	return shapeAndTemplate{s: allShapes,
		t: template.Must(template.New(name).Parse(temp))}
}

func shapedTemplateOf(s *shapes, name, temp string) shapeAndTemplate {
	return shapeAndTemplate{s: s,
		t: template.Must(template.New(name).Parse(temp))}
}

var sliceTemplate = templateOf("slice", `
// Load{{.Vec}}Slice loads {{.AOrAn}} {{.Vec}} from a slice of at least {{.Count}} {{.Type}}s
func Load{{.Vec}}Slice(s []{{.Type}}) {{.Vec}} {
	return Load{{.Vec}}((*[{{.Count}}]{{.Type}})(s))
}

// StoreSlice stores x into a slice of at least {{.Count}} {{.Type}}s
func (x {{.Vec}}) StoreSlice(s []{{.Type}}) {
	x.Store((*[{{.Count}}]{{.Type}})(s))
}
`)

var unaryTemplate = templateOf("unary_helpers", `
// test{{.Vec}}Unary tests the simd unary method f against the expected behavior generated by want
func test{{.Vec}}Unary(t *testing.T, f func(_ simd.{{.Vec}}) simd.{{.Vec}}, want func(_ []{{.Type}}) []{{.Type}}) {
	n := {{.Count}}
	t.Helper()
	forSlice(t, {{.Type}}s, n, func(x []{{.Type}}) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		g := make([]{{.Type}}, n)
		f(a).StoreSlice(g)
		w := want(x)
		return checkSlicesLogInput(t, g, w, func() {t.Helper(); t.Logf("x=%v", x)})
	})
}
`)

var unaryTemplateToInt32 = shapedTemplateOf(toInt32Shapes, "unary_int32_helpers", `
// test{{.Vec}}Unary tests the simd unary method f against the expected behavior generated by want
func test{{.Vec}}UnaryToInt32(t *testing.T, f func(x simd.{{.Vec}}) simd.Int32x{{.Count}}, want func(x []{{.Type}}) []int32) {
	n := {{.Count}}
	t.Helper()
	forSlice(t, {{.Type}}s, n, func(x []{{.Type}}) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		g := make([]int32, n)
		f(a).StoreSlice(g)
		w := want(x)
		return checkSlicesLogInput(t, g, w, func() {t.Helper(); t.Logf("x=%v", x)})
	})
}
`)

var binaryTemplate = templateOf("binary_helpers", `
// test{{.Vec}}Binary tests the simd binary method f against the expected behavior generated by want
func test{{.Vec}}Binary(t *testing.T, f func(_, _ simd.{{.Vec}}) simd.{{.Vec}}, want func(_, _ []{{.Type}}) []{{.Type}}) {
	n := {{.Count}}
	t.Helper()
	forSlicePair(t, {{.Type}}s, n, func(x, y []{{.Type}}) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		b := simd.Load{{.Vec}}Slice(y)
		g := make([]{{.Type}}, n)
		f(a, b).StoreSlice(g)
		w := want(x, y)
		return checkSlicesLogInput(t, g, w, func() {t.Helper(); t.Logf("x=%v", x); t.Logf("y=%v", y); })
	})
}
`)

var ternaryTemplate = templateOf("ternary_helpers", `
// test{{.Vec}}Ternary tests the simd ternary method f against the expected behavior generated by want
func test{{.Vec}}Ternary(t *testing.T, f func(_, _, _ simd.{{.Vec}}) simd.{{.Vec}}, want func(_, _, _ []{{.Type}}) []{{.Type}}) {
	n := {{.Count}}
	t.Helper()
	forSliceTriple(t, {{.Type}}s, n, func(x, y, z []{{.Type}}) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		b := simd.Load{{.Vec}}Slice(y)
		c := simd.Load{{.Vec}}Slice(z)
		g := make([]{{.Type}}, n)
		f(a, b, c).StoreSlice(g)
		w := want(x, y, z)
		return checkSlicesLogInput(t, g, w, func() {t.Helper(); t.Logf("x=%v", x); t.Logf("y=%v", y); t.Logf("z=%v", z); })
	})
}
`)

var compareTemplate = templateOf("compare_helpers", `
// test{{.Vec}}Compare tests the simd comparison method f against the expected behavior generated by want
func test{{.Vec}}Compare(t *testing.T, f func(_, _ simd.{{.Vec}}) simd.Mask{{.WxC}}, want func(_, _ []{{.Type}}) []int64) {
	n := {{.Count}}
	t.Helper()
	forSlicePair(t, {{.Type}}s, n, func(x, y []{{.Type}}) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		b := simd.Load{{.Vec}}Slice(y)
		g := make([]int{{.Width}}, n)
		f(a, b).AsInt{{.WxC}}().StoreSlice(g)
		w := want(x, y)
		return checkSlicesLogInput(t, s64(g), w, func() {t.Helper(); t.Logf("x=%v", x); t.Logf("y=%v", y); })
	})
}
`)

// TODO this has not been tested yet.
var compareMaskedTemplate = templateOf("comparemasked_helpers", `
// test{{.Vec}}CompareMasked tests the simd masked comparison method f against the expected behavior generated by want
// The mask is applied to the output of want; anything not in the mask, is zeroed.
func test{{.Vec}}CompareMasked(t *testing.T,
	f func(_, _ simd.{{.Vec}}, m simd.Mask{{.WxC}}) simd.Mask{{.WxC}},
	want func(_, _ []{{.Type}}) []int64) {
	n := {{.Count}}
	t.Helper()
	forSlicePairMasked(t, {{.Type}}s, n, func(x, y []{{.Type}}, m []bool) bool {
	 	t.Helper()
		a := simd.Load{{.Vec}}Slice(x)
		b := simd.Load{{.Vec}}Slice(y)
		k := simd.LoadInt{{.WxC}}Slice(toVect[int{{.Width}}](m)).AsMask{{.WxC}}()
		g := make([]int{{.Width}}, n)
		f(a, b, k).AsInt{{.WxC}}().StoreSlice(g)
		w := want(x, y)
		for i := range m {
			if !m[i] {
				w[i] = 0
			}
		}
		return checkSlicesLogInput(t, s64(g), w, func() {t.Helper(); t.Logf("x=%v", x); t.Logf("y=%v", y); t.Logf("m=%v", m); })
	})
}
`)

func main() {
	sl := flag.String("sl", "slice_amd64.go", "file name for slice operations")
	bh := flag.String("bh", "binary_helpers_test.go", "file name for binary test helpers")
	uh := flag.String("uh", "unary_helpers_test.go", "file name for unary test helpers")
	th := flag.String("th", "ternary_helpers_test.go", "file name for ternary test helpers")
	ch := flag.String("ch", "compare_helpers_test.go", "file name for compare test helpers")
	cmh := flag.String("cmh", "comparemasked_helpers_test.go", "file name for compare-masked test helpers")
	flag.Parse()

	if *sl != "" {
		one(*sl, prologue, sliceTemplate)
	}
	if *uh != "" {
		one(*uh, curryTestPrologue("unary simd methods"), unaryTemplate)
	}
	if *bh != "" {
		one(*bh, curryTestPrologue("binary simd methods"), binaryTemplate)
	}
	if *th != "" {
		one(*th, curryTestPrologue("ternary simd methods"), ternaryTemplate)
	}
	if *ch != "" {
		one(*ch, curryTestPrologue("simd methods that compare two operands"), compareTemplate)
	}
	if *cmh != "" {
		one(*cmh, curryTestPrologue("simd methods that compare two operands under a mask"), compareMaskedTemplate)
	}
}

func one(filename string, prologue func(s string, out io.Writer), sats ...shapeAndTemplate) {
	if filename == "" {
		return
	}

	ofile := os.Stdout

	if filename != "-" {
		var err error
		ofile, err = os.Create(filename)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not create the output file %s for the generated code, %v", filename, err)
			os.Exit(1)
		}
	}

	out := new(bytes.Buffer)

	prologue("go run genfiles.go", out)
	for _, sat := range sats {
		sat.forTemplates(out)
	}

	b, err := format.Source(out.Bytes())
	if err != nil {
		fmt.Fprintf(os.Stderr, "There was a problem formatting the generated code for %s, %v", filename, err)
		os.Exit(1)
	} else {
		ofile.Write(b)
		ofile.Close()
	}

}