1515
1616log = getLogger (__name__ )
1717
18- def check_cycle (reg , assignments ):
18+ def check_cycle (reg , assignments , mapping = None ):
1919 """Walk down the assignment list of a register,
2020 return the path walked if it is encountered again.
2121
@@ -37,31 +37,38 @@ def check_cycle(reg, assignments):
3737 >>> check_cycle('a', {'a': 'b', 'b': 'c', 'c': 'd', 'd': 'a'})
3838 ['a', 'b', 'c', 'd']
3939 """
40- return check_cycle_ (reg , assignments , [])
41-
42- def check_cycle_ (reg , assignments , path ):
43- target = assignments [reg ]
40+ if mapping is None :
41+ mapping = {}
42+ return check_cycle_ (reg , assignments , [], mapping )
43+
44+ def check_cycle_ (reg , assignments , path , mapping ):
45+ target = assignments .get (reg )
46+ if target is None :
47+ real_reg = mapping .get (reg , reg )
48+ reg , target = next ((k , v ) for k ,v in assignments .items () if mapping .get (k , k ) == real_reg )
4449 path .append (reg )
4550
51+ real_target = mapping .get (target , target )
52+
4653 # No cycle, some other value (e.g. 1)
47- if target not in assignments :
54+ if all ( real_target != mapping . get ( k , k ) for k in assignments ) :
4855 return []
4956
5057 # Found a cycle
51- if target in path :
58+ if any ( real_target == mapping . get ( k , k ) for k in path ) :
5259 # Does the cycle *start* with target?
5360 # This determines whether the original register is
5461 # in the cycle, or just depends on registers in one.
55- if target == path [0 ]:
62+ if real_target == mapping . get ( path [0 ], path [ 0 ]) :
5663 return path
5764
5865 # Just depends on one.
5966 return []
6067
6168 # Recurse
62- return check_cycle_ (target , assignments , path )
69+ return check_cycle_ (target , assignments , path , mapping )
6370
64- def extract_dependencies (reg , assignments ):
71+ def extract_dependencies (reg , assignments , mapping = None ):
6572 """Return a list of all registers which directly
6673 depend on the specified register.
6774
@@ -77,7 +84,9 @@ def extract_dependencies(reg, assignments):
7784 ['b', 'c']
7885 """
7986 # sorted() is only for determinism
80- return sorted ([k for k ,v in assignments .items () if v == reg ])
87+ if mapping is None :
88+ mapping = {}
89+ return sorted ([k for k ,v in assignments .items () if mapping .get (v , v ) == mapping .get (reg , reg )])
8190
8291
8392def resolve_order (reg , deps ):
@@ -110,7 +119,7 @@ def depends_on_cycle(reg, assignments, in_cycles):
110119 reg = assignments .get (reg , None )
111120 return False
112121
113- def regsort (in_out , all_regs , tmp = None , xchg = True , randomize = None ):
122+ def regsort (in_out , all_regs , mapping = None , tmp = None , xchg = True , randomize = None ):
114123 """
115124 Sorts register dependencies.
116125
@@ -233,6 +242,12 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
233242 if randomize is None :
234243 randomize = context .randomize
235244
245+ if mapping is None :
246+ if hasattr (all_regs , 'keys' ):
247+ mapping = all_regs
248+ else :
249+ mapping = {}
250+
236251 sentinel = object ()
237252
238253 # Drop all registers which will be set to themselves.
@@ -242,7 +257,7 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
242257
243258 # Collapse constant values
244259 #
245- # For eaxmple , {'eax': 0, 'ebx': 0} => {'eax': 0, 'ebx': 'eax'}
260+ # For example , {'eax': 0, 'ebx': 0} => {'eax': 0, 'ebx': 'eax'}
246261 v_k = defaultdict (list )
247262 for k ,v in sorted (in_out .items ()):
248263 if v not in all_regs and v != 0 :
@@ -263,7 +278,9 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
263278 # which are also 'outputs'.
264279 #
265280 # For example, {'eax': 1, 'ebx': 2, 'ecx': 'edx'}
266- if not any (v in in_out for k ,v in in_out .items ()):
281+ inps = {mapping .get (k , k ) for k in in_out }
282+ outs = {mapping .get (v , v ) for v in in_out .values ()}
283+ if not inps & outs :
267284 result = [('mov' , k ,in_out [k ]) for k in sorted (in_out )]
268285
269286 if randomize :
@@ -280,7 +297,7 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
280297 # Output: {'A': [], 'B': ['A', 'C'], 'C': []}
281298 #
282299 # In this case, both A and C must be set before B.
283- deps = {r : extract_dependencies (r , in_out ) for r in in_out }
300+ deps = {r : extract_dependencies (r , in_out , mapping ) for r in in_out }
284301
285302 # Final result which will be returned
286303 result = []
@@ -299,7 +316,7 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
299316
300317 while cycle_candidates :
301318 reg = cycle_candidates [0 ]
302- cycle = check_cycle (reg , in_out )
319+ cycle = check_cycle (reg , in_out , mapping )
303320
304321 if cycle :
305322 if randomize :
@@ -395,6 +412,10 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
395412
396413 if tmp :
397414 for cycle in cycles :
415+ if len (cycle ) == 1 :
416+ [reg ] = cycle
417+ result .append (('mov' , reg , in_out [reg ]))
418+ continue
398419
399420 first = cycle [0 ]
400421 last = cycle [- 1 ]
@@ -411,6 +432,9 @@ def regsort(in_out, all_regs, tmp = None, xchg = True, randomize = None):
411432 else :
412433 for cycle in cycles :
413434 size = len (cycle )
435+ if size == 1 :
436+ [reg ] = cycle
437+ result .append (('mov' , reg , in_out [reg ]))
414438 for i in range (size - 1 ):
415439 result .append (('xchg' , cycle [i ], cycle [(i + 1 ) % size ]))
416440
0 commit comments