mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
bikini
94 lines
4.5 KiB
Python
94 lines
4.5 KiB
Python
from __future__ import annotations
|
|
import ast
|
|
import copy
|
|
import random
|
|
|
|
def _is_pure(node: ast.AST) -> bool:
|
|
if isinstance(node, ast.Name):
|
|
return True
|
|
if isinstance(node, ast.Constant):
|
|
return isinstance(node.value, (int, bool))
|
|
if isinstance(node, ast.Attribute):
|
|
return _is_pure(node.value)
|
|
if isinstance(node, ast.Subscript):
|
|
return _is_pure(node.value) and _is_pure(node.slice)
|
|
if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd, ast.Invert)):
|
|
return _is_pure(node.operand)
|
|
return False
|
|
|
|
def _dup(node: ast.AST) -> ast.AST:
|
|
return copy.deepcopy(node)
|
|
|
|
class MBATransformer(ast.NodeTransformer):
|
|
|
|
def __init__(self, rng: random.Random, prob: float=0.7):
|
|
self.rng = rng
|
|
self.prob = prob
|
|
|
|
def visit_BinOp(self, node: ast.BinOp) -> ast.AST:
|
|
self.generic_visit(node)
|
|
if not (_is_pure(node.left) and _is_pure(node.right)):
|
|
return node
|
|
if self.rng.random() > self.prob:
|
|
return node
|
|
a, b = (node.left, node.right)
|
|
op = node.op
|
|
if isinstance(op, ast.BitXor):
|
|
return self._xor(a, b)
|
|
if isinstance(op, ast.BitOr):
|
|
return self._or(a, b)
|
|
if isinstance(op, ast.BitAnd):
|
|
return self._and(a, b)
|
|
if isinstance(op, ast.LShift) and isinstance(b, ast.Constant) and isinstance(b.value, int):
|
|
return self._lshift(a, b.value)
|
|
if isinstance(op, ast.RShift) and isinstance(b, ast.Constant) and isinstance(b.value, int):
|
|
return self._rshift(a, b.value)
|
|
return node
|
|
|
|
def visit_UnaryOp(self, node: ast.UnaryOp) -> ast.AST:
|
|
self.generic_visit(node)
|
|
if not _is_pure(node.operand):
|
|
return node
|
|
if self.rng.random() > self.prob:
|
|
return node
|
|
if isinstance(node.op, ast.Invert):
|
|
return self._invert(node.operand)
|
|
return node
|
|
|
|
def _xor(self, a: ast.AST, b: ast.AST) -> ast.AST:
|
|
choice = self.rng.choice(('subor', 'ornotand', 'addsub'))
|
|
if choice == 'subor':
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))
|
|
if choice == 'ornotand':
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)), op=ast.BitAnd(), right=ast.UnaryOp(op=ast.Invert(), operand=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b))))
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=ast.Constant(value=2), op=ast.Mult(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b))))
|
|
|
|
def _or(self, a: ast.AST, b: ast.AST) -> ast.AST:
|
|
if self.rng.random() < 0.5:
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitXor(), right=_dup(b)), op=ast.BitXor(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))
|
|
|
|
def _and(self, a: ast.AST, b: ast.AST) -> ast.AST:
|
|
if self.rng.random() < 0.5:
|
|
return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)))
|
|
return ast.UnaryOp(op=ast.Invert(), operand=ast.BinOp(left=ast.UnaryOp(op=ast.Invert(), operand=_dup(a)), op=ast.BitOr(), right=ast.UnaryOp(op=ast.Invert(), operand=_dup(b))))
|
|
|
|
def _invert(self, a: ast.AST) -> ast.AST:
|
|
return ast.BinOp(left=ast.UnaryOp(op=ast.USub(), operand=_dup(a)), op=ast.Sub(), right=ast.Constant(value=1))
|
|
|
|
def _lshift(self, a: ast.AST, n: int) -> ast.AST:
|
|
if n < 0 or n > 60:
|
|
return ast.BinOp(left=_dup(a), op=ast.LShift(), right=ast.Constant(value=n))
|
|
return ast.BinOp(left=_dup(a), op=ast.Mult(), right=ast.Constant(value=1 << n))
|
|
|
|
def _rshift(self, a: ast.AST, n: int) -> ast.AST:
|
|
if n < 0 or n > 60:
|
|
return ast.BinOp(left=_dup(a), op=ast.RShift(), right=ast.Constant(value=n))
|
|
return ast.BinOp(left=_dup(a), op=ast.FloorDiv(), right=ast.Constant(value=1 << n))
|
|
|
|
def visit_match_case(self, node: ast.match_case) -> ast.AST:
|
|
if node.guard is not None:
|
|
node.guard = self.visit(node.guard)
|
|
node.body = [self.visit(stmt) for stmt in node.body]
|
|
return node
|