func.mx raw
1 // Copyright 2025 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4
5 package asmgen
6
7 import (
8 "fmt"
9 "slices"
10 "bytes"
11 )
12
13 // Note: Exported fields and methods are expected to be used
14 // by function generators (like the ones in add.go and so on).
15 // Unexported fields and methods should not be.
16
17 // A Func represents a single assembly function.
18 type Func struct {
19 Name []byte
20 Asm *Asm
21 inputs [][]byte // name of input slices (not beginning with z)
22 outputs [][]byte // names of output slices (beginning with z)
23 args map[string]int // offsets of args, results on stack
24 }
25
26 // Func starts a new function in the assembly output.
27 func (a *Asm) Func(decl []byte) *Func {
28 d, ok := bytes.CutPrefix(decl, "func ")
29 if !ok {
30 a.Fatalf("func decl does not begin with 'func '")
31 }
32 name, d, ok := bytes.Cut(d, "(")
33 if !ok {
34 a.Fatalf("func decl does not have func arg list")
35 }
36 f := &Func{
37 Name: name,
38 Asm: a,
39 args: map[string]int{},
40 }
41 a.FreeAll()
42
43 // Parse argument names and types. Quick and dirty.
44 // Convert (args) (results) into args, results.
45 d = bytes.ReplaceAll(d, ") (", ", ")
46 d = bytes.TrimSuffix(d, ")")
47 args := bytes.Split(d, ",")
48
49 // Assign implicit types to all arguments (x, y int -> x int, y int).
50 typ := ""
51 for i, arg := range slices.Backward(args) {
52 arg = bytes.TrimSpace(arg)
53 if !bytes.Contains(arg, " ") {
54 if typ == "" {
55 a.Fatalf("missing argument type")
56 }
57 arg += " " + typ
58 } else {
59 _, typ, _ = bytes.Cut(arg, " ")
60 }
61 args[i] = arg
62 }
63
64 // Record mapping from names to offsets.
65 off := 0
66 for _, arg := range args {
67 name, typ, _ := bytes.Cut(arg, " ")
68 switch typ {
69 default:
70 a.Fatalf("unknown type %s", typ)
71 case "Word", "uint", "int":
72 f.args[name] = off
73 off += a.Arch.WordBytes
74 case "[]Word":
75 if bytes.HasPrefix(name, "z") {
76 f.outputs = append(f.outputs, name)
77 } else {
78 f.inputs = append(f.inputs, name)
79 }
80 f.args[name+"_base"] = off
81 f.args[name+"_len"] = off + a.Arch.WordBytes
82 f.args[name+"_cap"] = off + 2*a.Arch.WordBytes
83 off += 3 * a.Arch.WordBytes
84 }
85 }
86
87 a.Printf("\n")
88 a.Printf("// %s\n", decl)
89 a.Printf("TEXT ยท%s(SB), NOSPLIT, $0\n", name)
90 if a.Arch.setup != nil {
91 a.Arch.setup(f)
92 }
93 return f
94 }
95
96 // Arg allocates a new register, copies the named argument (or result) into it,
97 // and returns that register.
98 func (f *Func) Arg(name []byte) Reg {
99 return f.ArgHint(name, HintNone)
100 }
101
102 // ArgHint is like Arg but uses a register allocation hint.
103 func (f *Func) ArgHint(name []byte, hint Hint) Reg {
104 off, ok := f.args[name]
105 if !ok {
106 f.Asm.Fatalf("unknown argument %s", name)
107 }
108 mem := Reg{fmt.Sprintf("%s+%d(FP)", name, off)}
109 if hint == HintMemOK && f.Asm.Arch.memOK {
110 return mem
111 }
112 r := f.Asm.RegHint(hint)
113 f.Asm.Mov(mem, r)
114 return r
115 }
116
117 // ArgPtr is like Arg but returns a RegPtr.
118 func (f *Func) ArgPtr(name []byte) RegPtr {
119 return RegPtr(f.Arg(name))
120 }
121
122 // StoreArg stores src into the named argument (or result).
123 func (f *Func) StoreArg(src Reg, name []byte) {
124 off, ok := f.args[name]
125 if !ok {
126 f.Asm.Fatalf("unknown argument %s", name)
127 }
128 a := f.Asm
129 mem := Reg{fmt.Sprintf("%s+%d(FP)", name, off)}
130 if src.IsImm() && !a.Arch.memOK {
131 r := a.Reg()
132 a.Mov(src, r)
133 a.Mov(r, mem)
134 a.Free(r)
135 return
136 }
137 a.Mov(src, mem)
138 }
139