@@ -96,6 +96,8 @@ def assign(self, dst: StackEffect, src: StackEffect):
9696 cast = self .cast (dst , src )
9797 if m := re .match (r"^PEEK\((\d+)\)$" , dst .name ):
9898 self .emit (f"POKE({ m .group (1 )} , { cast } { src .name } );" )
99+ elif m := re .match (r"^REG\(oparg(\d+)\)$" , dst .name ):
100+ self .emit (f"Py_XSETREF({ dst .name } , { cast } { src .name } );" )
99101 else :
100102 self .emit (f"{ dst .name } = { cast } { src .name } ;" )
101103
@@ -109,6 +111,7 @@ class Instruction:
109111
110112 # Parts of the underlying instruction definition
111113 inst : parser .InstDef
114+ register : bool
112115 kind : typing .Literal ["inst" , "op" ]
113116 name : str
114117 block : parser .Block
@@ -121,13 +124,16 @@ class Instruction:
121124 cache_effects : list [parser .CacheEffect ]
122125 input_effects : list [StackEffect ]
123126 output_effects : list [StackEffect ]
127+ input_registers : list [str ] # Parallel to input_effects
128+ output_registers : list [str ] # Etc.
124129
125130 # Set later
126131 family : parser .Family | None = None
127132 predicted : bool = False
128133
129134 def __init__ (self , inst : parser .InstDef ):
130135 self .inst = inst
136+ self .register = inst .register
131137 self .kind = inst .kind
132138 self .name = inst .name
133139 self .block = inst .block
@@ -142,6 +148,14 @@ def __init__(self, inst: parser.InstDef):
142148 ]
143149 self .output_effects = inst .outputs # For consistency/completeness
144150
151+ def analyze_registers (self , a : "Analyzer" ) -> None :
152+ regs = iter (("REG(oparg1)" , "REG(oparg2)" , "REG(oparg3)" ))
153+ try :
154+ self .input_registers = [next (regs ) for _ in self .input_effects ]
155+ self .output_registers = [next (regs ) for _ in self .output_effects ]
156+ except StopIteration : # Running out of registers
157+ a .error (f"Instruction { self .name } has too many register effects" )
158+
145159 def write (self , out : Formatter ) -> None :
146160 """Write one instruction, sans prologue and epilogue."""
147161 # Write a static assertion that a family's cache size is correct
@@ -153,10 +167,16 @@ def write(self, out: Formatter) -> None:
153167 f'{ self .cache_offset } , "incorrect cache size");'
154168 )
155169
156- # Write input stack effect variable declarations and initializations
157- for i , ieffect in enumerate (reversed (self .input_effects ), 1 ):
158- src = StackEffect (f"PEEK({ i } )" , "" )
159- out .declare (ieffect , src )
170+ if not self .register :
171+ # Write input stack effect variable declarations and initializations
172+ for i , ieffect in enumerate (reversed (self .input_effects ), 1 ):
173+ src = StackEffect (f"PEEK({ i } )" , "" )
174+ out .declare (ieffect , src )
175+ else :
176+ # Write input register variable declarations and initializations
177+ for ieffect , reg in zip (self .input_effects , self .input_registers ):
178+ src = StackEffect (reg , "" )
179+ out .declare (ieffect , src )
160180
161181 # Write output stack effect variable declarations
162182 input_names = {ieffect .name for ieffect in self .input_effects }
@@ -170,18 +190,24 @@ def write(self, out: Formatter) -> None:
170190 if self .always_exits :
171191 return
172192
173- # Write net stack growth/shrinkage
174- diff = len (self .output_effects ) - len (self .input_effects )
175- out .stack_adjust (diff )
176-
177- # Write output stack effect assignments
178- unmoved_names : set [str ] = set ()
179- for ieffect , oeffect in zip (self .input_effects , self .output_effects ):
180- if ieffect .name == oeffect .name :
181- unmoved_names .add (ieffect .name )
182- for i , oeffect in enumerate (reversed (self .output_effects ), 1 ):
183- if oeffect .name not in unmoved_names :
184- dst = StackEffect (f"PEEK({ i } )" , "" )
193+ if not self .register :
194+ # Write net stack growth/shrinkage
195+ diff = len (self .output_effects ) - len (self .input_effects )
196+ out .stack_adjust (diff )
197+
198+ # Write output stack effect assignments
199+ unmoved_names : set [str ] = set ()
200+ for ieffect , oeffect in zip (self .input_effects , self .output_effects ):
201+ if ieffect .name == oeffect .name :
202+ unmoved_names .add (ieffect .name )
203+ for i , oeffect in enumerate (reversed (self .output_effects ), 1 ):
204+ if oeffect .name not in unmoved_names :
205+ dst = StackEffect (f"PEEK({ i } )" , "" )
206+ out .assign (dst , oeffect )
207+ else :
208+ # Write output register assignments
209+ for oeffect , reg in zip (self .output_effects , self .output_registers ):
210+ dst = StackEffect (reg , "" )
185211 out .assign (dst , oeffect )
186212
187213 # Write cache effect
@@ -218,24 +244,28 @@ def write_body(self, out: Formatter, dedent: int, cache_adjust: int = 0) -> None
218244 # ERROR_IF() must pop the inputs from the stack.
219245 # The code block is responsible for DECREF()ing them.
220246 # NOTE: If the label doesn't exist, just add it to ceval.c.
221- ninputs = len (self .input_effects )
222- # Don't pop common input/output effects at the bottom!
223- # These aren't DECREF'ed so they can stay.
224- for ieff , oeff in zip (self .input_effects , self .output_effects ):
225- if ieff .name == oeff .name :
226- ninputs -= 1
227- else :
228- break
247+ if not self .register :
248+ ninputs = len (self .input_effects )
249+ # Don't pop common input/output effects at the bottom!
250+ # These aren't DECREF'ed so they can stay.
251+ for ieff , oeff in zip (self .input_effects , self .output_effects ):
252+ if ieff .name == oeff .name :
253+ ninputs -= 1
254+ else :
255+ break
256+ else :
257+ ninputs = 0
229258 if ninputs :
230259 out .write_raw (
231260 f"{ extra } { space } if ({ cond } ) goto pop_{ ninputs } _{ label } ;\n "
232261 )
233262 else :
234263 out .write_raw (f"{ extra } { space } if ({ cond } ) goto { label } ;\n " )
235264 elif m := re .match (r"(\s*)DECREF_INPUTS\(\);\s*$" , line ):
236- space = m .group (1 )
237- for ieff in self .input_effects :
238- out .write_raw (f"{ extra } { space } Py_DECREF({ ieff .name } );\n " )
265+ if not self .register :
266+ space = m .group (1 )
267+ for ieff in self .input_effects :
268+ out .write_raw (f"{ extra } { space } Py_DECREF({ ieff .name } );\n " )
239269 else :
240270 out .write_raw (extra + line )
241271
@@ -387,6 +417,7 @@ def analyze(self) -> None:
387417 self .find_predictions ()
388418 self .map_families ()
389419 self .check_families ()
420+ self .analyze_register_instrs ()
390421 self .analyze_supers_and_macros ()
391422
392423 def find_predictions (self ) -> None :
@@ -453,6 +484,11 @@ def check_families(self) -> None:
453484 family ,
454485 )
455486
487+ def analyze_register_instrs (self ) -> None :
488+ for instr in self .instrs .values ():
489+ if instr .register :
490+ instr .analyze_registers (self )
491+
456492 def analyze_supers_and_macros (self ) -> None :
457493 """Analyze each super- and macro instruction."""
458494 self .super_instrs = {}
0 commit comments