// Copyright 2017 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.

// X86avxgen generates Go code for obj/x86 that adds AVX instructions support.
//
// Currently supports only AVX1 and AVX2 instructions.
// When x86.csv will contain AVX512 instructions and
// asm6.go is patched to support them,
// this program can be extended to generate the remainder.
//
// The output consists of multiple files:
// - cmd/internal/obj/x86/aenum.go
//	Add enum entries for new instructions.
// - cmd/internal/obj/x86/anames.go
//	Add new instruction names.
// - cmd/internal/obj/x86/vex_optabs.go
//	Add new instruction optabs.
// - cmd/asm/internal/asm/testdata/amd64enc.s
//	Uncomment tests for added instructions.
//
// Usage:
//	x86avxgen -goroot=$DEV_GOROOT [-csv=x86.csv] [-output=x86avxgen-output]
// $DEV_GOROOT is a path to Go repository working tree root.
//
// To get precise usage information, call this program without arguments.
package main

import (
	"bufio"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"golang.org/x/arch/x86/x86csv"
)

func main() {
	goroot := flag.String(
		"goroot", "",
		"Go sources root path")
	csv := flag.String(
		"csv", specFile,
		"Absolute path to x86spec CSV file")
	output := flag.String(
		"output", "x86avxgen-output",
		"Where to put output files")
	autopatchEnabled := flag.Bool(
		"autopatch", false,
		"Try automatic patching (writes to goroot, unsafe if it is not under VCS)")
	diagEnabled := flag.Bool(
		"diag", false,
		"Print debug information")
	flag.Parse()
	if len(os.Args) == 1 {
		fmt.Printf("%s: x86 AVX ytab generator", progName)
		flag.Usage()
		os.Exit(1)
	}
	if *goroot == "" {
		log.Fatal("goroot arg is mandatory")
	}
	if _, err := os.Stat(*csv); os.IsNotExist(err) {
		log.Fatalf("spec file %s not found", *csv)
	}

	r, err := specRowReader(*csv)
	if err != nil {
		log.Fatal(err)
	}
	if err := os.MkdirAll(*output, 0755); err != nil {
		log.Fatal(err)
	}

	opcodes, err := doGenerateVexOptabs(r, mustOpenFile(*output+"/"+filenameVexOptabs))
	if err != nil {
		log.Fatal(err)
	}
	if err := doGenerateAenum(*goroot, *output, opcodes); err != nil {
		log.Fatal(err)
	}
	if err := doGenerateAnames(*output); err != nil {
		log.Fatal(err)
	}
	if err := doGenerateTests(*goroot, *output, opcodes); err != nil {
		log.Fatal(err)
	}

	if *autopatchEnabled {
		if err := doAutopatch(*goroot, *output); err != nil {
			log.Fatal(err)
		}
	}

	if *diagEnabled {
		diag.Print()
	}
}

func mustOpenFile(path string) *os.File {
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatal(err)
	}
	return f
}

// filterVEX removes all non-VEX instructions from insts.
// Returns updates slice.
func filterVEX(insts []*x86csv.Inst) []*x86csv.Inst {
	vexInsts := insts[:0]
	for _, inst := range insts {
		// Checking CPUID for AVX is not good enough
		// in this case, because some instructions
		// have VEX prefix, but no AVX CPUID flag.
		if strings.HasPrefix(inst.Encoding, "VEX.") {
			vexInsts = append(vexInsts, inst)
		}
	}
	return vexInsts
}

