1package reg
2
3import (
4	"errors"
5	"fmt"
6)
7
8// Kind is a class of registers.
9type Kind uint8
10
11// Index of a register within a kind.
12type Index uint16
13
14// Family is a collection of Physical registers of a common kind.
15type Family struct {
16	Kind      Kind
17	registers []Physical
18}
19
20// define builds a register and adds it to the Family.
21func (f *Family) define(s Spec, idx Index, name string, flags ...Info) Physical {
22	r := newregister(f, s, idx, name, flags...)
23	f.add(r)
24	return r
25}
26
27// add r to the family.
28func (f *Family) add(r Physical) {
29	if r.Kind() != f.Kind {
30		panic("bad kind")
31	}
32	f.registers = append(f.registers, r)
33}
34
35// Virtual returns a virtual register from this family's kind.
36func (f *Family) Virtual(idx Index, s Spec) Virtual {
37	return NewVirtual(idx, f.Kind, s)
38}
39
40// Registers returns the registers in this family.
41func (f *Family) Registers() []Physical {
42	return append([]Physical(nil), f.registers...)
43}
44
45// Lookup returns the register with given physical index and spec. Returns nil if no such register exists.
46func (f *Family) Lookup(idx Index, s Spec) Physical {
47	for _, r := range f.registers {
48		if r.PhysicalIndex() == idx && r.Mask() == s.Mask() {
49			return r
50		}
51	}
52	return nil
53}
54
55// ID is a register identifier.
56type ID uint32
57
58// newid builds a new register ID from the virtual flag v, kind and index.
59func newid(v uint8, kind Kind, idx Index) ID {
60	return ID(v) | (ID(kind) << 8) | (ID(idx) << 16)
61}
62
63// IsVirtual reports whether this is an ID for a virtual register.
64func (id ID) IsVirtual() bool { return (id & 1) == 1 }
65
66// IsPhysical reports whether this is an ID for a physical register.
67func (id ID) IsPhysical() bool { return !id.IsVirtual() }
68
69// Kind extracts the kind from the register ID.
70func (id ID) Kind() Kind { return Kind(id >> 8) }
71
72// Index extracts the index from the register ID.
73func (id ID) Index() Index { return Index(id >> 16) }
74
75// Register represents a virtual or physical register.
76type Register interface {
77	ID() ID
78	Kind() Kind
79	Size() uint
80	Mask() uint16
81	Asm() string
82	as(Spec) Register
83	spec() Spec
84	register()
85}
86
87// Equal reports whether a and b are equal registers.
88func Equal(a, b Register) bool {
89	return (a.ID() == b.ID()) && (a.Mask() == b.Mask())
90}
91
92// Virtual is a register of a given type and size, not yet allocated to a physical register.
93type Virtual interface {
94	VirtualIndex() Index
95	Register
96}
97
98// ToVirtual converts r to Virtual if possible, otherwise returns nil.
99func ToVirtual(r Register) Virtual {
100	if v, ok := r.(Virtual); ok {
101		return v
102	}
103	return nil
104}
105
106type virtual struct {
107	idx  Index
108	kind Kind
109	Spec
110}
111
112// NewVirtual builds a Virtual register.
113func NewVirtual(idx Index, k Kind, s Spec) Virtual {
114	return virtual{
115		idx:  idx,
116		kind: k,
117		Spec: s,
118	}
119}
120
121func (v virtual) ID() ID              { return newid(1, v.kind, v.idx) }
122func (v virtual) VirtualIndex() Index { return v.idx }
123func (v virtual) Kind() Kind          { return v.kind }
124
125func (v virtual) Asm() string {
126	// TODO(mbm): decide on virtual register syntax
127	return fmt.Sprintf("<virtual:%v:%v:%v>", v.idx, v.Kind(), v.Size())
128}
129
130func (v virtual) as(s Spec) Register {
131	return virtual{
132		idx:  v.idx,
133		kind: v.kind,
134		Spec: s,
135	}
136}
137
138func (v virtual) spec() Spec { return v.Spec }
139func (v virtual) register()  {}
140
141// Info is a bitmask of register properties.
142type Info uint8
143
144// Defined register Info flags.
145const (
146	None       Info = 0
147	Restricted Info = 1 << iota
148)
149
150// Physical is a concrete register.
151type Physical interface {
152	PhysicalIndex() Index
153	Info() Info
154	Register
155}
156
157// ToPhysical converts r to Physical if possible, otherwise returns nil.
158func ToPhysical(r Register) Physical {
159	if p, ok := r.(Physical); ok {
160		return p
161	}
162	return nil
163}
164
165// register implements Physical.
166type register struct {
167	family *Family
168	idx    Index
169	name   string
170	info   Info
171	Spec
172}
173
174func newregister(f *Family, s Spec, idx Index, name string, flags ...Info) register {
175	r := register{
176		family: f,
177		idx:    idx,
178		name:   name,
179		info:   None,
180		Spec:   s,
181	}
182	for _, flag := range flags {
183		r.info |= flag
184	}
185	return r
186}
187
188func (r register) ID() ID               { return newid(0, r.Kind(), r.idx) }
189func (r register) PhysicalIndex() Index { return r.idx }
190func (r register) Kind() Kind           { return r.family.Kind }
191func (r register) Asm() string          { return r.name }
192func (r register) Info() Info           { return r.info }
193
194func (r register) as(s Spec) Register {
195	return r.family.Lookup(r.PhysicalIndex(), s)
196}
197
198func (r register) spec() Spec { return r.Spec }
199func (r register) register()  {}
200
201// Spec defines the size of a register as well as the bit ranges it occupies in
202// an underlying physical register.
203type Spec uint16
204
205// Spec values required for x86-64.
206const (
207	S0   Spec = 0x0 // zero value reserved for pseudo registers
208	S8L  Spec = 0x1
209	S8H  Spec = 0x2
210	S8        = S8L
211	S16  Spec = 0x3
212	S32  Spec = 0x7
213	S64  Spec = 0xf
214	S128 Spec = 0x1f
215	S256 Spec = 0x3f
216	S512 Spec = 0x7f
217)
218
219// Mask returns a mask representing which bytes of an underlying register are
220// used by this register. This is almost always the low bytes, except for the
221// case of the high-byte registers. If bit n of the mask is set, this means
222// bytes 2^(n-1) to 2^n-1 are used.
223func (s Spec) Mask() uint16 {
224	return uint16(s)
225}
226
227// Size returns the register width in bytes.
228func (s Spec) Size() uint {
229	x := uint(s)
230	return (x >> 1) + (x & 1)
231}
232
233// LookupPhysical returns the physical register with the given parameters, or nil if not found.
234func LookupPhysical(k Kind, idx Index, s Spec) Physical {
235	f := FamilyOfKind(k)
236	if f == nil {
237		return nil
238	}
239	return f.Lookup(idx, s)
240}
241
242// LookupID returns the physical register with the given id and spec, or nil if not found.
243func LookupID(id ID, s Spec) Physical {
244	if id.IsVirtual() {
245		return nil
246	}
247	return LookupPhysical(id.Kind(), id.Index(), s)
248}
249
250// Allocation records a register allocation.
251type Allocation map[ID]ID
252
253// NewEmptyAllocation builds an empty register allocation.
254func NewEmptyAllocation() Allocation {
255	return Allocation{}
256}
257
258// Merge allocations from b into a. Errors if there is disagreement on a common
259// register.
260func (a Allocation) Merge(b Allocation) error {
261	for id, p := range b {
262		if alt, found := a[id]; found && alt != p {
263			return errors.New("disagreement on overlapping register")
264		}
265		a[id] = p
266	}
267	return nil
268}
269
270// LookupDefault returns the register ID assigned by this allocation, returning
271// id if none is found.
272func (a Allocation) LookupDefault(id ID) ID {
273	if _, found := a[id]; found {
274		return a[id]
275	}
276	return id
277}
278
279// LookupRegister the allocation for register r, or return nil if there is none.
280func (a Allocation) LookupRegister(r Register) Physical {
281	// Return immediately if it is already a physical register.
282	if p := ToPhysical(r); p != nil {
283		return p
284	}
285
286	// Lookup an allocation for this virtual ID.
287	id, found := a[r.ID()]
288	if !found {
289		return nil
290	}
291
292	return LookupID(id, r.spec())
293}
294
295// LookupRegisterDefault returns the register assigned to r, or r itself if there is none.
296func (a Allocation) LookupRegisterDefault(r Register) Register {
297	if r == nil {
298		return nil
299	}
300	if p := a.LookupRegister(r); p != nil {
301		return p
302	}
303	return r
304}
305