1import re 2from analyzer import StackItem, Instruction, Uop 3from dataclasses import dataclass 4from cwriter import CWriter 5 6UNUSED = {"unused"} 7 8 9def maybe_parenthesize(sym: str) -> str: 10 """Add parentheses around a string if it contains an operator 11 and is not already parenthesized. 12 13 An exception is made for '*' which is common and harmless 14 in the context where the symbolic size is used. 15 """ 16 if sym.startswith("(") and sym.endswith(")"): 17 return sym 18 if re.match(r"^[\s\w*]+$", sym): 19 return sym 20 else: 21 return f"({sym})" 22 23 24def var_size(var: StackItem) -> str: 25 if var.condition: 26 # Special case simplifications 27 if var.condition == "0": 28 return "0" 29 elif var.condition == "1": 30 return var.size 31 elif var.condition == "oparg & 1" and var.size == "1": 32 return f"({var.condition})" 33 else: 34 return f"(({var.condition}) ? {var.size} : 0)" 35 else: 36 return var.size 37 38 39@dataclass 40class StackOffset: 41 "The stack offset of the virtual base of the stack from the physical stack pointer" 42 43 popped: list[str] 44 pushed: list[str] 45 46 @staticmethod 47 def empty() -> "StackOffset": 48 return StackOffset([], []) 49 50 def copy(self) -> "StackOffset": 51 return StackOffset(self.popped[:], self.pushed[:]) 52 53 def pop(self, item: StackItem) -> None: 54 self.popped.append(var_size(item)) 55 56 def push(self, item: StackItem) -> None: 57 self.pushed.append(var_size(item)) 58 59 def __sub__(self, other: "StackOffset") -> "StackOffset": 60 return StackOffset(self.popped + other.pushed, self.pushed + other.popped) 61 62 def __neg__(self) -> "StackOffset": 63 return StackOffset(self.pushed, self.popped) 64 65 def simplify(self) -> None: 66 "Remove matching values from both the popped and pushed list" 67 if not self.popped or not self.pushed: 68 return 69 # Sort the list so the lexically largest element is last. 70 popped = sorted(self.popped) 71 pushed = sorted(self.pushed) 72 self.popped = [] 73 self.pushed = [] 74 while popped and pushed: 75 pop = popped.pop() 76 push = pushed.pop() 77 if pop == push: 78 pass 79 elif pop > push: 80 # if pop > push, there can be no element in pushed matching pop. 81 self.popped.append(pop) 82 pushed.append(push) 83 else: 84 self.pushed.append(push) 85 popped.append(pop) 86 self.popped.extend(popped) 87 self.pushed.extend(pushed) 88 89 def to_c(self) -> str: 90 self.simplify() 91 int_offset = 0 92 symbol_offset = "" 93 for item in self.popped: 94 try: 95 int_offset -= int(item) 96 except ValueError: 97 symbol_offset += f" - {maybe_parenthesize(item)}" 98 for item in self.pushed: 99 try: 100 int_offset += int(item) 101 except ValueError: 102 symbol_offset += f" + {maybe_parenthesize(item)}" 103 if symbol_offset and not int_offset: 104 res = symbol_offset 105 else: 106 res = f"{int_offset}{symbol_offset}" 107 if res.startswith(" + "): 108 res = res[3:] 109 if res.startswith(" - "): 110 res = "-" + res[3:] 111 return res 112 113 def clear(self) -> None: 114 self.popped = [] 115 self.pushed = [] 116 117 118class SizeMismatch(Exception): 119 pass 120 121 122class Stack: 123 def __init__(self) -> None: 124 self.top_offset = StackOffset.empty() 125 self.base_offset = StackOffset.empty() 126 self.variables: list[StackItem] = [] 127 self.defined: set[str] = set() 128 129 def pop(self, var: StackItem) -> str: 130 self.top_offset.pop(var) 131 indirect = "&" if var.is_array() else "" 132 if self.variables: 133 popped = self.variables.pop() 134 if popped.size != var.size: 135 raise SizeMismatch( 136 f"Size mismatch when popping '{popped.name}' from stack to assign to {var.name}. " 137 f"Expected {var.size} got {popped.size}" 138 ) 139 if popped.name == var.name: 140 return "" 141 elif popped.name in UNUSED: 142 self.defined.add(var.name) 143 return ( 144 f"{var.name} = {indirect}stack_pointer[{self.top_offset.to_c()}];\n" 145 ) 146 elif var.name in UNUSED: 147 return "" 148 else: 149 self.defined.add(var.name) 150 return f"{var.name} = {popped.name};\n" 151 self.base_offset.pop(var) 152 if var.name in UNUSED: 153 return "" 154 else: 155 self.defined.add(var.name) 156 cast = f"({var.type})" if (not indirect and var.type) else "" 157 assign = ( 158 f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}];" 159 ) 160 if var.condition: 161 if var.condition == "1": 162 return f"{assign}\n" 163 elif var.condition == "0": 164 return "" 165 else: 166 return f"if ({var.condition}) {{ {assign} }}\n" 167 return f"{assign}\n" 168 169 def push(self, var: StackItem) -> str: 170 self.variables.append(var) 171 if var.is_array() and var.name not in self.defined and var.name not in UNUSED: 172 c_offset = self.top_offset.to_c() 173 self.top_offset.push(var) 174 self.defined.add(var.name) 175 return f"{var.name} = &stack_pointer[{c_offset}];\n" 176 else: 177 self.top_offset.push(var) 178 return "" 179 180 def flush(self, out: CWriter, cast_type: str = "PyObject *") -> None: 181 out.start_line() 182 for var in self.variables: 183 if not var.peek: 184 cast = f"({cast_type})" if var.type else "" 185 if var.name not in UNUSED and not var.is_array(): 186 if var.condition: 187 if var.condition == "0": 188 continue 189 elif var.condition != "1": 190 out.emit(f"if ({var.condition}) ") 191 out.emit( 192 f"stack_pointer[{self.base_offset.to_c()}] = {cast}{var.name};\n" 193 ) 194 self.base_offset.push(var) 195 if self.base_offset.to_c() != self.top_offset.to_c(): 196 print("base", self.base_offset.to_c(), "top", self.top_offset.to_c()) 197 assert False 198 number = self.base_offset.to_c() 199 if number != "0": 200 out.emit(f"stack_pointer += {number};\n") 201 self.variables = [] 202 self.base_offset.clear() 203 self.top_offset.clear() 204 out.start_line() 205 206 def peek_offset(self) -> str: 207 peek = self.base_offset.copy() 208 for var in self.variables: 209 if not var.peek: 210 break 211 peek.push(var) 212 return peek.to_c() 213 214 def as_comment(self) -> str: 215 return f"/* Variables: {[v.name for v in self.variables]}. Base offset: {self.base_offset.to_c()}. Top offset: {self.top_offset.to_c()} */" 216 217 218def get_stack_effect(inst: Instruction) -> Stack: 219 stack = Stack() 220 for uop in inst.parts: 221 if not isinstance(uop, Uop): 222 continue 223 for var in reversed(uop.stack.inputs): 224 stack.pop(var) 225 for i, var in enumerate(uop.stack.outputs): 226 stack.push(var) 227 return stack 228