#ifdef _Py_TIER2

#include "Python.h"

#include "opcode.h"

#include "pycore_dict.h"

#include "pycore_interp.h"

#include "pycore_opcode_metadata.h"

#include "pycore_opcode_utils.h"

#include "pycore_pystate.h" // _PyInterpreterState_GET()

#include "pycore_uop_metadata.h"

#include "pycore_dict.h"

#include "pycore_long.h"

#include "pycore_optimizer.h"

#include "pycore_object.h"

#include "pycore_dict.h"

#include "pycore_function.h"

#include "pycore_uop_metadata.h"

#include "pycore_uop_ids.h"

#include "pycore_range.h"

#include <stdarg.h>

#include <stdbool.h>

#include <stdint.h>

#include <stddef.h>

#ifdef Py_DEBUG

extern const char *_PyUOpName(int index);

extern void _PyUOpPrint(const _PyUOpInstruction *uop);

static const char *const DEBUG_ENV = "PYTHON_OPT_DEBUG";

static inline int get_lltrace(void) {

char *uop_debug = Py_GETENV(DEBUG_ENV);

int lltrace = 0;

if (uop_debug != NULL && *uop_debug >= '0') {

lltrace = *uop_debug - '0'; // TODO: Parse an int and all that

}

return lltrace;

}

#define DPRINTF(level, ...) \

if (get_lltrace() >= (level)) { printf(__VA_ARGS__); }

#define DPRINTF(level, ...)

static int

get_mutations(PyObject* dict) {

assert(PyDict_CheckExact(dict));

PyDictObject *d = (PyDictObject *)dict;

return (d->_ma_watcher_tag >> DICT_MAX_WATCHERS) & ((1 << DICT_WATCHED_MUTATION_BITS)-1);

}

static void

increment_mutations(PyObject* dict) {

assert(PyDict_CheckExact(dict));

PyDictObject *d = (PyDictObject *)dict;

d->_ma_watcher_tag += (1 << DICT_MAX_WATCHERS);

}

#define BUILTINS_WATCHER_ID 0

#define GLOBALS_WATCHER_ID 1

#define TYPE_WATCHER_ID 0

static int

globals_watcher_callback(PyDict_WatchEvent event, PyObject* dict,

PyObject* key, PyObject* new_value)

{

RARE_EVENT_STAT_INC(watched_globals_modification);

assert(get_mutations(dict) < _Py_MAX_ALLOWED_GLOBALS_MODIFICATIONS);

_Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), dict, 1);

increment_mutations(dict);

PyDict_Unwatch(GLOBALS_WATCHER_ID, dict);

return 0;

}

static int

type_watcher_callback(PyTypeObject* type)

{

_Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), type, 1);

PyType_Unwatch(TYPE_WATCHER_ID, (PyObject *)type);

return 0;

}

static PyObject *

convert_global_to_const(_PyUOpInstruction *inst, PyObject *obj)

{

assert(inst->opcode == _LOAD_GLOBAL_MODULE || inst->opcode == _LOAD_GLOBAL_BUILTINS || inst->opcode == _LOAD_ATTR_MODULE_FROM_KEYS);

assert(PyDict_CheckExact(obj));

PyDictObject *dict = (PyDictObject *)obj;

assert(dict->ma_keys->dk_kind == DICT_KEYS_UNICODE);

PyDictUnicodeEntry *entries = DK_UNICODE_ENTRIES(dict->ma_keys);

assert(inst->operand0 <= UINT16_MAX);

if ((int)inst->operand0 >= dict->ma_keys->dk_nentries) {

return NULL;

}

PyObject *res = entries[inst->operand0].me_value;

if (res == NULL) {

return NULL;

}

if (_Py_IsImmortal(res)) {

inst->opcode = _LOAD_CONST_INLINE_BORROW;

}

else {

inst->opcode = _LOAD_CONST_INLINE;

}

inst->operand0 = (uint64_t)res;

return res;

}

static int

incorrect_keys(_PyUOpInstruction *inst, PyObject *obj)

{

if (!PyDict_CheckExact(obj)) {

return 1;

}

PyDictObject *dict = (PyDictObject *)obj;

if (dict->ma_keys->dk_version != inst->operand0) {

return 1;

}

return 0;

}

