import re
from analyzer import StackItem, StackEffect, Instruction, Uop, PseudoInstruction
from collections import defaultdict
from dataclasses import dataclass
from cwriter import CWriter
from typing import Iterator, Tuple
UNUSED = {"unused"}
def maybe_parenthesize(sym: str) -> str:
"""Add parentheses around a string if it contains an operator
and is not already parenthesized.
An exception is made for '*' which is common and harmless
in the context where the symbolic size is used.
"""
if sym.startswith("(") and sym.endswith(")"):
return sym
if re.match(r"^[\s\w*]+$", sym):
return sym
else:
return f"({sym})"
def var_size(var: StackItem) -> str:
if var.condition:
# Special case simplifications
if var.condition == "0":
return "0"
elif var.condition == "1":
return var.get_size()
elif var.condition == "oparg & 1" and not var.size:
return f"({var.condition})"
else:
return f"(({var.condition}) ? {var.get_size()} : 0)"
elif var.size:
return var.size
else:
return "1"
@dataclass
class Local:
item: StackItem
cached: bool
in_memory: bool
defined: bool
def __repr__(self) -> str:
return f"Local('{self.item.name}', mem={self.in_memory}, defined={self.defined}, array={self.is_array()})"
def compact_str(self) -> str:
mtag = "M" if self.in_memory else ""
dtag = "D" if self.defined else ""
atag = "A" if self.is_array() else ""
return f"'{self.item.name}'{mtag}{dtag}{atag}"
@staticmethod
def unused(defn: StackItem) -> "Local":
return Local(defn, False, defn.is_array(), False)
@staticmethod
def undefined(defn: StackItem) -> "Local":
array = defn.is_array()
return Local(defn, not array, array, False)
@staticmethod
def redefinition(var: StackItem, prev: "Local") -> "Local":
assert var.is_array() == prev.is_array()
return Local(var, prev.cached, prev.in_memory, True)
@staticmethod
def from_memory(defn: StackItem) -> "Local":
return Local(defn, True, True, True)
def copy(self) -> "Local":
return Local(
self.item,
self.cached,
self.in_memory,
self.defined
)
@property
def size(self) -> str:
return self.item.size
@property
def name(self) -> str:
return self.item.name
@property
def condition(self) -> str | None:
return self.item.condition
def is_array(self) -> bool:
return self.item.is_array()
def __eq__(self, other: object) -> bool:
if not isinstance(other, Local):
return NotImplemented
return (
self.item is other.item
and self.cached is other.cached
and self.in_memory is other.in_memory
and self.defined is other.defined
)
@dataclass
class StackOffset:
"The stack offset of the virtual base of the stack from the physical stack pointer"
popped: list[str]
pushed: list[str]
@staticmethod
def empty() -> "StackOffset":
return StackOffset([], [])
def copy(self) -> "StackOffset":
return StackOffset(self.popped[:], self.pushed[:])
def pop(self, item: StackItem) -> None:
self.popped.append(var_size(item))
def push(self, item: StackItem) -> None:
self.pushed.append(var_size(item))
def __sub__(self, other: "StackOffset") -> "StackOffset":
return StackOffset(self.popped + other.pushed, self.pushed + other.popped)
def __neg__(self) -> "StackOffset":
return StackOffset(self.pushed, self.popped)
def simplify(self) -> None:
"Remove matching values from both the popped and pushed list"
if not self.popped:
self.pushed.sort()
return
if not self.pushed:
self.popped.sort()
return
# Sort the list so the lexically largest element is last.
popped = sorted(self.popped)
pushed = sorted(self.pushed)
self.popped = []
self.pushed = []
while popped and pushed:
pop = popped.pop()
push = pushed.pop()
if pop == push:
pass
elif pop > push:
# if pop > push, there can be no element in pushed matching pop.
self.popped.append(pop)
pushed.append(push)
else:
self.pushed.append(push)
popped.append(pop)
self.popped.extend(popped)
self.pushed.extend(pushed)
self.pushed.sort()
self.popped.sort()
def to_c(self) -> str:
self.simplify()
int_offset = 0
symbol_offset = ""
for item in self.popped:
try:
int_offset -= int(item)
except ValueError:
symbol_offset += f" - {maybe_parenthesize(item)}"
for item in self.pushed:
try:
int_offset += int(item)
except ValueError:
symbol_offset += f" + {maybe_parenthesize(item)}"
if symbol_offset and not int_offset:
res = symbol_offset
else:
res = f"{int_offset}{symbol_offset}"
if res.startswith(" + "):
res = res[3:]
if res.startswith(" - "):
res = "-" + res[3:]
return res
def as_int(self) -> int | None:
self.simplify()
int_offset = 0
for item in self.popped:
try:
int_offset -= int(item)
except ValueError:
return None
for item in self.pushed:
try:
int_offset += int(item)
except ValueError:
return None
return int_offset
def clear(self) -> None:
self.popped = []
self.pushed = []
def __bool__(self) -> bool:
self.simplify()
return bool(self.popped) or bool(self.pushed)
def __eq__(self, other: object) -> bool:
if not isinstance(other, StackOffset):
return NotImplemented
return self.to_c() == other.to_c()
class StackError(Exception):
pass
def array_or_scalar(var: StackItem | Local) -> str:
return "array" if var.is_array() else "scalar"
class Stack:
def __init__(self) -> None:
self.top_offset = StackOffset.empty()
self.base_offset = StackOffset.empty()
self.variables: list[Local] = []
self.defined: set[str] = set()
def pop(self, var: StackItem, extract_bits: bool = True) -> tuple[str, Local]:
self.top_offset.pop(var)
indirect = "&" if var.is_array() else ""
if self.variables:
popped = self.variables.pop()
if var.is_array() ^ popped.is_array():
raise StackError(
f"Array mismatch when popping '{popped.name}' from stack to assign to '{var.name}'. "
f"Expected {array_or_scalar(var)} got {array_or_scalar(popped)}"
)
if popped.size != var.size:
raise StackError(
f"Size mismatch when popping '{popped.name}' from stack to assign to '{var.name}'. "
f"Expected {var_size(var)} got {var_size(popped.item)}"
)
if var.name in UNUSED:
if popped.name not in UNUSED and popped.name in self.defined:
raise StackError(
f"Value is declared unused, but is already cached by prior operation as '{popped.name}'"
)
return "", popped
if not var.used:
return "", popped
self.defined.add(var.name)
if popped.defined:
if popped.name == var.name:
return "", popped
else:
defn = f"{var.name} = {popped.name};\n"
else:
if var.is_array():
defn = f"{var.name} = &stack_pointer[{self.top_offset.to_c()}];\n"
else:
defn = f"{var.name} = stack_pointer[{self.top_offset.to_c()}];\n"
popped.in_memory = True
return defn, Local.redefinition(var, popped)
self.base_offset.pop(var)
if var.name in UNUSED or not var.used:
return "", Local.unused(var)
self.defined.add(var.name)
cast = f"({var.type})" if (not indirect and var.type) else ""
bits = ".bits" if cast and extract_bits else ""
assign = f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}]{bits};"
if var.condition:
if var.condition == "1":
assign = f"{assign}\n"
elif var.condition == "0":
return "", Local.unused(var)
else:
assign = f"if ({var.condition}) {{ {assign} }}\n"
else:
assign = f"{assign}\n"
return assign, Local.from_memory(var)
def push(self, var: Local) -> None:
assert(var not in self.variables)
self.variables.append(var)
self.top_offset.push(var.item)
if var.item.used:
self.defined.add(var.name)
@staticmethod
def _do_emit(
out: CWriter,
var: StackItem,
base_offset: StackOffset,
cast_type: str = "uintptr_t",
extract_bits: bool = True,
) -> None:
cast = f"({cast_type})" if var.type else ""
bits = ".bits" if cast and extract_bits else ""
if var.condition == "0":
return
if var.condition and var.condition != "1":
out.emit(f"if ({var.condition}) ")
out.emit(f"stack_pointer[{base_offset.to_c()}]{bits} = {cast}{var.name};\n")
def _adjust_stack_pointer(self, out: CWriter, number: str) -> None:
if number != "0":
out.start_line()
out.emit(f"stack_pointer += {number};\n")
out.emit("assert(WITHIN_STACK_BOUNDS());\n")
def flush(
self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = True
) -> None:
out.start_line()
var_offset = self.base_offset.copy()
for var in self.variables:
if (
var.defined and
not var.in_memory
):
Stack._do_emit(out, var.item, var_offset, cast_type, extract_bits)
var.in_memory = True
var_offset.push(var.item)
number = self.top_offset.to_c()
self._adjust_stack_pointer(out, number)
self.base_offset -= self.top_offset
self.top_offset.clear()
out.start_line()
def is_flushed(self) -> bool:
return not self.variables and not self.base_offset and not self.top_offset
def peek_offset(self) -> str:
return self.top_offset.to_c()
def as_comment(self) -> str:
variables = ", ".join([v.compact_str() for v in self.variables])
return (
f"/* Variables: {variables}. base: {self.base_offset.to_c()}. top: {self.top_offset.to_c()} */"
)
def copy(self) -> "Stack":
other = Stack()
other.top_offset = self.top_offset.copy()
other.base_offset = self.base_offset.copy()
other.variables = [var.copy() for var in self.variables]
other.defined = set(self.defined)
return other
def __eq__(self, other: object) -> bool:
if not isinstance(other, Stack):
return NotImplemented
return (
self.top_offset == other.top_offset
and self.base_offset == other.base_offset
and self.variables == other.variables
)
def align(self, other: "Stack", out: CWriter) -> None:
if len(self.variables) != len(other.variables):
raise StackError("Cannot align stacks: differing variables")
if self.top_offset == other.top_offset:
return
diff = self.top_offset - other.top_offset
try:
self.top_offset -= diff
self.base_offset -= diff
self._adjust_stack_pointer(out, diff.to_c())
except ValueError:
raise StackError("Cannot align stacks: cannot adjust stack pointer")
def merge(self, other: "Stack", out: CWriter) -> None:
if len(self.variables) != len(other.variables):
raise StackError("Cannot merge stacks: differing variables")
for self_var, other_var in zip(self.variables, other.variables):
if self_var.name != other_var.name:
raise StackError(f"Mismatched variables on stack: {self_var.name} and {other_var.name}")
self_var.defined = self_var.defined and other_var.defined
self_var.in_memory = self_var.in_memory and other_var.in_memory
self.align(other, out)
def stacks(inst: Instruction | PseudoInstruction) -> Iterator[StackEffect]:
if isinstance(inst, Instruction):
for uop in inst.parts:
if isinstance(uop, Uop):
yield uop.stack
else:
assert isinstance(inst, PseudoInstruction)
yield inst.stack
def apply_stack_effect(stack: Stack, effect: StackEffect) -> None:
locals: dict[str, Local] = {}
for var in reversed(effect.inputs):
_, local = stack.pop(var)
if var.name != "unused":
locals[local.name] = local
for var in effect.outputs:
if var.name in locals:
local = locals[var.name]
else:
local = Local.unused(var)
stack.push(local)
def get_stack_effect(inst: Instruction | PseudoInstruction) -> Stack:
stack = Stack()
for s in stacks(inst):
apply_stack_effect(stack, s)
return stack
def get_stack_effects(inst: Instruction | PseudoInstruction) -> list[Stack]:
"""Returns a list of stack effects after each uop"""
result = []
stack = Stack()
for s in stacks(inst):
apply_stack_effect(stack, s)
result.append(stack.copy())
return result
@dataclass
class Storage:
stack: Stack
inputs: list[Local]
outputs: list[Local]
peeks: list[Local]
spilled: int = 0
@staticmethod
def needs_defining(var: Local) -> bool:
return (
not var.defined and
not var.is_array() and
var.name != "unused"
)
@staticmethod
def is_live(var: Local) -> bool:
return (
var.defined and
var.name != "unused"
)
def first_input_not_cleared(self) -> str:
for input in self.inputs:
if input.defined:
return input.name
return ""
def clear_inputs(self, reason:str) -> None:
while self.inputs:
tos = self.inputs.pop()
if self.is_live(tos) and not tos.is_array():
raise StackError(
f"Input '{tos.name}' is still live {reason}"
)
self.stack.pop(tos.item)
def clear_dead_inputs(self) -> None:
live = ""
while self.inputs:
tos = self.inputs[-1]
if self.is_live(tos):
live = tos.name
break
self.inputs.pop()
self.stack.pop(tos.item)
for var in self.inputs:
if not var.defined and not var.is_array() and var.name != "unused":
raise StackError(
f"Input '{var.name}' is not live, but '{live}' is"
)
def _push_defined_outputs(self) -> None:
defined_output = ""
for output in self.outputs:
if output.defined and not output.in_memory:
defined_output = output.name
if not defined_output:
return
self.clear_inputs(f"when output '{defined_output}' is defined")
undefined = ""
for out in self.outputs:
if out.defined:
if undefined:
f"Locals not defined in stack order. "
f"Expected '{undefined}' to be defined before '{out.name}'"
else:
undefined = out.name
while self.outputs and not self.needs_defining(self.outputs[0]):
out = self.outputs.pop(0)
self.stack.push(out)
def locals_cached(self) -> bool:
for out in self.outputs:
if out.defined:
return True
return False
def flush(self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = True) -> None:
self.clear_dead_inputs()
self._push_defined_outputs()
self.stack.flush(out, cast_type, extract_bits)
def save(self, out: CWriter) -> None:
assert self.spilled >= 0
if self.spilled == 0:
self.flush(out)
out.start_line()
out.emit("_PyFrame_SetStackPointer(frame, stack_pointer);\n")
self.spilled += 1
def reload(self, out: CWriter) -> None:
if self.spilled == 0:
raise StackError("Cannot reload stack as it hasn't been saved")
assert self.spilled > 0
self.spilled -= 1
if self.spilled == 0:
out.start_line()
out.emit("stack_pointer = _PyFrame_GetStackPointer(frame);\n")
@staticmethod
def for_uop(stack: Stack, uop: Uop, extract_bits: bool = True) -> tuple[list[str], "Storage"]:
code_list: list[str] = []
inputs: list[Local] = []
peeks: list[Local] = []
for input in reversed(uop.stack.inputs):
code, local = stack.pop(input, extract_bits)
code_list.append(code)
if input.peek:
peeks.append(local)
else:
inputs.append(local)
inputs.reverse()
peeks.reverse()
for peek in peeks:
stack.push(peek)
top_offset = stack.top_offset.copy()
for ouput in uop.stack.outputs:
if ouput.is_array() and ouput.used and not ouput.peek:
c_offset = top_offset.to_c()
top_offset.push(ouput)
code_list.append(f"{ouput.name} = &stack_pointer[{c_offset}];\n")
else:
top_offset.push(ouput)
for var in inputs:
stack.push(var)
outputs = [ Local.undefined(var) for var in uop.stack.outputs if not var.peek ]
return code_list, Storage(stack, inputs, outputs, peeks)
@staticmethod
def copy_list(arg: list[Local]) -> list[Local]:
return [ l.copy() for l in arg ]
def copy(self) -> "Storage":
new_stack = self.stack.copy()
variables = { var.name: var for var in new_stack.variables }
inputs = [ variables[var.name] for var in self.inputs]
assert [v.name for v in inputs] == [v.name for v in self.inputs], (inputs, self.inputs)
return Storage(
new_stack, inputs,
self.copy_list(self.outputs), self.copy_list(self.peeks)
)
def sanity_check(self) -> None:
names: set[str] = set()
for var in self.inputs:
if var.name in names:
raise StackError(f"Duplicate name {var.name}")
names.add(var.name)
names = set()
for var in self.outputs:
if var.name in names:
raise StackError(f"Duplicate name {var.name}")
names.add(var.name)
names = set()
for var in self.stack.variables:
if var.name in names:
raise StackError(f"Duplicate name {var.name}")
names.add(var.name)
def is_flushed(self) -> bool:
for var in self.outputs:
if var.defined and not var.in_memory:
return False
return self.stack.is_flushed()
def merge(self, other: "Storage", out: CWriter) -> None:
self.sanity_check()
if len(self.inputs) != len(other.inputs):
self.clear_dead_inputs()
other.clear_dead_inputs()
if len(self.inputs) != len(other.inputs):
diff = self.inputs[-1] if len(self.inputs) > len(other.inputs) else other.inputs[-1]
raise StackError(f"Unmergeable inputs. Differing state of '{diff.name}'")
for var, other_var in zip(self.inputs, other.inputs):
if var.defined != other_var.defined:
raise StackError(f"'{var.name}' is cleared on some paths, but not all")
if len(self.outputs) != len(other.outputs):
self._push_defined_outputs()
other._push_defined_outputs()
if len(self.outputs) != len(other.outputs):
var = self.outputs[0] if len(self.outputs) > len(other.outputs) else other.outputs[0]
raise StackError(f"'{var.name}' is set on some paths, but not all")
self.stack.merge(other.stack, out)
self.sanity_check()
def push_outputs(self) -> None:
if self.spilled:
raise StackError(f"Unbalanced stack spills")
self.clear_inputs("at the end of the micro-op")
if self.inputs:
raise StackError(f"Input variable '{self.inputs[-1].name}' is still live")
self._push_defined_outputs()
if self.outputs:
for out in self.outputs:
if self.needs_defining(out):
raise StackError(f"Output variable '{self.outputs[0].name}' is not defined")
self.stack.push(out)
self.outputs = []
def as_comment(self) -> str:
stack_comment = self.stack.as_comment()
next_line = "\n "
inputs = ", ".join([var.compact_str() for var in self.inputs])
outputs = ", ".join([var.compact_str() for var in self.outputs])
peeks = ", ".join([var.name for var in self.peeks])
return f"{stack_comment[:-2]}{next_line}inputs: {inputs}{next_line}outputs: {outputs}{next_line}peeks: {peeks} */"
Codebase Browser
cpython