diff --git a/uncommonlisp.py b/uncommonlisp.py index 9f0f5fe..1a9baa2 100644 --- a/uncommonlisp.py +++ b/uncommonlisp.py @@ -2229,6 +2229,326 @@ def bc_compile_proc(proc, env): return CompiledProc(code, proc.params, proc.rest, proc.env, proc.name) +############################################################################### +# JIT: Transpile bytecode to Python source, exec() it +############################################################################### + +def _jit_compile(proc): + """JIT a Proc/CompiledProc to a native Python function via AST transpilation.""" + if isinstance(proc, CompiledProc): + # Need the original AST — can't JIT from bytecode alone + return None + if not isinstance(proc, Proc): return None + if proc.rest: return None # rest args too complex + + params = [str(p) for p in proc.params] + name = proc.name or '_fn' + # Sanitize name for Python + pyname = ''.join(c if c.isalnum() or c == '_' else '_' for c in name) + if not pyname or pyname[0].isdigit(): pyname = '_' + pyname + + try: + body_src = _jit_expr(proc.body, params) + except _JitBail: + return None + + source = f'def {pyname}({", ".join(params)}):\n return {body_src}' + ns = {'Pair': Pair, 'NIL': NIL, 'VOID': VOID, 'S': S, 'Fraction': Fraction, + 'True': True, 'False': False} + # Add self-reference for recursion + try: + exec(source, ns) + fn = ns[pyname] + # For recursive functions, bind self + if name in source: + ns[pyname] = fn + exec(source, ns) + fn = ns[pyname] + def jit_wrapper(args, env): + return fn(*args) + jit_wrapper._jit_source = source + jit_wrapper._jit_name = name + return jit_wrapper + except Exception: + return None + + +class _JitBail(Exception): pass + +def _jit_expr(body, params): + """Transpile Scheme body (list of exprs) to a Python expression string.""" + if len(body) == 1: return _jit_one(body[0], params) + # begin: evaluate all, return last (only works if non-last are side-effect-free) + # For JIT, bail on side effects in non-tail position + return _jit_one(body[-1], params) + +def _jit_one(expr, params): + """Transpile a single Scheme expression to a Python expression string.""" + if expr is True: return 'True' + if expr is False: return 'False' + if expr is NIL: return 'NIL' + if isinstance(expr, (int, float)): return repr(expr) + if isinstance(expr, Fraction): return f'Fraction({expr.numerator},{expr.denominator})' + if isinstance(expr, str) and not isinstance(expr, Symbol): return repr(expr) + if isinstance(expr, Symbol): + s = str(expr) + if s in params or s.replace('-','_').replace('?','_p').replace('!','_b').isidentifier(): + return s + raise _JitBail() + if not isinstance(expr, Pair): raise _JitBail() + + head = expr.car; args = _L(expr.cdr) + + # Special forms + if head is S('if'): + test = _jit_one(args[0], params) + then = _jit_one(args[1], params) + els = _jit_one(args[2], params) if len(args) > 2 else 'VOID' + return f'({then} if {test} else {els})' + + if head is S('cond'): + return _jit_cond(args, params) + + if head is S('begin'): + return _jit_one(args[-1], params) + + if head is S('let'): + if isinstance(args[0], Symbol): + # Named let → while loop as a helper function + return _jit_named_let(args, params) + # Regular let → inline + binds = _L(args[0]); body = args[1:] + bind_strs = [] + new_params = list(params) + for b in binds: + bp = _L(b) + bind_strs.append(f'{bp[0]}={_jit_one(bp[1], params)}') + new_params.append(str(bp[0])) + body_str = _jit_expr(body, new_params) + return f'(lambda {",".join(str(_L(b)[0]) for b in binds)}: {body_str})({",".join(_jit_one(_L(b)[1], params) for b in binds)})' + + if head is S('and'): + if not args: return 'True' + parts = [_jit_one(a, params) for a in args] + return ' and '.join(f'({p})' for p in parts) + + if head is S('or'): + if not args: return 'False' + parts = [_jit_one(a, params) for a in args] + return ' or '.join(f'({p})' for p in parts) + + if head is S('quote'): + raise _JitBail() # can't represent arbitrary quoted data + + # Known pure functions → inline Python + _PYOP = { + S('+'): '+', S('-'): '-', S('*'): '*', + S('='): '==', S('<'): '<', S('>'): '>', + S('<='): '<=', S('>='): '>=', + } + if head in _PYOP and len(args) == 2: + a = _jit_one(args[0], params); b = _jit_one(args[1], params) + return f'({a} {_PYOP[head]} {b})' + if head is S('+') and len(args) == 1: return _jit_one(args[0], params) + if head is S('-') and len(args) == 1: return f'(-{_jit_one(args[0], params)})' + if head is S('not'): return f'(not {_jit_one(args[0], params)})' + if head is S('zero?'): return f'({_jit_one(args[0], params)} == 0)' + if head is S('null?'): return f'({_jit_one(args[0], params)} is NIL)' + if head is S('pair?'): return f'isinstance({_jit_one(args[0], params)}, Pair)' + if head is S('car'): return f'{_jit_one(args[0], params)}.car' + if head is S('cdr'): return f'{_jit_one(args[0], params)}.cdr' + if head is S('cons'): return f'Pair({_jit_one(args[0], params)},{_jit_one(args[1], params)})' + if head is S('remainder'): return f'({_jit_one(args[0], params)} % {_jit_one(args[1], params)})' + if head is S('modulo'): return f'({_jit_one(args[0], params)} % {_jit_one(args[1], params)})' + if head is S('abs'): return f'abs({_jit_one(args[0], params)})' + if head is S('expt'): return f'({_jit_one(args[0], params)} ** {_jit_one(args[1], params)})' + + # Generic function call + if isinstance(head, Symbol): + fn = str(head) + call_args = ', '.join(_jit_one(a, params) for a in args) + return f'{fn}({call_args})' + + raise _JitBail() + + +def _jit_cond(clauses, params): + """Transpile cond to nested ternary.""" + if not clauses: return 'VOID' + cl = _L(clauses[0]) + if cl[0] is S('else'): + return _jit_expr(cl[1:], params) + test = _jit_one(cl[0], params) + then = _jit_expr(cl[1:], params) if len(cl) > 1 else test + rest = _jit_cond(clauses[1:], params) + return f'({then} if {test} else {rest})' + + +def _jit_named_let(args, params): + """Transpile named let to a Python helper with while loop. + Returns a Python expression that calls the helper.""" + name = str(args[0]) + binds = _L(args[1]); body = args[2:] + bparams = [str(_L(b)[0]) for b in binds] + all_params = list(params) + bparams + [name] + + # The body becomes a while-True loop with return/continue + body_expr = _jit_expr(body, all_params) + + # For simple tail-recursive patterns (if test return_val (loop ...)), + # we can generate a while loop. But for the general case, + # use recursive Python function. + init_args = ', '.join(_jit_one(_L(b)[1], params) for b in binds) + # Generate as a local recursive function + raise _JitBail() # named-let needs statement-level code, not expression + + +def _jit_transpile(instrs, params, name, has_self_tc): + """Transpile bytecode to Python source lines using recursive descent.""" + lines = [f'def {name}({", ".join(params)}):'] + indent = ' ' + if has_self_tc: + lines.append(f'{indent}while True:') + indent = ' ' + + def _expr(ip): + """Transpile expression starting at ip, return (python_expr_string, next_ip).""" + if ip >= len(instrs): raise _JitBail() + op, arg = instrs[ip] + if op == OP_CONST: + if arg is True: return 'True', ip+1 + if arg is False: return 'False', ip+1 + if arg is NIL: return 'NIL', ip+1 + if isinstance(arg, str) and not isinstance(arg, Symbol): return repr(arg), ip+1 + return repr(arg), ip+1 + if op == OP_LOOKUP: return str(arg), ip+1 + if op == OP_LOOK_ADD1: return f'({arg} + 1)', ip+1 + if op == OP_LOOK_SUB1: return f'({arg} - 1)', ip+1 + if op == OP_VOID: return 'VOID', ip+1 + # Binary ops: left expr, right expr, op + if op in (OP_ADD, OP_SUB, OP_MUL, OP_NUM_EQ, OP_LT, OP_GT, OP_LE, OP_GE, + OP_CONS, OP_VEC_REF): + raise _JitBail() # handled by stack below + raise _JitBail() + + def _block(ip, end): + """Transpile a block of instructions [ip, end), return list of (stmt, next_ip).""" + stmts = []; stack = [] + while ip < end: + op, arg = instrs[ip] + if op == OP_CONST: + if arg is True: stack.append('True') + elif arg is False: stack.append('False') + elif arg is NIL: stack.append('NIL') + elif isinstance(arg, str) and not isinstance(arg, Symbol): + stack.append(repr(arg)) + else: stack.append(repr(arg)) + ip += 1 + elif op == OP_LOOKUP: stack.append(str(arg)); ip += 1 + elif op == OP_LOOK_ADD1: stack.append(f'({arg} + 1)'); ip += 1 + elif op == OP_LOOK_SUB1: stack.append(f'({arg} - 1)'); ip += 1 + elif op == OP_VOID: stack.append('VOID'); ip += 1 + elif op == OP_ADD: b=stack.pop(); a=stack.pop(); stack.append(f'({a} + {b})'); ip+=1 + elif op == OP_SUB: b=stack.pop(); a=stack.pop(); stack.append(f'({a} - {b})'); ip+=1 + elif op == OP_MUL: b=stack.pop(); a=stack.pop(); stack.append(f'({a} * {b})'); ip+=1 + elif op == OP_NEG: a=stack.pop(); stack.append(f'(-{a})'); ip+=1 + elif op == OP_ADD1: a=stack.pop(); stack.append(f'({a} + 1)'); ip+=1 + elif op == OP_SUB1: a=stack.pop(); stack.append(f'({a} - 1)'); ip+=1 + elif op == OP_NUM_EQ: b=stack.pop(); a=stack.pop(); stack.append(f'({a} == {b})'); ip+=1 + elif op == OP_LT: b=stack.pop(); a=stack.pop(); stack.append(f'({a} < {b})'); ip+=1 + elif op == OP_GT: b=stack.pop(); a=stack.pop(); stack.append(f'({a} > {b})'); ip+=1 + elif op == OP_LE: b=stack.pop(); a=stack.pop(); stack.append(f'({a} <= {b})'); ip+=1 + elif op == OP_GE: b=stack.pop(); a=stack.pop(); stack.append(f'({a} >= {b})'); ip+=1 + elif op == OP_NOT: a=stack.pop(); stack.append(f'(not {a})'); ip+=1 + elif op == OP_ZERO_P: a=stack.pop(); stack.append(f'({a} == 0)'); ip+=1 + elif op == OP_NULL_P: a=stack.pop(); stack.append(f'({a} is NIL)'); ip+=1 + elif op == OP_PAIR_P: a=stack.pop(); stack.append(f'isinstance({a}, Pair)'); ip+=1 + elif op == OP_CAR: a=stack.pop(); stack.append(f'{a}.car'); ip+=1 + elif op == OP_CDR: a=stack.pop(); stack.append(f'{a}.cdr'); ip+=1 + elif op == OP_CONS: d=stack.pop(); a=stack.pop(); stack.append(f'Pair({a},{d})'); ip+=1 + elif op == OP_VEC_REF: i=stack.pop(); v=stack.pop(); stack.append(f'{v}[{i}]'); ip+=1 + elif op == OP_POP: stack.pop() if stack else None; ip+=1 + elif op == OP_DUP: stack.append(stack[-1]); ip+=1 + elif op == OP_JUMP: ip = arg # forward jump = skip to target + elif op == OP_JUMP_IF_FALSE: + cond = stack.pop() + else_ip = arg + # Find JUMP at end of then-block → end of if + # Pattern: [then-block] JUMP end [else-block] end: + then_end = else_ip - 1 + if then_end >= 0 and instrs[then_end][0] == OP_JUMP: + end_ip = instrs[then_end][1] + then_stmts = _block(ip, then_end) + else_stmts = _block(else_ip, end_ip) + stmts.append(('if', cond, then_stmts, else_stmts)) + ip = end_ip + else: + # No else: if (not cond) skip + then_stmts = _block(ip, else_ip) + stmts.append(('if', cond, then_stmts, [])) + ip = else_ip + elif op == OP_RETURN: + val = stack.pop() if stack else 'VOID' + stmts.append(('return', val)); ip+=1 + elif op == OP_SELF_TAIL_CALL: + tc_n, tc_params = arg + pnames = [str(p) for p in tc_params] + args = []; + for _ in range(tc_n): args.insert(0, stack.pop()) + stmts.append(('self_tc', pnames, args)); ip+=1 + elif op == OP_CALL: + call_args = []; + for _ in range(arg): call_args.insert(0, stack.pop()) + func = stack.pop() + stack.append(f'{func}({",".join(call_args)})'); ip+=1 + elif op == OP_TAIL_CALL: + call_args = [] + for _ in range(arg): call_args.insert(0, stack.pop()) + func = stack.pop() + stmts.append(('return', f'{func}({",".join(call_args)})')); ip+=1 + elif op == OP_SET: + val = stack.pop() + stmts.append(('assign', str(arg), val)); ip+=1 + elif op == OP_DEFINE: + val = stack.pop() + stmts.append(('assign', str(arg), val)); ip+=1 + else: + raise _JitBail() + return stmts + + def _emit(stmts, ind): + for s in stmts: + if s[0] == 'return': + lines.append(f'{ind}return {s[1]}') + elif s[0] == 'self_tc': + pnames, args = s[1], s[2] + lines.append(f'{ind}{", ".join(pnames)} = {", ".join(args)}') + lines.append(f'{ind}continue') + elif s[0] == 'assign': + lines.append(f'{ind}{s[1]} = {s[2]}') + elif s[0] == 'if': + _, cond, then_s, else_s = s + lines.append(f'{ind}if {cond}:') + if then_s: _emit(then_s, ind + ' ') + else: lines.append(f'{ind} pass') + if else_s: + lines.append(f'{ind}else:') + _emit(else_s, ind + ' ') + + stmts = _block(0, len(instrs)) + _emit(stmts, indent) + return lines + + +def _jit_try(proc, env): + """Try to JIT a procedure. Returns JIT'd callable or original proc.""" + if not isinstance(proc, (CompiledProc, Proc)): return proc + if isinstance(proc, Proc): + proc = bc_compile_proc(proc, env) + jit_fn = _jit_compile(proc) + return jit_fn if jit_fn else proc + + def _disassemble(proc): """Return human-readable bytecode listing.""" if isinstance(proc, Proc): @@ -3022,6 +3342,8 @@ def make_global_env(): # ── Bytecode compiler ──────────────────────────────────────────────────── d(S('compile'), lambda a, e: bc_compile_proc(a[0], e)) d(S('compiled?'), lambda a, _: isinstance(a[0], CompiledProc)) + d(S('jit'), lambda a, e: _jit_try(a[0], e)) + d(S('jit-source'), lambda a, e: getattr(_jit_compile(a[0] if isinstance(a[0], CompiledProc) else bc_compile_proc(a[0], e)), '_jit_source', False) if isinstance(a[0], (Proc, CompiledProc)) else False) d(S('disassemble'),lambda a, _: (print(_disassemble(a[0])) or VOID)) d(S('save-compiled'), lambda a, _: save_compiled(_str_val(a[0]), a[1]) or VOID) d(S('load-compiled'), lambda a, e: load_compiled(_str_val(a[0]), e))