static int

check_next_uop(_PyUOpInstruction *buffer, int size, int pc, uint16_t expected)

{

if (pc + 1 >= size) {

DPRINTF(1, "Cannot rewrite %s at pc %d: buffer too small\n",

_PyOpcode_uop_name[buffer[pc].opcode], pc);

return 0;

}

uint16_t next_opcode = buffer[pc + 1].opcode;

if (next_opcode != expected) {

DPRINTF(1,

"Cannot rewrite %s at pc %d: unexpected next opcode %s, "

"expected %s\n",

_PyOpcode_uop_name[buffer[pc].opcode], pc,

_PyOpcode_uop_name[next_opcode], _PyOpcode_uop_name[expected]);

return 0;

}

return 1;

}

static int

remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,

int buffer_size, _PyBloomFilter *dependencies)

{

PyInterpreterState *interp = _PyInterpreterState_GET();

PyObject *builtins = frame->f_builtins;

if (builtins != interp->builtins) {

OPT_STAT_INC(remove_globals_builtins_changed);

return 1;

}

PyObject *globals = frame->f_globals;

PyFunctionObject *function = _PyFrame_GetFunction(frame);

assert(PyFunction_Check(function));

assert(function->func_builtins == builtins);

assert(function->func_globals == globals);

uint32_t function_version = _PyFunction_GetVersionForCurrentState(function);

/* In order to treat globals as constants, we need to

/* These values represent stacks of booleans (one bool per bit).

uint32_t function_checked = 0;

uint32_t builtins_watched = 0;

uint32_t globals_watched = 0;

uint32_t prechecked_function_version = 0;

if (interp->dict_state.watchers[GLOBALS_WATCHER_ID] == NULL) {

interp->dict_state.watchers[GLOBALS_WATCHER_ID] = globals_watcher_callback;

}

if (interp->type_watchers[TYPE_WATCHER_ID] == NULL) {

interp->type_watchers[TYPE_WATCHER_ID] = type_watcher_callback;

}

for (int pc = 0; pc < buffer_size; pc++) {

_PyUOpInstruction *inst = &buffer[pc];

int opcode = inst->opcode;

switch(opcode) {

case _GUARD_BUILTINS_VERSION_PUSH_KEYS:

if (incorrect_keys(inst, builtins)) {

OPT_STAT_INC(remove_globals_incorrect_keys);

return 0;

}

if (interp->rare_events.builtin_dict >= _Py_MAX_ALLOWED_BUILTINS_MODIFICATIONS) {

continue;

}

if (!check_next_uop(buffer, buffer_size, pc,

_LOAD_GLOBAL_BUILTINS_FROM_KEYS)) {

continue;

}

if ((builtins_watched & 1) == 0) {

PyDict_Watch(BUILTINS_WATCHER_ID, builtins);

builtins_watched |= 1;

}

if (function_checked & 1) {

buffer[pc].opcode = NOP;

}

else {

buffer[pc].opcode = _CHECK_FUNCTION;

buffer[pc].operand0 = function_version;

function_checked |= 1;

}

// We're no longer pushing the builtins keys; rewrite the

// instruction that consumed the keys to load them from the

// frame.

buffer[pc + 1].opcode = _LOAD_GLOBAL_BUILTINS;

break;

case _GUARD_GLOBALS_VERSION:

case _GUARD_GLOBALS_VERSION_PUSH_KEYS:

if (incorrect_keys(inst, globals)) {

OPT_STAT_INC(remove_globals_incorrect_keys);

return 0;

}

uint64_t watched_mutations = get_mutations(globals);

if (watched_mutations >= _Py_MAX_ALLOWED_GLOBALS_MODIFICATIONS) {

continue;

}

if (opcode == _GUARD_GLOBALS_VERSION_PUSH_KEYS &&

!check_next_uop(buffer, buffer_size, pc,

_LOAD_GLOBAL_MODULE_FROM_KEYS)) {

continue;

}

if ((globals_watched & 1) == 0) {

PyDict_Watch(GLOBALS_WATCHER_ID, globals);

_Py_BloomFilter_Add(dependencies, globals);

globals_watched |= 1;

}

if (function_checked & 1) {

buffer[pc].opcode = NOP;

}

else {

buffer[pc].opcode = _CHECK_FUNCTION;

buffer[pc].operand0 = function_version;

function_checked |= 1;

}

if (opcode == _GUARD_GLOBALS_VERSION_PUSH_KEYS) {

// We're no longer pushing the globals keys; rewrite the

// instruction that consumed the keys to load them from the

// frame.

buffer[pc + 1].opcode = _LOAD_GLOBAL_MODULE;

}

break;

case _LOAD_GLOBAL_BUILTINS:

if (function_checked & globals_watched & builtins_watched & 1) {

convert_global_to_const(inst, builtins);

}

break;

case _LOAD_GLOBAL_MODULE:

if (function_checked & globals_watched & 1) {

convert_global_to_const(inst, globals);

}

break;

case _PUSH_FRAME:

{

builtins_watched <<= 1;

globals_watched <<= 1;

function_checked <<= 1;

uint64_t operand = buffer[pc].operand0;

if (operand == 0 || (operand & 1)) {

// It's either a code object or NULL, so bail

return 1;

}

PyFunctionObject *func = (PyFunctionObject *)operand;

if (func == NULL) {

return 1;

}

assert(PyFunction_Check(func));

function_version = func->func_version;

if (prechecked_function_version == function_version) {

function_checked |= 1;

}

prechecked_function_version = 0;

globals = func->func_globals;

builtins = func->func_builtins;

if (builtins != interp->builtins) {

OPT_STAT_INC(remove_globals_builtins_changed);

return 1;

}

break;

}

case _RETURN_VALUE:

{

builtins_watched >>= 1;

globals_watched >>= 1;

function_checked >>= 1;

uint64_t operand = buffer[pc].operand0;

if (operand == 0 || (operand & 1)) {

// It's either a code object or NULL, so bail

return 1;

}

PyFunctionObject *func = (PyFunctionObject *)operand;

if (func == NULL) {

return 1;

}

assert(PyFunction_Check(func));

function_version = func->func_version;

globals = func->func_globals;

builtins = func->func_builtins;

break;

}

case _CHECK_FUNCTION_EXACT_ARGS:

prechecked_function_version = (uint32_t)buffer[pc].operand0;

break;

default:

if (is_terminator(inst)) {

return 1;

}

break;

}

}

return 0;

}

#define STACK_LEVEL() ((int)(stack_pointer - ctx->frame->stack))

#define STACK_SIZE() ((int)(ctx->frame->stack_len))

#define WITHIN_STACK_BOUNDS() \

(STACK_LEVEL() >= 0 && STACK_LEVEL() <= STACK_SIZE())

#define GETLOCAL(idx) ((ctx->frame->locals[idx]))

#define REPLACE_OP(INST, OP, ARG, OPERAND) \

INST->opcode = OP; \

INST->oparg = ARG; \

INST->operand0 = OPERAND;

#define sym_is_not_null _Py_uop_sym_is_not_null

#define sym_is_const _Py_uop_sym_is_const

#define sym_get_const _Py_uop_sym_get_const

#define sym_new_unknown _Py_uop_sym_new_unknown

#define sym_new_not_null _Py_uop_sym_new_not_null

#define sym_new_type _Py_uop_sym_new_type

#define sym_is_null _Py_uop_sym_is_null

#define sym_new_const _Py_uop_sym_new_const

#define sym_new_null _Py_uop_sym_new_null

#define sym_has_type _Py_uop_sym_has_type

#define sym_get_type _Py_uop_sym_get_type

#define sym_matches_type _Py_uop_sym_matches_type

#define sym_matches_type_version _Py_uop_sym_matches_type_version

#define sym_set_null(SYM) _Py_uop_sym_set_null(ctx, SYM)

#define sym_set_non_null(SYM) _Py_uop_sym_set_non_null(ctx, SYM)

#define sym_set_type(SYM, TYPE) _Py_uop_sym_set_type(ctx, SYM, TYPE)

#define sym_set_type_version(SYM, VERSION) _Py_uop_sym_set_type_version(ctx, SYM, VERSION)

#define sym_set_const(SYM, CNST) _Py_uop_sym_set_const(ctx, SYM, CNST)

#define sym_is_bottom _Py_uop_sym_is_bottom

#define sym_truthiness _Py_uop_sym_truthiness

#define frame_new _Py_uop_frame_new

#define frame_pop _Py_uop_frame_pop

#define sym_new_tuple _Py_uop_sym_new_tuple

#define sym_tuple_getitem _Py_uop_sym_tuple_getitem

#define sym_tuple_length _Py_uop_sym_tuple_length

#define sym_is_immortal _Py_uop_sym_is_immortal

static int

optimize_to_bool(

JitOptSymbol **result_ptr)

{

if (sym_matches_type(value, &PyBool_Type)) {

REPLACE_OP(this_instr, _NOP, 0, 0);

*result_ptr = value;

return 1;

}

int truthiness = sym_truthiness(value);

if (truthiness >= 0) {

PyObject *load = truthiness ? Py_True : Py_False;

REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);

*result_ptr = sym_new_const(ctx, load);

return 1;

}

return 0;

}

static void

eliminate_pop_guard(_PyUOpInstruction *this_instr, bool exit)

{

REPLACE_OP(this_instr, _POP_TOP, 0, 0);

if (exit) {

REPLACE_OP((this_instr+1), _EXIT_TRACE, 0, 0);

this_instr[1].target = this_instr->target;

}

}

static PyCodeObject *

get_code(_PyUOpInstruction *op)

{

assert(op->opcode == _PUSH_FRAME || op->opcode == _RETURN_VALUE || op->opcode == _RETURN_GENERATOR);

PyCodeObject *co = NULL;

uint64_t operand = op->operand0;

if (operand == 0) {

return NULL;

}

if (operand & 1) {

co = (PyCodeObject *)(operand & ~1);

}

else {

PyFunctionObject *func = (PyFunctionObject *)operand;

assert(PyFunction_Check(func));

co = (PyCodeObject *)func->func_code;

}

assert(PyCode_Check(co));

return co;

}

static PyCodeObject *

get_code_with_logging(_PyUOpInstruction *op)

{

PyCodeObject *co = NULL;

uint64_t push_operand = op->operand0;

if (push_operand & 1) {

co = (PyCodeObject *)(push_operand & ~1);

DPRINTF(3, "code=%p ", co);

assert(PyCode_Check(co));

}

else {

PyFunctionObject *func = (PyFunctionObject *)push_operand;

DPRINTF(3, "func=%p ", func);

if (func == NULL) {

DPRINTF(3, "\n");

DPRINTF(1, "Missing function\n");

return NULL;

}

co = (PyCodeObject *)func->func_code;

DPRINTF(3, "code=%p ", co);

}

return co;

}

static int

optimize_uops(

)

{

JitOptContext context;

JitOptContext *ctx = &context;

uint32_t opcode = UINT16_MAX;

int curr_space = 0;

int max_space = 0;

_PyUOpInstruction *first_valid_check_stack = NULL;

_PyUOpInstruction *corresponding_check_stack = NULL;

_Py_uop_abstractcontext_init(ctx);

_Py_UOpsAbstractFrame *frame = _Py_uop_frame_new(ctx, co, curr_stacklen, NULL, 0);

if (frame == NULL) {

return -1;

}

ctx->curr_frame_depth++;

ctx->frame = frame;

ctx->done = false;

ctx->out_of_space = false;

ctx->contradiction = false;

_PyUOpInstruction *this_instr = NULL;

for (int i = 0; !ctx->done; i++) {

assert(i < trace_len);

this_instr = &trace[i];

int oparg = this_instr->oparg;

opcode = this_instr->opcode;

JitOptSymbol **stack_pointer = ctx->frame->stack_pointer;

#ifdef Py_DEBUG

if (get_lltrace() >= 3) {

printf("%4d abs: ", (int)(this_instr - trace));

_PyUOpPrint(this_instr);

printf(" ");

}

switch (opcode) {

#include "optimizer_cases.c.h"

default:

DPRINTF(1, "\nUnknown opcode in abstract interpreter\n");

Py_UNREACHABLE();

}

assert(ctx->frame != NULL);

DPRINTF(3, " stack_level %d\n", STACK_LEVEL());

ctx->frame->stack_pointer = stack_pointer;

assert(STACK_LEVEL() >= 0);

}

if (ctx->out_of_space) {

DPRINTF(3, "\n");

DPRINTF(1, "Out of space in abstract interpreter\n");

}

if (ctx->contradiction) {

// Attempted to push a "bottom" (contradiction) symbol onto the stack.

// This means that the abstract interpreter has hit unreachable code.

// We *could* generate an _EXIT_TRACE or _FATAL_ERROR here, but hitting

// bottom indicates type instability, so we are probably better off

// retrying later.

DPRINTF(3, "\n");

DPRINTF(1, "Hit bottom in abstract interpreter\n");

_Py_uop_abstractcontext_fini(ctx);

return 0;

}

/* Either reached the end or cannot optimize further, but there

_Py_uop_abstractcontext_fini(ctx);

if (first_valid_check_stack != NULL) {

assert(first_valid_check_stack->opcode == _CHECK_STACK_SPACE);

assert(max_space > 0);

assert(max_space <= INT_MAX);

assert(max_space <= INT32_MAX);

first_valid_check_stack->opcode = _CHECK_STACK_SPACE_OPERAND;

first_valid_check_stack->operand0 = max_space;

}

return trace_len;

error:

DPRINTF(3, "\n");

DPRINTF(1, "Encountered error in abstract interpreter\n");

if (opcode <= MAX_UOP_ID) {

OPT_ERROR_IN_OPCODE(opcode);

}

_Py_uop_abstractcontext_fini(ctx);

return -1;

}

static int

remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)

{

/* Remove _SET_IP and _CHECK_VALIDITY where possible.

int last_set_ip = -1;

bool may_have_escaped = true;

for (int pc = 0; pc < buffer_size; pc++) {

int opcode = buffer[pc].opcode;

switch (opcode) {

case _START_EXECUTOR:

may_have_escaped = false;

break;

case _SET_IP:

buffer[pc].opcode = _NOP;

last_set_ip = pc;

break;

case _CHECK_VALIDITY:

if (may_have_escaped) {

may_have_escaped = false;

}

else {

buffer[pc].opcode = _NOP;

}

break;

case _CHECK_VALIDITY_AND_SET_IP:

if (may_have_escaped) {

may_have_escaped = false;

buffer[pc].opcode = _CHECK_VALIDITY;

}

else {

buffer[pc].opcode = _NOP;

}

last_set_ip = pc;

break;

case _POP_TOP:

{

_PyUOpInstruction *last = &buffer[pc-1];

while (last->opcode == _NOP) {

last--;

}

if (last->opcode == _LOAD_CONST_INLINE ||

last->opcode == _LOAD_CONST_INLINE_BORROW ||

last->opcode == _LOAD_FAST ||

last->opcode == _COPY

) {

last->opcode = _NOP;

buffer[pc].opcode = _NOP;

}

if (last->opcode == _REPLACE_WITH_TRUE) {

last->opcode = _NOP;

}

break;

}

case _JUMP_TO_TOP:

case _EXIT_TRACE:

case _DYNAMIC_EXIT:

return pc + 1;

default:

{

/* _PUSH_FRAME doesn't escape or error, but it

bool needs_ip = opcode == _PUSH_FRAME;

if (_PyUop_Flags[opcode] & HAS_ESCAPES_FLAG) {

needs_ip = true;

may_have_escaped = true;

}

if (needs_ip && last_set_ip >= 0) {

if (buffer[last_set_ip].opcode == _CHECK_VALIDITY) {

buffer[last_set_ip].opcode = _CHECK_VALIDITY_AND_SET_IP;

}

else {

assert(buffer[last_set_ip].opcode == _NOP);

buffer[last_set_ip].opcode = _SET_IP;

}

last_set_ip = -1;

}

}

}

}

Py_UNREACHABLE();

}

_Py_uop_analyze_and_optimize(

)

{

OPT_STAT_INC(optimizer_attempts);

int err = remove_globals(frame, buffer, length, dependencies);

if (err <= 0) {

return err;

}

length = optimize_uops(

_PyFrame_GetCode(frame), buffer,

length, curr_stacklen, dependencies);

if (length <= 0) {

return length;

}

length = remove_unneeded_uops(buffer, length);

assert(length > 0);

OPT_STAT_INC(optimizer_successes);

return length;

}

#endif /* _Py_TIER2 */