//! Public trace entrypoint for `pyre`'s JIT portal.
//!
//! RPython MetaInterp._interpret() parity: trace_bytecode creates a
//! PyreMetaInterp and delegates to interpret(). The interpret loop
//! calls MIFrame::trace_code_step() for each bytecode, combining
//! concrete execution and symbolic IR recording.

use majit_metainterp::{TraceAction, TraceCtx};
use pyre_interpreter::CodeObject;

use crate::metainterp::{MetaInterpFrame, PyreMetaInterp};
use crate::state::PyreSym;

/// Trace an entire loop body starting at `start_pc`.
///
/// RPython MetaInterp._interpret() parity: creates a PyreMetaInterp
/// with a single frame and delegates to interpret(). The interpret
/// loop calls MIFrame::trace_code_step() for each bytecode, combining
/// concrete execution and symbolic IR recording.
pub fn trace_bytecode(
    ctx: &mut TraceCtx,
    sym: &mut PyreSym,
    code: &CodeObject,
    start_pc: usize,
    mut concrete_frame: Box<pyre_interpreter::pyframe::PyFrame>,
) -> (TraceAction, Box<pyre_interpreter::pyframe::PyFrame>) {
    // RPython MetaInterp._interpret() parity: root frame owns a concrete
    // PyFrame snapshot. MetaInterp drives both symbolic tracing AND
    // concrete execution — the interpreter does not run during tracing.
    concrete_frame.next_instr = start_pc;
    let cf_addr = &*concrete_frame as *const pyre_interpreter::pyframe::PyFrame as usize;
    let frame = MetaInterpFrame {
        sym: sym as *mut PyreSym,
        owned_sym: None,
        jitcode: code as *const CodeObject,
        pc: start_pc,
        greenkey: None,
        concrete_frame: cf_addr,
        owned_concrete_frame: Some(concrete_frame),
        parent_frames: Vec::new(),
        drop_frame_opref: None,
        caller_result_stack_idx: None,
        arg_state: pyre_interpreter::bytecode::OpArgState::default(),
    };

    let mut metainterp = PyreMetaInterp::new(code as *const CodeObject, std::ptr::null_mut());
    metainterp.framestack.push(frame);

    // RPython pyjitpl.py:2971-2973: register the initial merge point so
    // reached_loop_header recognizes the trace start backedge and closes
    // the loop instead of unrolling it as a first-visit inner loop.
    let start_key = crate::driver::make_green_key(code as *const CodeObject, start_pc);
    {
        let input_args: Vec<majit_ir::OpRef> = (0..ctx.num_inputs())
            .map(|i| majit_ir::OpRef(i as u32))
            .collect();
        let input_types = ctx.inputarg_types();
        ctx.add_merge_point(start_key, input_args, input_types, start_pc);
    }

    let action = metainterp.interpret(ctx);

    // Recover the root frame's owned_concrete_frame for writeback.
    let executed_frame = metainterp
        .framestack
        .pop()
        .and_then(|f| f.owned_concrete_frame);

    // RPython pyjitpl.py:3160: greenkey = original_boxes[:num_green_args]
    // — always uses the OUTER (start) key, never retargets to inner loop.
    // header_pc records the actual close position for cut_trace_from.
    match &action {
        TraceAction::CloseLoopWithArgs {
            loop_header_pc: Some(target_pc),
            ..
        } if *target_pc != start_pc => {
            let key = crate::driver::make_green_key(code as *const CodeObject, start_pc);
            ctx.set_green_key(key);
            ctx.header_pc = *target_pc;
            ctx.cut_inner_green_key = Some(crate::driver::make_green_key(
                code as *const CodeObject,
                *target_pc,
            ));
        }
        TraceAction::CloseLoop | TraceAction::CloseLoopWithArgs { .. } => {
            let key = crate::driver::make_green_key(code as *const CodeObject, start_pc);
            ctx.set_green_key(key);
            ctx.header_pc = start_pc;
        }
        _ => {}
    }

    // On abort, root frame may still be on the stack.
    let root_frame = if let Some(frame) = executed_frame {
        frame
    } else {
        metainterp
            .framestack
            .pop()
            .and_then(|frame| frame.owned_concrete_frame)
            .expect("trace_bytecode must return the root concrete frame")
    };
    (action, root_frame)
}

#[cfg(test)]
mod tests {
    use crate::metainterp::semantic_fallthrough_pc;
    use pyre_interpreter::bytecode::Instruction;
    use pyre_interpreter::compile_exec;
    use pyre_interpreter::decode_instruction_at;

    #[test]
    fn test_semantic_fallthrough_pc_skips_branch_trivia() {
        let mut source = String::from("def f(x, y):\n    if x < y:\n");
        for i in 0..400 {
            source.push_str(&format!("        z{i} = {i}\n"));
        }
        source.push_str("    return 0\n");
        source.push_str("f(1, 2)\n");

        let module = compile_exec(&source).expect("test code should compile");
        let code = module
            .constants
            .iter()
            .find_map(|constant| match constant {
                pyre_interpreter::ConstantData::Code { code } if code.obj_name.as_str() == "f" => {
                    Some((**code).clone())
                }
                _ => None,
            })
            .expect("test source should contain function code");

        let branch_pc = (0..code.instructions.len())
            .find(|&pc| {
                matches!(
                    decode_instruction_at(&code, pc),
                    Some((Instruction::PopJumpIfFalse { .. }, _))
                )
            })
            .expect("test bytecode should contain POP_JUMP_IF_FALSE");

        let fallthrough_pc = semantic_fallthrough_pc(&code, branch_pc);
        let fallthrough_instruction = decode_instruction_at(&code, fallthrough_pc)
            .map(|(instruction, _)| instruction)
            .expect("semantic fallthrough should decode");

        assert!(
            !matches!(
                fallthrough_instruction,
                Instruction::ExtendedArg
                    | Instruction::Resume { .. }
                    | Instruction::Nop
                    | Instruction::Cache
                    | Instruction::NotTaken
            ),
            "semantic fallthrough must skip bytecode trivia"
        );
    }
}