func doGenerateVexOptabs(r *x86csv.Reader, w io.Writer) (opcodes []string, err error) {
	insts, err := r.ReadAll()
	if err != nil {
		return nil, err
	}
	insts = filterVEX(insts)

	var buf bytes.Buffer

	visitOptab := func(o optab) {
		diag.optabsGenerated++

		opcodes = append(opcodes, o.as)

		tmpl := "\t{A%s, %s, Pvex, [23]uint8{%s}},\n"
		fmt.Fprintf(&buf, tmpl, o.as, o.ytabID, strings.Join(o.op, ","))
	}

	doGroups(insts, func(op string, insts []*x86csv.Inst) {
		diag.optabsTotal++

		if ot, ok := precomputedOptabs[op]; ok {
			log.Printf("notice: using precomputed %s optab", op)
			visitOptab(ot)
			return
		}

		key := ytabKey(op, insts)
		ytabID := ytabMap[key]
		if ytabID == "" {
			diag.ytabMisses[key]++
			log.Printf("warning: skip %s: no ytabID for '%s' key", op, key)
			return
		}
		var encParts []string
		for _, inst := range insts {
			enc := parseEncoding(inst.Encoding)

			encParts = append(encParts, vexExpr(enc.vex))
			encParts = append(encParts, "0x"+enc.opbyte)
			if enc.opdigit != "" {
				encParts = append(encParts, "0"+enc.opdigit)
			}
		}
		visitOptab(optab{
			as:     op,
			ytabID: ytabID,
			op:     encParts,
		})
	})

	tmpl := `// Code generated by %s. DO NOT EDIT.

package x86

var vexOptab = []Optab{
%s
}
`
	code := []byte(fmt.Sprintf(tmpl, progName, buf.String()))

	prettyCode, err := format.Source(code)
	if err != nil {
		return nil, err
	}

	_, err = w.Write(prettyCode)

	return opcodes, err
}

func doGenerateAenum(goroot, output string, newNames []string) error {
	w, err := os.Create(output + "/" + filenameAenum)
	if err != nil {
		return err
	}
	defer w.Close()
	r, err := os.Open(goroot + "/" + pathAenum)
	if err != nil {
		return err
	}
	defer r.Close()

	return generateAenum(r, w, newNames)
}

func doGenerateAnames(output string) error {
	// Runs "go generate" over previously generated aenum file.
	path := output + "/" + filenameAenum
	cmd := exec.Command("go", "generate", path)
	var buf bytes.Buffer
	cmd.Stderr = &buf
	err := cmd.Run()
	if err != nil {
		return errors.New(err.Error() + ": " + buf.String())
	}
	return nil
}

// testLineReplacer is used in uncommentedTestLine function.
var testLineReplacer = strings.NewReplacer(
	"//TODO: ", "",

	// Fix register references.
	"XMM", "X",
	"YMM", "Y",
)

func uncommentedTestLine(line string) string {
	// Sync with x86/x86test/print.go.
	const x86testFmt = "\t%-39s // %s"

	line = testLineReplacer.Replace(line)
	i := strings.Index(line, " // ")
	return fmt.Sprintf(x86testFmt, line[len("\t"):i], line[i+len(" // "):])
}

// stringsSet returns a map mapping each x in xs to true.
func stringsSet(xs []string) map[string]bool {
	set := make(map[string]bool, len(xs))
	for _, x := range xs {
		set[x] = true
	}
	return set
}

func doGenerateTests(goroot, output string, newNames []string) error {
	testsFile, err := os.Open(goroot + "/" + pathTests)
	if err != nil {
		return err
	}
	defer testsFile.Close()

	var rxCommentedTestCase = regexp.MustCompile(`//TODO: ([A-Z][A-Z0-9]+)`)

	newNamesSet := stringsSet(newNames)

	var buf bytes.Buffer
	scanner := bufio.NewScanner(testsFile)
	for scanner.Scan() {
		line := scanner.Text()
		m := rxCommentedTestCase.FindStringSubmatch(line)
		if m != nil {
			name := string(m[1])
			if newNamesSet[name] {
				line = uncommentedTestLine(line)
			}
		}
		buf.WriteString(line)
		buf.WriteByte('\n')
	}

	return ioutil.WriteFile(output+"/"+filenameTests, buf.Bytes(), 0644)
}

func doAutopatch(goroot, output string) error {
	from := [...]string{
		output + "/" + filenameVexOptabs,
		output + "/" + filenameAenum,
		output + "/" + filenameAnames,
		output + "/" + filenameTests,
	}
	to := [...]string{
		goroot + "/" + pathVexOptabs,
		goroot + "/" + pathAenum,
		goroot + "/" + pathAnames,
		goroot + "/" + pathTests,
	}

	// No recovery if rename will fail.
	// There is a warning in "autopatch" description.
	for i := range from {
		if err := os.Rename(from[i], to[i]); err != nil {
			return err
		}
	}

	return nil
}

func specRowReader(path string) (*x86csv.Reader, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	return x86csv.NewReader(bufio.NewReader(f)), nil
}