tt/cmd/cmd.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"robaertschi.xyz/robaertschi/tt/asm/amd64"
	"robaertschi.xyz/robaertschi/tt/lexer"
	"robaertschi.xyz/robaertschi/tt/parser"
	"robaertschi.xyz/robaertschi/tt/token"
	"robaertschi.xyz/robaertschi/tt/ttir"
	"robaertschi.xyz/robaertschi/tt/typechecker"
)

type prefixWriter struct {
	output              io.Writer
	outputPrefix        []byte
	outputPrefixWritten bool
}

func NewPrefixWriter(output io.Writer, prefix []byte) *prefixWriter {
	return &prefixWriter{
		output:       output,
		outputPrefix: prefix,
	}
}

func NewPrefixWriterString(output io.Writer, prefix string) *prefixWriter {
	return &prefixWriter{
		output:       output,
		outputPrefix: []byte(prefix),
	}
}

func (w *prefixWriter) Write(p []byte) (n int, err error) {

	toWrites := bytes.SplitAfter(p, []byte{'\n'})

	for _, toWrite := range toWrites {
		if len(toWrite) <= 0 {
			continue
		}
		if !w.outputPrefixWritten {
			w.outputPrefixWritten = true
			w.output.Write(w.outputPrefix)
		}

		if bytes.Contains(toWrite, []byte{'\n'}) {
			w.outputPrefixWritten = false
		}

		var written int
		written, err = w.output.Write(toWrite)
		n += written
		if err != nil {
			return
		}
	}

	return
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s [flags] input\nPossible flags:\n", os.Args[0])
		flag.PrintDefaults()
	}

	var output string
	flag.StringVar(&output, "o", "", "Output a executable named `file`")
	flag.StringVar(&output, "output", "", "Output a executable named `file`")
	onlyEmitAsm := flag.Bool("S", false, "Only emit the asembly file and exit")
	flag.Parse()

	input := flag.Arg(0)
	if input == "" {
		flag.Usage()
		os.Exit(1)
	}

	if output == "" {
		output = strings.TrimSuffix(input, filepath.Ext(input))
	}
	asmOutputName := strings.TrimSuffix(input, filepath.Ext(input)) + ".asm"

	file, err := os.Open(input)
	if err != nil {
		fmt.Printf("Could not open file %q because: %e", input, err)
		os.Exit(1)
	}
	defer file.Close()

	inputText, err := io.ReadAll(file)
	if err != nil {
		fmt.Printf("Could not read file %q because: %e", input, err)
		os.Exit(1)
	}

	l, err := lexer.New(string(inputText), input)
	if err != nil {
		fmt.Printf("Error while creating lexer: %e", err)
		os.Exit(1)
	}

	l.WithErrorCallback(func(l token.Loc, s string, a ...any) {
		fmt.Printf("%s:%d:%d: %s\n", l.File, l.Line, l.Col, fmt.Sprintf(s, a...))
	})

	p := parser.New(l)
	p.WithErrorCallback(func(t token.Token, s string, a ...any) {
		loc := t.Loc
		fmt.Printf("%s:%d:%d: %s\n", loc.File, loc.Line, loc.Col, fmt.Sprintf(s, a...))
	})

	program := p.ParseProgram()

	if p.Errors() > 0 {
		fmt.Printf("Parser encountered 1 or more errors, quiting...\n")
		os.Exit(1)
	}

	tprogram, err := typechecker.New().CheckProgram(program)
	if err != nil {
		fmt.Printf("Typechecker failed with %e\n", err)
		os.Exit(1)
	}

	ir := ttir.EmitProgram(tprogram)
	asm := amd64.CgProgram(ir)

	asmOutput := asm.Emit()

	asmOutputFile, err := os.Create(asmOutputName)
	if err != nil {
		fmt.Printf("Failed to create/truncate asm file %q because: %v\n", asmOutputName, err)
		os.Exit(1)
	}

	_, err = asmOutputFile.WriteString(asmOutput)
	asmOutputFile.Close()
	if err != nil {
		fmt.Printf("Failed to write to file %q because: %v\n", asmOutputName, err)
		os.Exit(1)
	}

	fasmPath, err := exec.LookPath("fasm")
	if err != nil {
		fasmPath, err = exec.LookPath("fasm2")
		if err != nil {
			fmt.Printf("Could not find fasm or fasm2, please install any those two using your systems package manager or from https://flatassembler.net\n")
			os.Exit(1)
		}
	}

	if !*onlyEmitAsm {
		args := []string{asmOutputName, output}
		cmd := exec.Command(fasmPath, args...)
		cmd.Stdout = NewPrefixWriterString(os.Stdout, "fasm output: ")
		cmd.Stderr = cmd.Stdout

		err = cmd.Run()
		if err != nil {
			fmt.Printf("Failed to run fasm because: %v\nCheck the asm file %q for errors and report these to the author!\n", err, asmOutputName)
			os.Exit(1)
		}

		removeErr := os.Remove(asmOutputName)
		if removeErr != nil {
			fmt.Printf("Failed to remove %q file, please remove it yourself. Err: %v\n", asmOutputName, err)
		}
	}
}