Function extract_annotations_from_annotate_code

crates/jit/tests/common.rs:80–172

Extract annotations from an annotate function's bytecode. The annotate function uses BUILD_MAP with key-value pairs loaded before it. Keys are parameter names (from LOAD_CONST), values are type names (from LOAD_NAME/LOAD_GLOBAL).

(code: &CodeObject)

/// The annotate function uses BUILD_MAP with key-value pairs loaded before it.
/// Keys are parameter names (from LOAD_CONST), values are type names (from LOAD_NAME/LOAD_GLOBAL).
fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap<Wtf8Buf, StackValue> {
    let mut annotations = HashMap::new();
    let mut stack: Vec<(bool, usize)> = Vec::new(); // (is_const, index)
    let mut op_arg_state = OpArgState::default();

    for &word in code.instructions.iter() {
        let (instruction, arg) = op_arg_state.get(word);

        match instruction {
            Instruction::LoadConst { consti } => {
                stack.push((true, consti.get(arg).as_usize()));
            }
            Instruction::LoadName { namei } => {
                stack.push((false, namei.get(arg) as usize));
            }
            Instruction::LoadGlobal { namei } => {
                stack.push((false, (namei.get(arg) >> 1) as usize));
            }
            Instruction::BuildMap { count } => {
                let count = count.get(arg) as usize;
                // Stack has key-value pairs in order: k1, v1, k2, v2, ...
                // So we need count * 2 items from the stack
                let start = stack.len().saturating_sub(count * 2);
                let pairs: Vec<_> = stack.drain(start..).collect();

                for chunk in pairs.chunks(2) {
                    if chunk.len() == 2 {
                        let (key_is_const, key_idx) = chunk[0];
                        let (val_is_const, val_idx) = chunk[1];

                        // Key should be a const string (parameter name)
                        if key_is_const
                            && let ConstantData::Str { value } =
                                &code.constants[(key_idx as u32).into()]
                        {
                            let param_name = value;
                            // Value can be a name (type ref) or a const string (forward ref)
                            let type_name = if val_is_const {
                                match code.constants.get(val_idx) {
                                    Some(ConstantData::Str { value }) => value
                                        .as_str()
                                        .map(|s| s.to_owned())
                                        .unwrap_or_else(|_| value.to_string_lossy().into_owned()),
                                    Some(other) => panic!(
                                        "Unsupported annotation const for '{:?}' at idx {}: {:?}",
                                        param_name, val_idx, other
                                    ),
                                    None => panic!(
                                        "Annotation const idx out of bounds for '{:?}': {} (len={})",
                                        param_name,
                                        val_idx,
                                        code.constants.len()
                                    ),
                                }
                            } else {
                                match code.names.get(val_idx) {
                                    Some(name) => name.clone(),
                                    None => panic!(