asm.go raw

   1  package main
   2  
   3  import (
   4  	"fmt"
   5  	"os"
   6  	"os/exec"
   7  	"path/filepath"
   8  	"strings"
   9  )
  10  
  11  func extractASM(outdir string) error {
  12  	for _, opt := range optLevels {
  13  		label := optLabel(opt)
  14  		nFuncs, err := emitAndExtractASM(label, outdir)
  15  		if err != nil {
  16  			return fmt.Errorf("asm %s: %w", label, err)
  17  		}
  18  		fmt.Printf("  ASM %s: %d functions\n", label, nFuncs)
  19  	}
  20  	return nil
  21  }
  22  
  23  func emitAndExtractASM(label, outdir string) (int, error) {
  24  	modIR := "module.ll"
  25  	if label != "" {
  26  		modIR = "module." + label + ".ll"
  27  	}
  28  	irPath := filepath.Join(outdir, modIR)
  29  	irData, err := os.ReadFile(irPath)
  30  	if err != nil {
  31  		return 0, fmt.Errorf("read IR: %w", err)
  32  	}
  33  
  34  	cmd := exec.Command("llc", "-filetype=asm", "-o", "-")
  35  	cmd.Stdin = strings.NewReader(string(irData))
  36  	asmData, err := cmd.Output()
  37  	if err != nil {
  38  		if exitErr, ok := err.(*exec.ExitError); ok {
  39  			return 0, fmt.Errorf("llc asm: %s", string(exitErr.Stderr))
  40  		}
  41  		return 0, fmt.Errorf("llc asm: %w", err)
  42  	}
  43  
  44  	asmFilename := "module.s"
  45  	if label != "" {
  46  		asmFilename = "module." + label + ".s"
  47  	}
  48  	if err := os.WriteFile(filepath.Join(outdir, asmFilename), asmData, 0644); err != nil {
  49  		return 0, err
  50  	}
  51  
  52  	return extractFunctionASM(asmData, label, outdir)
  53  }
  54  
  55  func extractFunctionASM(asmData []byte, label, outdir string) (int, error) {
  56  	segDir := filepath.Join(outdir, "asm")
  57  	if err := os.MkdirAll(segDir, 0755); err != nil {
  58  		return 0, err
  59  	}
  60  
  61  	n := 0
  62  	lines := strings.Split(string(asmData), "\n")
  63  	i := 0
  64  	for i < len(lines) {
  65  		line := lines[i]
  66  		name := extractASMFuncLabel(line)
  67  		if name == "" {
  68  			i++
  69  			continue
  70  		}
  71  
  72  		// Backtrack to include preamble directives (.hidden, .globl, .p2align, .type).
  73  		start := i
  74  		for start > 0 {
  75  			prev := lines[start-1]
  76  			if strings.HasPrefix(prev, "\t.hidden") || strings.HasPrefix(prev, "\t.globl") ||
  77  				strings.HasPrefix(prev, "\t.p2align") || strings.HasPrefix(prev, "\t.type") {
  78  				start--
  79  				continue
  80  			}
  81  			break
  82  		}
  83  
  84  		// Scan forward to "# -- End function" or next function preamble.
  85  		i++
  86  		for i < len(lines) {
  87  			if strings.Contains(lines[i], "# -- End function") {
  88  				i++
  89  				break
  90  			}
  91  			i++
  92  		}
  93  
  94  		funcASM := strings.Join(lines[start:i], "\n") + "\n"
  95  		safeName := sanitizeIRName(name)
  96  		filename := safeName + ".s"
  97  		if label != "" {
  98  			filename = safeName + "." + label + ".s"
  99  		}
 100  
 101  		if err := os.WriteFile(filepath.Join(segDir, filename), []byte(funcASM), 0644); err != nil {
 102  			return 0, err
 103  		}
 104  		n++
 105  	}
 106  	return n, nil
 107  }
 108  
 109  func extractASMFuncLabel(line string) string {
 110  	if len(line) == 0 || line[0] == '.' || line[0] == '\t' || line[0] == ' ' || line[0] == '#' {
 111  		return ""
 112  	}
 113  	// Quoted labels: "pkg/path.Func":  # @"pkg/path.Func"
 114  	if line[0] == '"' {
 115  		colonIdx := strings.Index(line, "\":")
 116  		if colonIdx < 1 {
 117  			return ""
 118  		}
 119  		return line[1:colonIdx]
 120  	}
 121  	if !strings.HasSuffix(line, ":") {
 122  		return ""
 123  	}
 124  	name := strings.TrimSuffix(line, ":")
 125  	if strings.HasPrefix(name, ".") {
 126  		return ""
 127  	}
 128  	return name
 129  }
 130