main.go

package main

import (
	"bufio"
	"bytes"
	"debug/dwarf"
	"debug/macho"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
	"text/template"
)

var (
	funcmap = template.FuncMap{
		"EmitOpcodes": emitOpcodes,
	}

	functemplate   = template.Must(template.New("functempl").Funcs(funcmap).Parse(funcTempl))
	headertemplate = template.Must(template.New("headertempl").Parse(headerTempl))
)

const (
	m9prefix  = "m9:"
	funcTempl = `TEXT {{.Symbol.Markup}}
{{EmitOpcodes .Symbol .Data}}
`
	headerTempl = `/*
 * Generated by mach9 {{.}}; DO NOT EDIT.
 */

#include "textflag.h"

`
)

type symbol struct {
	// These fields come from the Mach-O symbol table
	Name   string
	Offset int
	Data   []byte
	Type   uint8

	// This comes from the DWARF debugging info
	DeclLine int

	// This comes from the source file
	Markup string
}

type templateData struct {
	Symbols []symbol

	Source []string

	LineEntries []dwarf.LineEntry

	Invocation string
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("mach9: ")

	if len(os.Args) < 2 {
		log.Fatal("No object file provided")
	}

	mf, err := macho.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	defer mf.Close()

	// Find the __text section, this contains the assembly opcodes
	text, textidx := findSectionAndIdx("__text", mf)
	if text == nil {
		log.Fatal("No __text section")
	}

	// We need the symbols table to help identify where each function begins
	// and ends in the assembly
	if mf.Symtab == nil {
		log.Fatal("Missing symbol table, can't do anything")
	}

	// Extract the assembled opcodes
	code := make([]byte, text.Size)
	if n, err := text.ReadAt(code, 0); n < int(text.Size) || err != nil {
		if err == nil {
			log.Fatal("Failed to read all bytes in section")
		}
		log.Fatal(err)
	}

	var symbols []symbol

	for _, sym := range mf.Symtab.Syms {
		// See notes.md: symbols which start with l are private link labels and
		// can be ignored. I'm not confident so documenting here but not
		// implementing.

		// Ignore symbols that aren't in the text section
		if sym.Sect != uint8(textidx+1) {
			continue
		}

		symbols = append(symbols, symbol{
			Name: sym.Name, Offset: int(sym.Value), Type: sym.Type,
		})
	}
	// Add a sentinel that represents the end of the code buffer to simplify the
	// next loop.
	symbols = append(symbols, symbol{Name: "", Offset: len(code)})

	// Sort the symbols in order how they appear in the file
	slices.SortFunc(symbols, func(a, b symbol) int {
		return int(a.Offset - b.Offset)
	})

	for i := 0; i < len(symbols)-1; i++ {
		start := &symbols[i]
		end := &symbols[i+1]

		start.Data = code[start.Offset:end.Offset]
	}

	symbols = slices.DeleteFunc[[]symbol](symbols, func(sym symbol) bool {
		// The bottom bit of Type is set if the symbol is an external symbol,
		// one that can be referenced by the linker and other programs. See
		// page 44 https://github.com/aidansteele/osx-abi-macho-file-format-reference/blob/master/Mach-O_File_Format.pdf
		return sym.Type&1 == 0
	})

	var src []string
	cu, le, dwarfOK := parseDWARF(mf, symbols)
	if dwarfOK {
		if lines, err := readSourceFile(cu); err == nil {
			extractDecl(lines, symbols)
			src = prepareSource(lines)
		}
	}

	// Every symbol requires m9 markup
	invalid := false
	for _, sym := range symbols {
		if sym.Markup == "" {
			fmt.Fprintf(os.Stderr, "Symbol %q missing m9 declaration\n", sym.Name)
			invalid = true
		}
	}
	if invalid {
		os.Exit(1)
	}

	outData := &templateData{
		Symbols:     symbols,
		Source:      src,
		LineEntries: le,
		Invocation:  strings.Join(os.Args[1:], " "),
	}
	generateOutput(os.Stdout, outData)
}

func extractDecl(lines []string, symbols []symbol) {
	for i := range symbols {
		sym := &symbols[i]

		line := lines[sym.DeclLine-1]
		idx := strings.Index(line, m9prefix)
		if idx != -1 {
			idx += len(m9prefix)
			sym.Markup = strings.TrimLeft(line[idx:], " \t")
		}
	}
}

func generateOutput(w io.Writer, data *templateData) error {
	if err := headertemplate.Execute(w, data.Invocation); err != nil {
		return err
	}

	for _, sym := range data.Symbols {
		td := struct {
			Symbol *symbol
			Data   *templateData
		}{&sym, data}
		if err := functemplate.Execute(w, &td); err != nil {
			return err
		}
	}

	return nil
}

