1package ir
2
3import (
4	"errors"
5
6	"github.com/mmcloughlin/avo/attr"
7	"github.com/mmcloughlin/avo/buildtags"
8	"github.com/mmcloughlin/avo/gotypes"
9	"github.com/mmcloughlin/avo/operand"
10	"github.com/mmcloughlin/avo/reg"
11)
12
13// Node is a part of a Function.
14type Node interface {
15	node()
16}
17
18// Label within a function.
19type Label string
20
21func (l Label) node() {}
22
23// Comment represents a multi-line comment.
24type Comment struct {
25	Lines []string
26}
27
28func (c *Comment) node() {}
29
30// NewComment builds a Comment consisting of the provided lines.
31func NewComment(lines ...string) *Comment {
32	return &Comment{
33		Lines: lines,
34	}
35}
36
37// Instruction is a single instruction in a function.
38type Instruction struct {
39	Opcode   string
40	Operands []operand.Op
41
42	Inputs  []operand.Op
43	Outputs []operand.Op
44
45	IsTerminal       bool
46	IsBranch         bool
47	IsConditional    bool
48	CancellingInputs bool
49
50	// ISA is the list of required instruction set extensions.
51	ISA []string
52
53	// CFG.
54	Pred []*Instruction
55	Succ []*Instruction
56
57	// LiveIn/LiveOut are sets of live register IDs pre/post execution.
58	LiveIn  reg.MaskSet
59	LiveOut reg.MaskSet
60}
61
62func (i *Instruction) node() {}
63
64// IsUnconditionalBranch reports whether i is an unconditional branch.
65func (i Instruction) IsUnconditionalBranch() bool {
66	return i.IsBranch && !i.IsConditional
67}
68
69// TargetLabel returns the label referenced by this instruction. Returns nil if
70// no label is referenced.
71func (i Instruction) TargetLabel() *Label {
72	if !i.IsBranch {
73		return nil
74	}
75	if len(i.Operands) == 0 {
76		return nil
77	}
78	if ref, ok := i.Operands[0].(operand.LabelRef); ok {
79		lbl := Label(ref)
80		return &lbl
81	}
82	return nil
83}
84
85// Registers returns all registers involved in the instruction.
86func (i Instruction) Registers() []reg.Register {
87	var rs []reg.Register
88	for _, op := range i.Operands {
89		rs = append(rs, operand.Registers(op)...)
90	}
91	return rs
92}
93
94// InputRegisters returns all registers read by this instruction.
95func (i Instruction) InputRegisters() []reg.Register {
96	var rs []reg.Register
97	for _, op := range i.Inputs {
98		rs = append(rs, operand.Registers(op)...)
99	}
100	if i.CancellingInputs && rs[0] == rs[1] {
101		rs = []reg.Register{}
102	}
103	for _, op := range i.Outputs {
104		if operand.IsMem(op) {
105			rs = append(rs, operand.Registers(op)...)
106		}
107	}
108	return rs
109}
110
111// OutputRegisters returns all registers written by this instruction.
112func (i Instruction) OutputRegisters() []reg.Register {
113	var rs []reg.Register
114	for _, op := range i.Outputs {
115		if r, ok := op.(reg.Register); ok {
116			rs = append(rs, r)
117		}
118	}
119	return rs
120}
121
122// Section is a part of a file.
123type Section interface {
124	section()
125}
126
127// File represents an assembly file.
128type File struct {
129	Constraints buildtags.Constraints
130	Includes    []string
131	Sections    []Section
132}
133
134// NewFile initializes an empty file.
135func NewFile() *File {
136	return &File{}
137}
138
139// AddSection appends a Section to the file.
140func (f *File) AddSection(s Section) {
141	f.Sections = append(f.Sections, s)
142}
143
144// Functions returns all functions in the file.
145func (f *File) Functions() []*Function {
146	var fns []*Function
147	for _, s := range f.Sections {
148		if fn, ok := s.(*Function); ok {
149			fns = append(fns, fn)
150		}
151	}
152	return fns
153}
154
155// Pragma represents a function compiler directive.
156type Pragma struct {
157	Directive string
158	Arguments []string
159}
160
161// Function represents an assembly function.
162type Function struct {
163	Name       string
164	Attributes attr.Attribute
165	Pragmas    []Pragma
166	Doc        []string
167	Signature  *gotypes.Signature
168	LocalSize  int
169
170	Nodes []Node
171
172	// LabelTarget maps from label name to the following instruction.
173	LabelTarget map[Label]*Instruction
174
175	// Register allocation.
176	Allocation reg.Allocation
177
178	// ISA is the list of required instruction set extensions.
179	ISA []string
180}
181
182func (f *Function) section() {}
183
184// NewFunction builds an empty function of the given name.
185func NewFunction(name string) *Function {
186	return &Function{
187		Name:      name,
188		Signature: gotypes.NewSignatureVoid(),
189	}
190}
191
192// AddPragma adds a pragma to this function.
193func (f *Function) AddPragma(directive string, args ...string) {
194	f.Pragmas = append(f.Pragmas, Pragma{
195		Directive: directive,
196		Arguments: args,
197	})
198}
199
200// SetSignature sets the function signature.
201func (f *Function) SetSignature(s *gotypes.Signature) {
202	f.Signature = s
203}
204
205// AllocLocal allocates size bytes in this function's stack.
206// Returns a reference to the base pointer for the newly allocated region.
207func (f *Function) AllocLocal(size int) operand.Mem {
208	ptr := operand.NewStackAddr(f.LocalSize)
209	f.LocalSize += size
210	return ptr
211}
212
213// AddInstruction appends an instruction to f.
214func (f *Function) AddInstruction(i *Instruction) {
215	f.AddNode(i)
216}
217
218// AddLabel appends a label to f.
219func (f *Function) AddLabel(l Label) {
220	f.AddNode(l)
221}
222
223// AddComment adds comment lines to f.
224func (f *Function) AddComment(lines ...string) {
225	f.AddNode(NewComment(lines...))
226}
227
228// AddNode appends a Node to f.
229func (f *Function) AddNode(n Node) {
230	f.Nodes = append(f.Nodes, n)
231}
232
233// Instructions returns just the list of instruction nodes.
234func (f *Function) Instructions() []*Instruction {
235	var is []*Instruction
236	for _, n := range f.Nodes {
237		i, ok := n.(*Instruction)
238		if ok {
239			is = append(is, i)
240		}
241	}
242	return is
243}
244
245// Labels returns just the list of label nodes.
246func (f *Function) Labels() []Label {
247	var lbls []Label
248	for _, n := range f.Nodes {
249		lbl, ok := n.(Label)
250		if ok {
251			lbls = append(lbls, lbl)
252		}
253	}
254	return lbls
255}
256
257// Stub returns the Go function declaration.
258func (f *Function) Stub() string {
259	return "func " + f.Name + f.Signature.String()
260}
261
262// FrameBytes returns the size of the stack frame in bytes.
263func (f *Function) FrameBytes() int {
264	return f.LocalSize
265}
266
267// ArgumentBytes returns the size of the arguments in bytes.
268func (f *Function) ArgumentBytes() int {
269	return f.Signature.Bytes()
270}
271
272// Datum represents a data element at a particular offset of a data section.
273type Datum struct {
274	Offset int
275	Value  operand.Constant
276}
277
278// NewDatum builds a Datum from the given constant.
279func NewDatum(offset int, v operand.Constant) Datum {
280	return Datum{
281		Offset: offset,
282		Value:  v,
283	}
284}
285
286// Interval returns the range of bytes this datum will occupy within its section.
287func (d Datum) Interval() (int, int) {
288	return d.Offset, d.Offset + d.Value.Bytes()
289}
290
291// Overlaps returns true
292func (d Datum) Overlaps(other Datum) bool {
293	s, e := d.Interval()
294	so, eo := other.Interval()
295	return !(eo <= s || e <= so)
296}
297
298// Global represents a DATA section.
299type Global struct {
300	Symbol     operand.Symbol
301	Attributes attr.Attribute
302	Data       []Datum
303	Size       int
304}
305
306// NewGlobal constructs an empty DATA section.
307func NewGlobal(sym operand.Symbol) *Global {
308	return &Global{
309		Symbol: sym,
310	}
311}
312
313// NewStaticGlobal is a convenience for building a static DATA section.
314func NewStaticGlobal(name string) *Global {
315	return NewGlobal(operand.NewStaticSymbol(name))
316}
317
318func (g *Global) section() {}
319
320// Base returns a pointer to the start of the data section.
321func (g *Global) Base() operand.Mem {
322	return operand.NewDataAddr(g.Symbol, 0)
323}
324
325// Grow ensures that the data section has at least the given size.
326func (g *Global) Grow(size int) {
327	if g.Size < size {
328		g.Size = size
329	}
330}
331
332// AddDatum adds d to this data section, growing it if necessary. Errors if the datum overlaps with existing data.
333func (g *Global) AddDatum(d Datum) error {
334	for _, other := range g.Data {
335		if d.Overlaps(other) {
336			return errors.New("overlaps existing datum")
337		}
338	}
339	g.add(d)
340	return nil
341}
342
343// Append the constant to the end of the data section.
344func (g *Global) Append(v operand.Constant) {
345	g.add(Datum{
346		Offset: g.Size,
347		Value:  v,
348	})
349}
350
351func (g *Global) add(d Datum) {
352	_, end := d.Interval()
353	g.Grow(end)
354	g.Data = append(g.Data, d)
355}
356