1 package main
2
3 import (
4 "encoding/json"
5 "fmt"
6 "go/types"
7 "os"
8 "os/exec"
9 "path/filepath"
10 "runtime"
11 "strings"
12 "time"
13
14 "moxie/compiler"
15 "moxie/compileopts"
16 "moxie/goenv"
17 "moxie/loader"
18
19 llvm "tinygo.org/x/go-llvm"
20 )
21
22 var optLevels = []string{"0", "1", "2", "s"}
23
24 // extractIR compiles a package at all four optimization levels and extracts
25 // per-function IR. Uses the compiler directly — no linking, no DCE.
26 func extractIR(pkgName, outdir string) error {
27 wrapDir, err := os.MkdirTemp("", "mxcorpus-wrap-*")
28 if err != nil {
29 return err
30 }
31 defer os.RemoveAll(wrapDir)
32
33 wrapMain := filepath.Join(wrapDir, "main.mx")
34 wrapMod := filepath.Join(wrapDir, "moxie.mod")
35
36 mainSrc := fmt.Sprintf("package main\n\nimport _ %q\n\nfunc main() {}\n", pkgName)
37 if err := os.WriteFile(wrapMain, []byte(mainSrc), 0644); err != nil {
38 return err
39 }
40
41 var modSrc string
42 if flagModDir != "" {
43 modFile := filepath.Join(flagModDir, "moxie.mod")
44 modData, err := os.ReadFile(modFile)
45 if err != nil {
46 return fmt.Errorf("read %s: %w", modFile, err)
47 }
48 var modName string
49 for _, line := range strings.Split(string(modData), "\n") {
50 if strings.HasPrefix(line, "module ") {
51 modName = strings.TrimSpace(line[7:])
52 break
53 }
54 }
55 if modName == "" {
56 return fmt.Errorf("no module name in %s", modFile)
57 }
58 modSrc = fmt.Sprintf("module mxcorpus_wrap\n\ngo 1.25.0\n\nreplace %s => %s\n", modName, flagModDir)
59 } else {
60 moxieRoot := goenv.Get("MOXIEROOT")
61 modSrc = fmt.Sprintf("module mxcorpus_wrap\n\ngo 1.25.0\n\nreplace %s => %s/src/%s\n", pkgName, moxieRoot, pkgName)
62 }
63 if err := os.WriteFile(wrapMod, []byte(modSrc), 0644); err != nil {
64 return err
65 }
66
67 baseIR, err := compilePackageIR(wrapDir, pkgName)
68 if err != nil {
69 return fmt.Errorf("compile: %w", err)
70 }
71
72 symbols := map[string]map[string]string{}
73 for _, opt := range optLevels {
74 label := optLabel(opt)
75 nFuncs, funcs, err := compileAndExtractIR(opt, label, outdir, baseIR)
76 if err != nil {
77 return fmt.Errorf("ir opt=%s: %w", opt, err)
78 }
79 fmt.Printf(" IR %s: %d functions\n", label, nFuncs)
80 m := map[string]string{}
81 for _, name := range funcs {
82 m[name] = sanitizeIRName(name)
83 }
84 symbols[label] = m
85 }
86
87 symData, err := json.MarshalIndent(symbols, "", " ")
88 if err != nil {
89 return fmt.Errorf("symbols json: %w", err)
90 }
91 if err := os.WriteFile(filepath.Join(outdir, "symbols.json"), symData, 0644); err != nil {
92 return err
93 }
94 fmt.Printf(" SYMBOLS: %d opt levels\n", len(symbols))
95 return nil
96 }
97
98 // compilePackageIR loads the program, compiles the target package, and returns
99 // unoptimized IR. Called once; the result is then optimized at each opt level.
100 func compilePackageIR(wrapDir, targetPkg string) ([]byte, error) {
101 config, err := makeConfig("0", wrapDir)
102 if err != nil {
103 return nil, err
104 }
105
106 compilerConfig := &compiler.Config{
107 Triple: config.Triple(),
108 CPU: config.CPU(),
109 Features: config.Features(),
110 ABI: config.ABI(),
111 GOOS: config.GOOS(),
112 GOARCH: config.GOARCH(),
113 BuildMode: config.BuildMode(),
114 CodeModel: config.CodeModel(),
115 RelocationModel: config.RelocationModel(),
116 MoxieVersion: goenv.Version(),
117 Scheduler: config.Scheduler(),
118 DefaultStackSize: config.StackSize(),
119 MaxStackAlloc: config.MaxStackAlloc(),
120 Debug: true,
121 PanicStrategy: config.PanicStrategy(),
122 }
123
124 machine, err := compiler.NewTargetMachine(compilerConfig)
125 if err != nil {
126 return nil, err
127 }
128 defer machine.Dispose()
129
130 lprogram, err := loader.Load(config, ".", types.Config{
131 Sizes: compiler.Sizes(machine),
132 })
133 if err != nil {
134 return nil, fmt.Errorf("load: %w", err)
135 }
136 if err := lprogram.Parse(); err != nil {
137 return nil, fmt.Errorf("parse: %w", err)
138 }
139
140 program := lprogram.LoadSSA()
141
142 var targetLoaderPkg *loader.Package
143 for _, pkg := range lprogram.Sorted() {
144 if pkg.ImportPath == targetPkg {
145 targetLoaderPkg = pkg
146 break
147 }
148 }
149 if targetLoaderPkg == nil {
150 return nil, fmt.Errorf("package %q not found in program", targetPkg)
151 }
152
153 mod, errs := compiler.CompilePackage(
154 targetLoaderPkg.ImportPath,
155 targetLoaderPkg,
156 program.Package(targetLoaderPkg.Pkg),
157 machine,
158 compilerConfig,
159 false,
160 nil,
161 )
162 if errs != nil {
163 return nil, fmt.Errorf("compile: %v", errs)
164 }
165 defer mod.Dispose()
166
167 if err := llvm.VerifyModule(mod, llvm.PrintMessageAction); err != nil {
168 return nil, fmt.Errorf("verify: %w", err)
169 }
170
171 return []byte(mod.String()), nil
172 }
173
174 func compileAndExtractIR(opt, label, outdir string, baseIR []byte) (int, []string, error) {
175 if err := os.MkdirAll(outdir, 0755); err != nil {
176 return 0, nil, err
177 }
178
179 var irData []byte
180 if opt == "0" {
181 irData = baseIR
182 } else {
183 optFlag := "-O" + opt
184 cmd := exec.Command("opt", "-S", "-passes=default<"+optFlag[1:]+">", "-o", "-")
185 cmd.Stdin = strings.NewReader(string(baseIR))
186 out, err := cmd.Output()
187 if err != nil {
188 if exitErr, ok := err.(*exec.ExitError); ok {
189 return 0, nil, fmt.Errorf("opt %s: %s", optFlag, string(exitErr.Stderr))
190 }
191 return 0, nil, fmt.Errorf("opt %s: %w", optFlag, err)
192 }
193 irData = out
194 }
195
196 modFilename := "module.ll"
197 if label != "" {
198 modFilename = "module." + label + ".ll"
199 }
200 if err := os.WriteFile(filepath.Join(outdir, modFilename), irData, 0644); err != nil {
201 return 0, nil, err
202 }
203
204 return extractFunctionIR(irData, label, outdir)
205 }
206
207 func extractFunctionIR(irData []byte, label, outdir string) (int, []string, error) {
208 segDir := filepath.Join(outdir, "ir")
209 if err := os.MkdirAll(segDir, 0755); err != nil {
210 return 0, nil, err
211 }
212
213 n := 0
214 var names []string
215 lines := strings.Split(string(irData), "\n")
216 i := 0
217 for i < len(lines) {
218 line := lines[i]
219 if !strings.HasPrefix(line, "define ") {
220 i++
221 continue
222 }
223
224 name := extractFuncName(line)
225 start := i
226 braceDepth := 0
227 for i < len(lines) {
228 for _, ch := range lines[i] {
229 if ch == '{' {
230 braceDepth++
231 } else if ch == '}' {
232 braceDepth--
233 }
234 }
235 i++
236 if braceDepth == 0 {
237 break
238 }
239 }
240
241 if name == "" {
242 continue
243 }
244
245 funcIR := strings.Join(lines[start:i], "\n") + "\n"
246 safeName := sanitizeIRName(name)
247 filename := safeName + ".ll"
248 if label != "" {
249 filename = safeName + "." + label + ".ll"
250 }
251
252 if err := os.WriteFile(filepath.Join(segDir, filename), []byte(funcIR), 0644); err != nil {
253 return 0, nil, err
254 }
255 names = append(names, name)
256 n++
257 }
258 return n, names, nil
259 }
260
261 func extractFuncName(line string) string {
262 atIdx := strings.Index(line, "@")
263 if atIdx < 0 {
264 return ""
265 }
266 rest := line[atIdx+1:]
267 // Quoted name: @"(*pkg.T).Method"(args...)
268 if len(rest) > 0 && rest[0] == '"' {
269 endQuote := strings.Index(rest[1:], "\"")
270 if endQuote < 0 {
271 return ""
272 }
273 return rest[1 : endQuote+1]
274 }
275 parenIdx := strings.Index(rest, "(")
276 if parenIdx < 0 {
277 return ""
278 }
279 return rest[:parenIdx]
280 }
281
282 func sanitizeIRName(name string) string {
283 name = strings.ReplaceAll(name, "/", "_")
284 name = strings.ReplaceAll(name, " ", "_")
285 name = strings.ReplaceAll(name, "*", "_")
286 name = strings.ReplaceAll(name, "(", "")
287 name = strings.ReplaceAll(name, ")", "")
288 name = strings.ReplaceAll(name, ".", "_")
289 return name
290 }
291
292 func optLabel(opt string) string {
293 switch opt {
294 case "0":
295 return "O0"
296 case "1":
297 return "O1"
298 case "2":
299 return "O2"
300 case "s":
301 return "Os"
302 }
303 return "O" + opt
304 }
305
306 func makeConfig(opt, dir string) (*compileopts.Config, error) {
307 options := &compileopts.Options{
308 GOOS: goenv.Get("GOOS"),
309 GOARCH: goenv.Get("GOARCH"),
310 Opt: opt,
311 Debug: true,
312 Directory: dir,
313 InterpTimeout: 180 * time.Second,
314 Semaphore: make(chan struct{}, runtime.GOMAXPROCS(0)),
315 }
316
317 spec, err := compileopts.LoadTarget(options)
318 if err != nil {
319 return nil, err
320 }
321
322 _, gorootMinor, err := goenv.GetGorootVersion()
323 if err != nil {
324 return nil, err
325 }
326
327 return &compileopts.Config{
328 Options: options,
329 Target: spec,
330 GoMinorVersion: gorootMinor,
331 }, nil
332 }
333