// Returns the path to the source file for the compile unit and whether DWARF
// info was available. It also updates the symbols in the symmap with the
// declaration line of the symbol.
func parseDWARF(mf *macho.File, symbols []symbol) (string, []dwarf.LineEntry, bool) {
	dwarfdata, err := mf.DWARF()
	if err != nil {
		return "", nil, false
	}

	var cuPath string
	var lentries []dwarf.LineEntry

	reader := dwarfdata.Reader()
	for {
		entry, err := reader.Next()
		if err != nil {
			return "", nil, false
		}
		if entry == nil {
			break
		}

		switch entry.Tag {
		case dwarf.TagCompileUnit:
			cuName, ok := entry.Val(dwarf.AttrName).(string)
			if !ok {
				continue
			}
			cuCompDir, ok := entry.Val(dwarf.AttrCompDir).(string)
			if !ok {
				continue
			}
			cuPath = filepath.Join(cuCompDir, cuName)

			// Read out line entry table for this CU
			lr, err := dwarfdata.LineReader(entry)
			if err != nil {
				continue
			}
			var lentry dwarf.LineEntry
			for lr.Next(&lentry) != io.EOF {
				if !lentry.EndSequence {
					lentries = append(lentries, lentry)
				}
			}

		case dwarf.TagLabel:
			labelName, ok := entry.Val(dwarf.AttrName).(string)
			if !ok {
				continue
			}

			declLine, ok := entry.Val(dwarf.AttrDeclLine).(int64)
			if !ok {
				continue
			}

			for i := range symbols {
				sym := &symbols[i]
				if sym.Name == labelName {
					sym.DeclLine = int(declLine)
					break
				}
			}
		}
	}

	return cuPath, lentries, true
}

// Builds a string that contains the opcodes as literal bytes in Plan9 assembler format
func emitOpcodes(sym *symbol, data *templateData) string {
	builder := strings.Builder{}

	address := sym.Offset

	// TODO - Simplify this code, can the three steps be combined?
	code := sym.Data
	n := len(code)
	off := 0
	s := n / 4
	if s > 0 {
		for i := 0; i < s; i++ {
			line := findSourceLine(address+off, data)
			opcodes := code[off : off+4]
			builder.WriteString(fmt.Sprintf("\tWORD $%#8x", binary.LittleEndian.Uint32(opcodes)))
			if line != "" {
				builder.WriteString(" // " + line)
			}
			builder.WriteString("\n")

			off += 4
		}
		n -= s * 4
	}
	s = n / 2
	if s > 0 {
		for i := 0; i < s; i++ {
			line := findSourceLine(address+off, data)
			opcodes := code[off : off+2]
			builder.WriteString(fmt.Sprintf("\tWORD $%#4x\n", binary.LittleEndian.Uint16(opcodes)))
			if line != "" {
				builder.WriteString(" // " + line)
			}
			builder.WriteString("\n")

			off += 2
		}
		n -= s * 2
	}

	for i := 0; i < n; i++ {
		line := findSourceLine(address+off, data)
		opcode := code[off]
		builder.WriteString(fmt.Sprintf("\tWORD $%#2x\n", opcode))
		if line != "" {
			builder.WriteString(" // " + line)
		}
		builder.WriteString("\n")

		off++
	}

	return builder.String()
}

func findSourceLine(offset int, data *templateData) string {
	for i := range data.LineEntries {
		if uint64(offset) == data.LineEntries[i].Address {
			return data.Source[data.LineEntries[i].Line]
		}
	}

	return ""
}

func findSectionAndIdx(name string, mf *macho.File) (*macho.Section, int) {
	for i, s := range mf.Sections {
		if s.Name == name {
			return s, i
		}
	}
	return nil, -1
}

// Returns the source file as an array of strings, one for each line of the
// file. The first element of the array is empty to allow convenient 1-based
// indexing.
func readSourceFile(file string) ([]string, error) {
	src, err := os.ReadFile(file)
	if err != nil {
		return nil, err
	}

	lines := make([]string, 1)
	scanner := bufio.NewScanner(bytes.NewBuffer(src))
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	return lines, nil
}

var dupws = regexp.MustCompile(`\s+`)

// This removes comment lines and duplicate whitespace to compact the source
func prepareSource(lines []string) []string {
	out := make([]string, len(lines))

	for i, l := range lines {
		if len(l) > 0 {
			l = deleteAfterComment(l)
			out[i] = strings.TrimPrefix(dupws.ReplaceAllString(l, " "), " ")
		}
	}

	return out
}

func deleteAfterComment(text string) string {
	for _, delim := range []string{"//", ";"} {
		index := strings.Index(text, delim)
		if index != -1 {
			return strings.TrimRight(text[:index], " ")
		}
	}
	return text
}