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 NeedsStackObjects: config.NeedsStackObjects(),
121 Debug: true,
122 PanicStrategy: config.PanicStrategy(),
123 }
124
125 machine, err := compiler.NewTargetMachine(compilerConfig)
126 if err != nil {
127 return nil, err
128 }
129 defer machine.Dispose()
130
131 lprogram, err := loader.Load(config, ".", types.Config{
132 Sizes: compiler.Sizes(machine),
133 })
134 if err != nil {
135 return nil, fmt.Errorf("load: %w", err)
136 }
137 if err := lprogram.Parse(); err != nil {
138 return nil, fmt.Errorf("parse: %w", err)
139 }
140
141 program := lprogram.LoadSSA()
142
143 var targetLoaderPkg *loader.Package
144 for _, pkg := range lprogram.Sorted() {
145 if pkg.ImportPath == targetPkg {
146 targetLoaderPkg = pkg
147 break
148 }
149 }
150 if targetLoaderPkg == nil {
151 return nil, fmt.Errorf("package %q not found in program", targetPkg)
152 }
153
154 mod, errs := compiler.CompilePackage(
155 targetLoaderPkg.ImportPath,
156 targetLoaderPkg,
157 program.Package(targetLoaderPkg.Pkg),
158 machine,
159 compilerConfig,
160 false,
161 nil,
162 )
163 if errs != nil {
164 return nil, fmt.Errorf("compile: %v", errs)
165 }
166 defer mod.Dispose()
167
168 if err := llvm.VerifyModule(mod, llvm.PrintMessageAction); err != nil {
169 return nil, fmt.Errorf("verify: %w", err)
170 }
171
172 return []byte(mod.String()), nil
173 }
174
175 func compileAndExtractIR(opt, label, outdir string, baseIR []byte) (int, []string, error) {
176 if err := os.MkdirAll(outdir, 0755); err != nil {
177 return 0, nil, err
178 }
179
180 var irData []byte
181 if opt == "0" {
182 irData = baseIR
183 } else {
184 optFlag := "-O" + opt
185 cmd := exec.Command("opt", "-S", "-passes=default<"+optFlag[1:]+">", "-o", "-")
186 cmd.Stdin = strings.NewReader(string(baseIR))
187 out, err := cmd.Output()
188 if err != nil {
189 if exitErr, ok := err.(*exec.ExitError); ok {
190 return 0, nil, fmt.Errorf("opt %s: %s", optFlag, string(exitErr.Stderr))
191 }
192 return 0, nil, fmt.Errorf("opt %s: %w", optFlag, err)
193 }
194 irData = out
195 }
196
197 modFilename := "module.ll"
198 if label != "" {
199 modFilename = "module." + label + ".ll"
200 }
201 if err := os.WriteFile(filepath.Join(outdir, modFilename), irData, 0644); err != nil {
202 return 0, nil, err
203 }
204
205 return extractFunctionIR(irData, label, outdir)
206 }
207
208 func extractFunctionIR(irData []byte, label, outdir string) (int, []string, error) {
209 segDir := filepath.Join(outdir, "ir")
210 if err := os.MkdirAll(segDir, 0755); err != nil {
211 return 0, nil, err
212 }
213
214 n := 0
215 var names []string
216 lines := strings.Split(string(irData), "\n")
217 i := 0
218 for i < len(lines) {
219 line := lines[i]
220 if !strings.HasPrefix(line, "define ") {
221 i++
222 continue
223 }
224
225 name := extractFuncName(line)
226 start := i
227 braceDepth := 0
228 for i < len(lines) {
229 for _, ch := range lines[i] {
230 if ch == '{' {
231 braceDepth++
232 } else if ch == '}' {
233 braceDepth--
234 }
235 }
236 i++
237 if braceDepth == 0 {
238 break
239 }
240 }
241
242 if name == "" {
243 continue
244 }
245
246 funcIR := strings.Join(lines[start:i], "\n") + "\n"
247 safeName := sanitizeIRName(name)
248 filename := safeName + ".ll"
249 if label != "" {
250 filename = safeName + "." + label + ".ll"
251 }
252
253 if err := os.WriteFile(filepath.Join(segDir, filename), []byte(funcIR), 0644); err != nil {
254 return 0, nil, err
255 }
256 names = append(names, name)
257 n++
258 }
259 return n, names, nil
260 }
261
262 func extractFuncName(line string) string {
263 atIdx := strings.Index(line, "@")
264 if atIdx < 0 {
265 return ""
266 }
267 rest := line[atIdx+1:]
268 // Quoted name: @"(*pkg.T).Method"(args...)
269 if len(rest) > 0 && rest[0] == '"' {
270 endQuote := strings.Index(rest[1:], "\"")
271 if endQuote < 0 {
272 return ""
273 }
274 return rest[1 : endQuote+1]
275 }
276 parenIdx := strings.Index(rest, "(")
277 if parenIdx < 0 {
278 return ""
279 }
280 return rest[:parenIdx]
281 }
282
283 func sanitizeIRName(name string) string {
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 name = strings.ReplaceAll(name, ".", "_")
290 return name
291 }
292
293 func optLabel(opt string) string {
294 switch opt {
295 case "0":
296 return "O0"
297 case "1":
298 return "O1"
299 case "2":
300 return "O2"
301 case "s":
302 return "Os"
303 }
304 return "O" + opt
305 }
306
307 func makeConfig(opt, dir string) (*compileopts.Config, error) {
308 options := &compileopts.Options{
309 GOOS: goenv.Get("GOOS"),
310 GOARCH: goenv.Get("GOARCH"),
311 Opt: opt,
312 Debug: true,
313 Directory: dir,
314 InterpTimeout: 180 * time.Second,
315 Semaphore: make(chan struct{}, runtime.GOMAXPROCS(0)),
316 }
317
318 spec, err := compileopts.LoadTarget(options)
319 if err != nil {
320 return nil, err
321 }
322
323 _, gorootMinor, err := goenv.GetGorootVersion()
324 if err != nil {
325 return nil, err
326 }
327
328 return &compileopts.Config{
329 Options: options,
330 Target: spec,
331 GoMinorVersion: gorootMinor,
332 }, nil
333 }
334