@@ -58,11 +58,6 @@ pub fn RegisterManager(
5858 return @alignCast (@fieldParentPtr ("register_manager" , self ));
5959 }
6060
61- fn excludeRegister (reg : Register , register_class : RegisterBitSet ) bool {
62- const index = indexOfRegIntoTracked (reg ) orelse return true ;
63- return ! register_class .isSet (index );
64- }
65-
6661 fn markRegIndexAllocated (self : * Self , tracked_index : TrackedIndex ) void {
6762 self .allocated_registers .set (tracked_index );
6863 }
@@ -234,28 +229,20 @@ pub fn RegisterManager(
234229 ) ? [count ]Register {
235230 comptime assert (count > 0 and count <= tracked_registers .len );
236231
237- var free_and_not_locked_registers = self .free_registers ;
238- free_and_not_locked_registers .setIntersection (register_class );
239-
240- var unlocked_registers = self .locked_registers ;
241- unlocked_registers .toggleAll ();
242-
243- free_and_not_locked_registers .setIntersection (unlocked_registers );
244-
245- if (free_and_not_locked_registers .count () < count ) return null ;
232+ var free_and_unlocked_registers = self .locked_registers ;
233+ free_and_unlocked_registers .toggleAll ();
234+ free_and_unlocked_registers .setIntersection (self .free_registers );
235+ free_and_unlocked_registers .setIntersection (register_class );
246236
247237 var regs : [count ]Register = undefined ;
248238 var i : usize = 0 ;
249- for (tracked_registers ) | reg | {
250- if (i >= count ) break ;
251- if (excludeRegister (reg , register_class )) continue ;
252- if (self .isRegLocked (reg )) continue ;
253- if (! self .isRegFree (reg )) continue ;
254-
255- regs [i ] = reg ;
239+ var it = free_and_unlocked_registers .iterator (.{});
240+ while (it .next ()) | reg_index | {
241+ regs [i ] = regAtTrackedIndex (@intCast (reg_index ));
256242 i += 1 ;
243+ if (i >= count ) break ;
257244 }
258- assert (i == count );
245+ if (i < count ) return null ;
259246
260247 for (regs , insts ) | reg , inst | {
261248 log .debug ("tryAllocReg {} for inst {?}" , .{ reg , inst });
@@ -290,46 +277,27 @@ pub fn RegisterManager(
290277 ) AllocationError ! [count ]Register {
291278 comptime assert (count > 0 and count <= tracked_registers .len );
292279
293- var locked_registers = self .locked_registers ;
294- locked_registers .setIntersection (register_class );
295-
296- if (count > register_class .count () - locked_registers .count ()) return error .OutOfRegisters ;
297-
298280 const result = self .tryAllocRegs (count , insts , register_class ) orelse blk : {
281+ var unlocked_registers = self .locked_registers ;
282+ unlocked_registers .toggleAll ();
283+ unlocked_registers .setIntersection (register_class );
284+
299285 // We'll take over the first count registers. Spill
300286 // the instructions that were previously there to a
301287 // stack allocations.
302288 var regs : [count ]Register = undefined ;
303289 var i : usize = 0 ;
304- for (tracked_registers ) | reg | {
305- if (i >= count ) break ;
306- if (excludeRegister (reg , register_class )) break ;
307- if (self .isRegLocked (reg )) continue ;
308-
309- log .debug ("allocReg {} for inst {?}" , .{ reg , insts [i ] });
310- regs [i ] = reg ;
311- self .markRegAllocated (reg );
312- const index = indexOfRegIntoTracked (reg ).? ; // indexOfReg() on a callee-preserved reg should never return null
313- if (insts [i ]) | inst | {
314- // Track the register
315- if (self .isRegFree (reg )) {
316- self .markRegUsed (reg );
317- } else {
318- const spilled_inst = self .registers [index ];
319- try self .getFunction ().spillInstruction (reg , spilled_inst );
320- }
321- self .registers [index ] = inst ;
322- } else {
323- // Don't track the register
324- if (! self .isRegFree (reg )) {
325- const spilled_inst = self .registers [index ];
326- try self .getFunction ().spillInstruction (reg , spilled_inst );
327- self .freeReg (reg );
328- }
329- }
330-
290+ var it = unlocked_registers .iterator (.{});
291+ while (it .next ()) | reg_index | {
292+ const tracked_index : TrackedIndex = @intCast (reg_index );
293+ if (! self .isRegIndexFree (tracked_index ) and
294+ self .registers [tracked_index ].unwrap () == .target ) continue ;
295+ try self .getRegIndex (tracked_index , insts [i ]);
296+ regs [i ] = regAtTrackedIndex (tracked_index );
331297 i += 1 ;
298+ if (i >= count ) break ;
332299 }
300+ if (i < count ) return error .OutOfRegisters ;
333301
334302 break :blk regs ;
335303 };
@@ -351,7 +319,7 @@ pub fn RegisterManager(
351319 /// Spills the register if it is currently allocated. If a
352320 /// corresponding instruction is passed, will also track this
353321 /// register.
354- fn getRegIndex (
322+ pub fn getRegIndex (
355323 self : * Self ,
356324 tracked_index : TrackedIndex ,
357325 inst : ? Air.Inst.Index ,
0 commit comments