ir_rewrite.mx raw
1 package main
2
3 import (
4 "git.smesh.lol/moxie/pkg/mxutil"
5 )
6
7 // joinChunks concatenates a slice of byte slices in one allocation.
8 // Uses byte-by-byte copy because s[off:] creates a copy under
9 // subslice-copies semantics, so copy(s[off:], src) writes to a
10 // temporary instead of s.
11 func joinChunks(chunks [][]byte) (r []byte) {
12 n := int32(0)
13 for _, c := range chunks {
14 n += len(c)
15 }
16 out := []byte{:n}
17 off := int32(0)
18 for _, c := range chunks {
19 for j := int32(0); j < len(c); j++ {
20 out[off+j] = c[j]
21 }
22 off += len(c)
23 }
24 return out
25 }
26
27 // commaSep interleaves ", " between segments.
28 func commaSep(segs []string) (r []string) {
29 var out []string
30 for i, s := range segs {
31 if i > 0 {
32 push(out, ", ")
33 }
34 push(out, s)
35 }
36 return out
37 }
38
39 // joinStrs concatenates string segments in one allocation.
40 // Replaces Schlemiel x = x | y loops.
41 func joinStrs(segs []string) (r string) {
42 n := int32(0)
43 for _, s := range segs {
44 n += len(s)
45 }
46 out := []byte{:n}
47 off := int32(0)
48 for _, s := range segs {
49 for j := int32(0); j < len(s); j++ {
50 out[off+j] = s[j]
51 }
52 off += len(s)
53 }
54 return string(out)
55 }
56
57 // srcTrimStart returns index of first non-whitespace byte in src[from:to].
58 func srcTrimStart(src []byte, from, to int32) (s int32) {
59 s = from
60 for s < to && (src[s] == ' ' || src[s] == '\t' || src[s] == '\r') {
61 s++
62 }
63 return s
64 }
65
66 // srcTrimEnd returns index one past last non-whitespace byte in src[from:to].
67 func srcTrimEnd(src []byte, from, to int32) (e int32) {
68 e = to
69 for e > from && (src[e-1] == ' ' || src[e-1] == '\t' || src[e-1] == '\r') {
70 e--
71 }
72 return e
73 }
74
75 // srcHasPrefix checks whether src[start:end] begins with prefix.
76 func srcHasPrefix(src []byte, start, end int32, prefix []byte) (ok bool) {
77 plen := int32(len(prefix))
78 if end-start < plen {
79 return false
80 }
81 for i := int32(0); i < plen; i++ {
82 if src[start+i] != prefix[i] {
83 return false
84 }
85 }
86 return true
87 }
88
89 // srcFindByte finds byte c in src[from:to], returns index or -1.
90 func srcFindByte(src []byte, from, to int32, c byte) (idx int32) {
91 for i := from; i < to; i++ {
92 if src[i] == c {
93 return i
94 }
95 }
96 return -1
97 }
98
99 func scanEmbedPragmas(src []byte, pkgDir string) {
100 cctx.embedVars = map[string]*embedInfo{}
101 embedPrefix := []byte("//moxie:embed_le ")
102 embedPrefix2 := []byte("//:embed ")
103 varPrefix := []byte("var ")
104 commentPrefix := []byte("//")
105 pendingPath := ""
106 slen := int32(len(src))
107 i := int32(0)
108 for i < slen {
109 // Find line boundaries without subslicing.
110 le := srcFindByte(src, i, slen, '\n')
111 if le < 0 {
112 le = slen
113 }
114 ts := srcTrimStart(src, i, le)
115 te := srcTrimEnd(src, i, le)
116 tlen := te - ts
117 if len(pendingPath) > 0 {
118 if tlen == 0 || srcHasPrefix(src, ts, te, commentPrefix) {
119 if le < slen {
120 i = le + 1
121 } else {
122 i = slen
123 }
124 continue
125 }
126 if srcHasPrefix(src, ts, te, varPrefix) {
127 rs := ts + 4
128 sp := srcFindByte(src, rs, te, ' ')
129 if sp > rs {
130 varName := string(src[rs:sp])
131 vs := srcTrimStart(src, sp+1, te)
132 ve := srcTrimEnd(src, vs, te)
133 typeStr := string(src[vs:ve])
134 cctx.embedVars[varName] = &embedInfo{path: pendingPath, elemType: typeStr}
135 }
136 }
137 pendingPath = ""
138 } else if srcHasPrefix(src, ts, te, embedPrefix) {
139 ps := srcTrimStart(src, ts+int32(len(embedPrefix)), te)
140 pe := srcTrimEnd(src, ps, te)
141 if pe > ps {
142 pendingPath = string(src[ps:pe])
143 if pkgDir != "" {
144 pendingPath = mxutil.JoinPath(pkgDir, pendingPath)
145 }
146 }
147 } else if srcHasPrefix(src, ts, te, embedPrefix2) {
148 ps := srcTrimStart(src, ts+int32(len(embedPrefix2)), te)
149 pe := srcTrimEnd(src, ps, te)
150 if pe > ps {
151 pendingPath = string(src[ps:pe])
152 if pkgDir != "" {
153 pendingPath = mxutil.JoinPath(pkgDir, pendingPath)
154 }
155 }
156 }
157 if le < slen {
158 i = le + 1
159 } else {
160 i = slen
161 }
162 }
163 }
164
165 func scanExportPragmas(src []byte) (m map[string]string) {
166 result := map[string]string{}
167 exportPrefix := []byte("//export ")
168 funcPrefix := []byte("func ")
169 commentPrefix := []byte("//")
170 pendingExport := ""
171 slen := int32(len(src))
172 i := int32(0)
173 for i < slen {
174 le := srcFindByte(src, i, slen, '\n')
175 if le < 0 {
176 le = slen
177 }
178 ts := srcTrimStart(src, i, le)
179 te := srcTrimEnd(src, i, le)
180 tlen := te - ts
181 if len(pendingExport) > 0 {
182 if tlen == 0 || srcHasPrefix(src, ts, te, commentPrefix) {
183 if le < slen { i = le + 1 } else { i = slen }
184 continue
185 }
186 if srcHasPrefix(src, ts, te, funcPrefix) {
187 rs := ts + 5
188 paren := srcFindByte(src, rs, te, '(')
189 if paren > rs {
190 ns := srcTrimStart(src, rs, paren)
191 ne := srcTrimEnd(src, ns, paren)
192 if ne > ns {
193 result[string(src[ns:ne])] = pendingExport
194 }
195 }
196 }
197 pendingExport = ""
198 } else if srcHasPrefix(src, ts, te, exportPrefix) {
199 ps := srcTrimStart(src, ts+9, te)
200 pe := srcTrimEnd(src, ps, te)
201 if pe > ps {
202 pendingExport = string(src[ps:pe])
203 }
204 }
205 if le < slen { i = le + 1 } else { i = slen }
206 }
207 return result
208 }
209
210 func scanLinknameMap(src []byte) (m map[string]string) {
211 result := map[string]string{}
212 prefix := []byte("//:linkname ")
213 slen := int32(len(src))
214 i := int32(0)
215 for i < slen {
216 le := srcFindByte(src, i, slen, '\n')
217 if le < 0 {
218 le = slen
219 }
220 ts := srcTrimStart(src, i, le)
221 te := srcTrimEnd(src, i, le)
222 if srcHasPrefix(src, ts, te, prefix) {
223 rs := srcTrimStart(src, ts+12, te)
224 re := srcTrimEnd(src, rs, te)
225 if re > rs {
226 // Find space between localName and targetName.
227 sp := srcFindByte(src, rs, re, ' ')
228 if sp < 0 {
229 sp = srcFindByte(src, rs, re, '\t')
230 }
231 if sp > rs {
232 localName := string(src[rs:sp])
233 ts2 := srcTrimStart(src, sp+1, re)
234 if ts2 < re {
235 targetName := string(src[ts2:re])
236 result[localName] = targetName
237 }
238 }
239 }
240 }
241 if le < slen { i = le + 1 } else { i = slen }
242 }
243 return result
244 }
245
246
247 // wasmImportMap maps IR symbol (e.g. "time.timezoneOffsetMinutes") -> [module, import_name].
248 // Populated from //:wasmimport pragmas.
249
250 // scanWasmImportMap scans source for //:wasmimport <module> <name> pragmas.
251 // The pragma applies to the next function declaration (body-less).
252 // pkgPath is the package import path used to construct full IR symbols.
253 func scanWasmImportMap(src []byte, pkgPath string) (m map[string][2]string) {
254 if !mxutil.HasPrefix(cctx.compileTargetTriple, "wasm32") {
255 return nil
256 }
257 result := map[string][2]string{}
258 wasmPrefix := []byte("//:wasmimport ")
259 funcPrefix := []byte("func ")
260 slen := int32(len(src))
261 i := int32(0)
262 for i < slen {
263 le := srcFindByte(src, i, slen, '\n')
264 if le < 0 {
265 le = slen
266 }
267 ts := srcTrimStart(src, i, le)
268 te := srcTrimEnd(src, i, le)
269 if !srcHasPrefix(src, ts, te, wasmPrefix) {
270 if le < slen { i = le + 1 } else { i = slen }
271 continue
272 }
273 // Parse module and import name from pragma line.
274 rs := ts + int32(len(wasmPrefix))
275 sp := srcFindByte(src, rs, te, ' ')
276 if sp < 0 {
277 sp = srcFindByte(src, rs, te, '\t')
278 }
279 if sp < 0 {
280 if le < slen { i = le + 1 } else { i = slen }
281 continue
282 }
283 mod := string(src[rs:sp])
284 ns := srcTrimStart(src, sp+1, te)
285 ne := srcTrimEnd(src, ns, te)
286 impName := string(src[ns:ne])
287 // Find next func declaration.
288 j := le
289 if j < slen { j++ }
290 for j < slen {
291 fle := srcFindByte(src, j, slen, '\n')
292 if fle < 0 {
293 fle = slen
294 }
295 fts := srcTrimStart(src, j, fle)
296 fte := srcTrimEnd(src, j, fle)
297 if !srcHasPrefix(src, fts, fte, funcPrefix) {
298 if fle < slen { j = fle + 1 } else { j = slen }
299 continue
300 }
301 // Extract function name from "func NAME(" or "func (recv) NAME("
302 rs2 := fts + 5
303 if rs2 < fte && src[rs2] == '(' {
304 depth := int32(1)
305 k := rs2 + 1
306 for k < fte && depth > 0 {
307 if src[k] == '(' { depth++ }
308 if src[k] == ')' { depth-- }
309 k++
310 }
311 rs2 = srcTrimStart(src, k, fte)
312 }
313 paren := srcFindByte(src, rs2, fte, '(')
314 if paren < 0 {
315 break
316 }
317 fns := srcTrimStart(src, rs2, paren)
318 fne := srcTrimEnd(src, fns, paren)
319 if fne > fns {
320 funcName := string(src[fns:fne])
321 sym := pkgPath | "." | funcName
322 result[sym] = [2]string{mod, impName}
323 }
324 break
325 }
326 if le < slen { i = le + 1 } else { i = slen }
327 }
328 if len(result) == 0 {
329 return nil
330 }
331 return result
332 }
333
334 func stripDuplicatePackageClauses(src []byte) (buf []byte) {
335 found := false
336 var out []byte
337 i := 0
338 for i < len(src) {
339 nlIdx := lIndexByte(src[i:], '\n')
340 var line []byte
341 var lineEnd int32
342 if nlIdx < 0 {
343 line = src[i:]
344 lineEnd = len(src)
345 } else {
346 line = src[i : i+nlIdx]
347 lineEnd = i + nlIdx + 1
348 }
349 trimmed := lTrimSpace(line)
350 if lHasPrefix(trimmed, "package ") {
351 if found {
352 if out == nil {
353 out = []byte{:len(src)}
354 copy(out, src)
355 }
356 for k := i; k < i+len(line); k++ {
357 out[k] = ' '
358 }
359 i = lineEnd
360 continue
361 }
362 found = true
363 }
364 i = lineEnd
365 }
366 if out == nil {
367 return src
368 }
369 return out
370 }
371
372 func stripDuplicatePackageClauses2(src []byte) (buf []byte) {
373 // Fast path: count package clauses. If only one, return src unchanged.
374 count := int32(0)
375 for i := int32(0); i < len(src); i++ {
376 if i == 0 || src[i-1] == '\n' {
377 j := i
378 for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
379 j++
380 }
381 if len(src)-j >= 8 && src[j] == 'p' && src[j+1] == 'a' && src[j+2] == 'c' && src[j+3] == 'k' && src[j+4] == 'a' && src[j+5] == 'g' && src[j+6] == 'e' && src[j+7] == ' ' {
382 count++
383 }
384 }
385 }
386 if count <= 1 {
387 return src
388 }
389 // Multiple package clauses: copy src, blank out duplicates in-place.
390 out := []byte{:len(src)}
391 copy(out, src)
392 found := false
393 for i := int32(0); i < len(out); i++ {
394 if i == 0 || out[i-1] == '\n' {
395 ls := i
396 for ls < len(out) && (out[ls] == ' ' || out[ls] == '\t') {
397 ls++
398 }
399 if len(out)-ls >= 8 && out[ls] == 'p' && out[ls+1] == 'a' && out[ls+2] == 'c' && out[ls+3] == 'k' && out[ls+4] == 'a' && out[ls+5] == 'g' && out[ls+6] == 'e' && out[ls+7] == ' ' {
400 if found {
401 // Blank out this line.
402 for k := i; k < len(out) && out[k] != '\n'; k++ {
403 out[k] = ' '
404 }
405 }
406 found = true
407 }
408 }
409 }
410 return out
411 }
412
413 func rewriteSliceMakeLiterals(src []byte) (buf []byte) {
414 var chunks [][]byte
415 i := 0
416 for i < len(src) {
417 start := lIndex(src[i:], []byte("{:"))
418 if start < 0 {
419 push(chunks, src[i:])
420 break
421 }
422 start = start + i
423 lbrack := findSliceTypeStart(src, start)
424 if lbrack < 0 {
425 push(chunks, src[i:start+2])
426 i = start + 2
427 continue
428 }
429 closeIdx := findMatchingBrace(src, start)
430 if closeIdx < 0 {
431 push(chunks, src[i:start+2])
432 i = start + 2
433 continue
434 }
435 inner := src[start+2 : closeIdx]
436 colonIdx := lIndexByte(inner, ':')
437 typeText := src[lbrack:start]
438 push(chunks, src[i:lbrack])
439 if colonIdx < 0 {
440 push(chunks, []byte("__slicealloc("))
441 push(chunks, typeText)
442 push(chunks, []byte(", "))
443 push(chunks, lTrimSpace(inner))
444 push(chunks, []byte(")"))
445 } else {
446 lenExpr := lTrimSpace(inner[:colonIdx])
447 capExpr := lTrimSpace(inner[colonIdx+1:])
448 push(chunks, []byte("__slicealloc("))
449 push(chunks, typeText)
450 push(chunks, []byte(", "))
451 push(chunks, lenExpr)
452 push(chunks, []byte(", "))
453 push(chunks, capExpr)
454 push(chunks, []byte(")"))
455 }
456 i = closeIdx + 1
457 }
458 if len(chunks) == 0 {
459 return src
460 }
461 return joinChunks(chunks)
462 }
463
464 func findSliceTypeStart(src []byte, braceIdx int32) (n int32) {
465 j := braceIdx - 1
466 for j >= 0 && (src[j] == ' ' || src[j] == '\t' || src[j] == '\n') {
467 j--
468 }
469 if j < 0 {
470 return -1
471 }
472 depth := 0
473 parenDepth := 0
474 candidate := -1
475 for j >= 0 {
476 ch := src[j]
477 if ch == ')' {
478 parenDepth++
479 } else if ch == '(' {
480 if parenDepth > 0 {
481 parenDepth--
482 } else {
483 if candidate >= 0 {
484 return candidate
485 }
486 return -1
487 }
488 } else if parenDepth > 0 {
489 j--
490 continue
491 }
492 if ch == ']' {
493 depth++
494 } else if ch == '[' {
495 depth--
496 if depth == 0 {
497 candidate = j
498 }
499 } else if depth == 0 && parenDepth == 0 {
500 if candidate >= 0 {
501 return candidate
502 }
503 if ch == ' ' || ch == '\t' || ch == '\n' || ch == '*' || ch == '(' || ch == ')' {
504 j--
505 continue
506 }
507 if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_' || ch == '.' {
508 j--
509 continue
510 }
511 return -1
512 }
513 j--
514 }
515 if candidate >= 0 {
516 return candidate
517 }
518 return -1
519 }
520
521 func findMatchingBrace(src []byte, openIdx int32) (n int32) {
522 depth := 1
523 for i := openIdx + 1; i < len(src); i++ {
524 if src[i] == '{' {
525 depth++
526 } else if src[i] == '}' {
527 depth--
528 if depth == 0 {
529 return i
530 }
531 }
532 }
533 return -1
534 }
535
536 func rewriteChanMakeLiterals(src []byte) (buf []byte) {
537 chanKw := []byte("chan ")
538 var chunks [][]byte
539 started := false
540 i := 0
541 for i < len(src) {
542 idx := lIndex(src[i:], chanKw)
543 if idx < 0 {
544 if started {
545 push(chunks, src[i:])
546 }
547 break
548 }
549 idx = idx + i
550 j := idx + 5
551 for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
552 j++
553 }
554 if j >= len(src) {
555 if started {
556 push(chunks, src[i:])
557 }
558 break
559 }
560 if src[j] == '<' && j+1 < len(src) && src[j+1] == '-' {
561 if started {
562 push(chunks, src[i:j+2])
563 }
564 i = j + 2
565 continue
566 }
567 for j < len(src) && (src[j] != '{' && src[j] != '\n' && src[j] != ';' && src[j] != '(' && src[j] != ')') {
568 if src[j] == ' ' || src[j] == '\t' {
569 break
570 }
571 j++
572 }
573 if j >= len(src) || src[j] != '{' {
574 if !started {
575 i = idx + 4
576 } else {
577 push(chunks, src[i:idx+4])
578 i = idx + 4
579 }
580 continue
581 }
582 braceOpen := j
583 endsStruct := braceOpen >= 6 && string(src[braceOpen-6:braceOpen]) == "struct"
584 endsIface := braceOpen >= 9 && string(src[braceOpen-9:braceOpen]) == "interface"
585 if endsStruct || endsIface {
586 typeClose := findMatchingBrace(src, braceOpen)
587 if typeClose < 0 {
588 if !started {
589 i = idx + 4
590 } else {
591 push(chunks, src[i:idx+4])
592 i = idx + 4
593 }
594 continue
595 }
596 k := typeClose + 1
597 if k >= len(src) || src[k] != '{' {
598 if !started {
599 i = typeClose + 1
600 } else {
601 push(chunks, src[i:typeClose+1])
602 i = typeClose + 1
603 }
604 continue
605 }
606 braceOpen = k
607 }
608 closeIdx := findMatchingBrace(src, braceOpen)
609 if closeIdx < 0 {
610 if !started {
611 i = idx + 4
612 } else {
613 push(chunks, src[i:idx+4])
614 i = idx + 4
615 }
616 continue
617 }
618 inner := lTrimSpace(src[braceOpen+1 : closeIdx])
619 chanType := src[idx : braceOpen]
620 for len(chanType) > 0 && (chanType[len(chanType)-1] == ' ' || chanType[len(chanType)-1] == '\t') {
621 chanType = chanType[:len(chanType)-1]
622 }
623 if !started {
624 push(chunks, src[:idx])
625 started = true
626 } else {
627 push(chunks, src[i:idx])
628 }
629 if len(inner) == 0 {
630 push(chunks, []byte("__slicealloc("))
631 push(chunks, chanType)
632 push(chunks, []byte(")"))
633 } else {
634 push(chunks, []byte("__slicealloc("))
635 push(chunks, chanType)
636 push(chunks, []byte(", "))
637 push(chunks, inner)
638 push(chunks, []byte(")"))
639 }
640 i = closeIdx + 1
641 }
642 if !started {
643 return src
644 }
645 return joinChunks(chunks)
646 }
647