Function get_jit_args

crates/vm/src/builtins/function/jit.rs:154–224

(
    func: &PyFunction,
    func_args: &FuncArgs,
    jitted_code: &'a CompiledCode,
    vm: &VirtualMachine,
)

/// `fill_locals_from_args` which will raise the actual exception if needed.
#[cfg(feature = "jit")]
pub(crate) fn get_jit_args<'a>(
    func: &PyFunction,
    func_args: &FuncArgs,
    jitted_code: &'a CompiledCode,
    vm: &VirtualMachine,
) -> Result<Args<'a>, ArgsError> {
    let mut jit_args = jitted_code.args_builder();
    let nargs = func_args.args.len();

    let code: &Py<PyCode> = &func.code;
    let arg_names = code.arg_names();
    let arg_count = code.arg_count;
    let posonlyarg_count = code.posonlyarg_count;

    if nargs > arg_count as usize || nargs < posonlyarg_count as usize {
        return Err(ArgsError::WrongNumberOfArgs);
    }

    // Add positional arguments
    for i in 0..nargs {
        jit_args.set(i, get_jit_value(vm, &func_args.args[i])?)?;
    }

    // Handle keyword arguments
    for (name, value) in &func_args.kwargs {
        let arg_pos =
            |args: &[&PyStrInterned], name: &str| args.iter().position(|arg| arg.as_str() == name);
        if let Some(arg_idx) = arg_pos(arg_names.args, name) {
            if jit_args.is_set(arg_idx) {
                return Err(ArgsError::ArgPassedMultipleTimes);
            }
            jit_args.set(arg_idx, get_jit_value(vm, value)?)?;
        } else if let Some(kwarg_idx) = arg_pos(arg_names.kwonlyargs, name) {
            let arg_idx = kwarg_idx + arg_count as usize;
            if jit_args.is_set(arg_idx) {
                return Err(ArgsError::ArgPassedMultipleTimes);
            }
            jit_args.set(arg_idx, get_jit_value(vm, value)?)?;
        } else {
            return Err(ArgsError::NotAKeywordArg);
        }
    }

    let (defaults, kwdefaults) = func.defaults_and_kwdefaults.lock().clone();

    // fill in positional defaults
    if let Some(defaults) = defaults {
        for (i, default) in defaults.iter().enumerate() {
            let arg_idx = i + arg_count as usize - defaults.len();
            if !jit_args.is_set(arg_idx) {
                jit_args.set(arg_idx, get_jit_value(vm, default)?)?;
            }
        }
    }

    // fill in keyword only defaults
    if let Some(kw_only_defaults) = kwdefaults {
        for (i, name) in arg_names.kwonlyargs.iter().enumerate() {