alef 0.63.0

/// Rewrite opaque impl-block methods whose return type is a capsule type.
///
/// The generic method generator emits `Ok({CapsuleType} { inner: Arc::new(result) })` for
/// methods returning capsule-configured types.  Because capsule types have no `#[pyclass]`
/// struct, that code does not compile.  This function replaces each such method with a
/// capsule-aware body that either calls `into_raw()` + `PyCapsule_New` (Capsule variant) or
/// constructs the Python object via the dependency capsule (ConstructFrom variant), mirroring
/// what `capsule::gen_capsule_function` does for free functions.
pub(super) fn rewrite_capsule_methods(
    impl_block: String,
    typ: &crate::core::ir::TypeDef,
    capsule_types: &std::collections::HashMap<String, crate::core::config::CapsuleTypeConfig>,
    error_converters: &[String],
) -> String {
    use crate::codegen::type_mapper::TypeMapper as _;
    use crate::core::ir::TypeRef;
    use heck::ToSnakeCase;

    let mut result = impl_block;

    for method in &typ.methods {
        let capsule_ret_name: Option<&str> = match &method.return_type {
            TypeRef::Named(n) if capsule_types.contains_key(n.as_str()) => Some(n.as_str()),
            TypeRef::Optional(inner) => {
                if let TypeRef::Named(n) = inner.as_ref() {
                    if capsule_types.contains_key(n.as_str()) {
                        Some(n.as_str())
                    } else {
                        None
                    }
                } else {
                    None
                }
            }
            _ => None,
        };

        let has_capsule_param = method
            .params
            .iter()
            .any(|p| matches!(&p.ty, TypeRef::Named(n) if capsule_types.contains_key(n.as_str())));

        if capsule_ret_name.is_none() && !has_capsule_param {
            continue;
        }

        let cfg = capsule_ret_name.map(|n| &capsule_types[n]);

        let old_sig_search = if let Some(ret_name) = capsule_ret_name {
            format!("-> PyResult<{ret_name}>")
        } else {
            format!("pub fn {}(", method.name)
        };

        if capsule_ret_name.is_some() && !result.contains(&old_sig_search) {
            continue;
        }

        let mut capsule_param_extract = String::new();
        let mut call_args_parts: Vec<String> = Vec::new();

        for p in &method.params {
            let param_is_capsule = matches!(&p.ty, TypeRef::Named(n) if capsule_types.contains_key(n.as_str()));

            if param_is_capsule {
                if let TypeRef::Named(capsule_name) = &p.ty {
                    capsule_param_extract.push_str(&crate::backends::pyo3::template_env::render(
                        "pyo3_capsule_param_extract.jinja",
                        minijinja::context! {
                            param_name => p.name.as_str(),
                            capsule_name => capsule_name,
                        },
                    ));
                    call_args_parts.push(p.name.clone());
                } else {
                    call_args_parts.push(p.name.clone());
                }
            } else {
                let needs_borrow = p.is_ref && matches!(p.ty, TypeRef::String | TypeRef::Char);
                if needs_borrow {
                    call_args_parts.push(format!("&{}", p.name));
                } else {
                    call_args_parts.push(p.name.clone());
                }
            }
        }
        let call_args_str = call_args_parts.join(", ");

        let mapper = crate::backends::pyo3::type_map::Pyo3Mapper::new();
        let mut sig_params = vec!["&self".to_string(), "py: pyo3::Python<'_>".to_string()];
        for p in &method.params {
            let param_type = if matches!(&p.ty, TypeRef::Named(n) if capsule_types.contains_key(n.as_str())) {
                "pyo3::Py<pyo3::PyAny>".to_string()
            } else {
                mapper.map_type(&p.ty)
            };
            sig_params.push(format!("{}: {}", p.name, param_type));
        }

        // Build the #[pyo3(signature = (...))] attribute (skipped when there are no params).
        let sig_attr = if method.params.is_empty() {
            String::new()
        } else {
            let names = method
                .params
                .iter()
                .map(|p| p.name.as_str())
                .collect::<Vec<_>>()
                .join(", ");
            format!("    #[pyo3(signature = ({names}))]\n")
        };

        let core_call = format!("self.inner.{}({})", method.name, call_args_str);

        let err_map_suffix = if method.error_type.is_some() {
            let converter = method
                .error_type
                .as_ref()
                .and_then(|et| {
                    let short = et.split("::").last().unwrap_or(et.as_str());
                    let candidate = format!("{}_to_py_err", short.to_snake_case());
                    if error_converters.iter().any(|c| c == &candidate) {
                        Some(candidate)
                    } else {
                        None
                    }
                })
                .unwrap_or_else(|| "|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())".to_string());
            format!(".map_err({converter})?")
        } else {
            String::new()
        };

        let params_str = sig_params.join(", ");
        let method_name = &method.name;

        let new_body = if cfg.is_none() {
            let return_annotation = if matches!(method.return_type, TypeRef::Unit) {
                "".to_string()
            } else {
                format!(" -> PyResult<{}>", mapper.map_type(&method.return_type))
            };
            format!(
                r#"    {sig_attr}    #[allow(clippy::missing_errors_doc)]
    pub fn {method_name}({params_str}){return_annotation} {{
{capsule_param_extract}        {core_call}{err_map_suffix}
    }}"#,
            )
        } else if let Some(cfg) = cfg {
            match cfg {
                crate::core::config::CapsuleTypeConfig::Capsule(capsule_name_str) => {
                    let capsule_cstr = capsule_name_str.replace('.', "_").to_ascii_uppercase();
                    let construct = match capsule_name_str.rsplit_once('.') {
                    Some((module_path, class_name)) => format!(
                        r#"        // SAFETY: capsule_ptr is a valid, non-null Python object pointer we just created above.
        let _capsule_obj = unsafe {{ pyo3::Bound::from_owned_ptr(py, capsule_ptr) }};
        let _ts_mod = py.import("{module_path}")?;
        let _cls = _ts_mod.getattr("{class_name}")?;
        Ok(_cls.call1((_capsule_obj,))?.unbind())"#,
                    ),
                    None => {
                        "        // SAFETY: capsule_ptr is a valid, non-null Python object pointer we just created above.\n        Ok(unsafe { pyo3::Bound::from_owned_ptr(py, capsule_ptr) }.unbind())".to_string()
                    }
                };
                    format!(
                        r#"    {sig_attr}    #[allow(clippy::missing_errors_doc)]
    pub fn {method_name}({params_str}) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {{
        const {capsule_cstr}_NAME: &::std::ffi::CStr = c"{capsule_name_str}";
{capsule_param_extract}        let result = {core_call}{err_map_suffix};
        let raw_ptr = result.into_raw();
        // SAFETY: raw_ptr is a valid pointer derived from into_raw() on a value with program lifetime.
        let capsule_ptr = unsafe {{ pyo3::ffi::PyCapsule_New(raw_ptr as *mut _, {capsule_cstr}_NAME.as_ptr(), None) }};
        if capsule_ptr.is_null() {{
            return Err(pyo3::exceptions::PyRuntimeError::new_err("Failed to create PyCapsule"));
        }}
{construct}
    }}"#,
                    )
                }
                crate::core::config::CapsuleTypeConfig::ConstructFrom {
                    python_type,
                    construct_from,
                } => {
                    let dep_snake = construct_from.to_snake_case();
                    let first_str_param = method.params.iter().find(|p| matches!(p.ty, TypeRef::String));

                    if let Some(sp) = first_str_param {
                        let dep_expr = format!("get_{dep_snake}(py, {}.clone())?.bind(py).clone()", sp.name);
                        if let Some((module_path, class_name)) = python_type.rsplit_once('.') {
                            format!(
                                r#"    {sig_attr}    #[allow(clippy::missing_errors_doc)]
    pub fn {method_name}({params_str}) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {{
        // Construct {python_type} via Python-side factory.
        let _dep = {dep_expr};
        let _ts_mod = py.import("{module_path}")?;
        let _cls = _ts_mod.getattr("{class_name}")?;
        Ok(_cls.call1((_dep,))?.unbind())
    }}"#,
                            )
                        } else {
                            format!(
                                r#"    {sig_attr}    #[allow(clippy::missing_errors_doc)]
    pub fn {method_name}({params_str}) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {{
        // Construct {python_type} via Python-side factory.
        let _dep = {dep_expr};
        let _cls = py.eval(c"{python_type}", None, None)?;
        Ok(_cls.call1((_dep,))?.unbind())
    }}"#,
                            )
                        }
                    } else {
                        format!(
                            r#"    {sig_attr}    #[allow(clippy::missing_errors_doc)]
    pub fn {method_name}({params_str}) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {{
        // Unsupported: no parameter found to obtain the {construct_from} capsule.
        Err(pyo3::exceptions::PyRuntimeError::new_err("capsule_types construct_from: no parameter available to obtain {construct_from} capsule"))
    }}"#,
                        )
                    }
                }
            }
        } else {
            unreachable!("Method capsule config should be present when cfg.is_none() is false.");
        };

        let method_start_marker = format!("pub fn {method_name}(");
        if let Some(start_idx) = result.find(&method_start_marker) {
            let attr_start = find_method_attrs_start(&result, start_idx);
            if let Some(end_idx) = find_method_end(&result, start_idx) {
                result = format!("{}{}{}", &result[..attr_start], new_body, &result[end_idx..]);
            }
        }
    }

    result
}

/// Returns true when `line` (trimmed) consists entirely of `#[…]` attribute patterns and
/// intervening whitespace — i.e. it contains no non-attribute tokens such as `impl Foo {`.
///
/// This correctly handles:
/// - A single attribute: `#[pyo3(signature = (name))]`  → true
/// - Multiple attributes on one line: `#[allow(dead_code)]  #[pyo3(get)]`  → true
/// - A block-attr + impl opener on one line: `#[pymethods]impl Foo {`  → false
fn is_method_attr_line(line: &str) -> bool {
    let mut rest = line.trim();
    if rest.is_empty() {
        return false;
    }
    loop {
        rest = rest.trim_start();
        if rest.is_empty() {
            return true;
        }
        if !rest.starts_with("#[") {
            return false;
        }
        // Consume the `#[…]` span, respecting nested brackets.
        let mut depth = 0usize;
        let mut consumed = 0usize;
        let mut found_close = false;
        for (i, ch) in rest.char_indices() {
            match ch {
                '[' => depth += 1,
                ']' => {
                    depth = depth.saturating_sub(1);
                    if depth == 0 {
                        consumed = i + 1;
                        found_close = true;
                        break;
                    }
                }
                _ => {}
            }
        }
        if !found_close {
            return false;
        }
        rest = &rest[consumed..];
    }
}

/// Find the byte index of the start of the attribute block that precedes the `pub fn` at
/// `fn_idx`.  Walks backward line-by-line past `#[…]` attribute lines and blank lines.
/// Stops as soon as it encounters a line that is not purely made of `#[…]` attributes
/// (e.g. `#[pymethods]impl Foo {`).  Returns the byte index of the first character of the
/// first method-attribute line (or `fn_idx` when there are none).
fn find_method_attrs_start(code: &str, fn_idx: usize) -> usize {
    let before = &code[..fn_idx];
    let line_starts: Vec<usize> = std::iter::once(0)
        .chain(before.match_indices('\n').map(|(i, _)| i + 1))
        .collect();

    let mut attr_start_byte = fn_idx;
    for &line_byte_start in line_starts.iter().rev() {
        let line = &before[line_byte_start..before.len().min(attr_start_byte)];
        let trimmed = line.trim_end_matches('\n').trim();
        if trimmed.is_empty() || is_method_attr_line(trimmed) {
            attr_start_byte = line_byte_start;
        } else {
            break;
        }
    }
    attr_start_byte
}

/// Find the byte index just after the closing `}` of a Rust method block whose `pub fn`
/// starts at byte `fn_idx` in `code`.
fn find_method_end(code: &str, fn_idx: usize) -> Option<usize> {
    let slice = &code[fn_idx..];
    let mut depth = 0usize;
    let mut found_open = false;
    let mut byte_offset = 0usize;
    for ch in slice.chars() {
        match ch {
            '{' => {
                depth += 1;
                found_open = true;
            }
            '}' if found_open => {
                depth = depth.saturating_sub(1);
                if depth == 0 {
                    byte_offset += ch.len_utf8();
                    return Some(fn_idx + byte_offset);
                }
            }
            _ => {}
        }
        byte_offset += ch.len_utf8();
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::config::CapsuleTypeConfig;
    use crate::core::ir::{MethodDef, ParamDef, TypeDef, TypeRef};
    use std::collections::HashMap;

    fn capsule_map(entries: &[(&str, CapsuleTypeConfig)]) -> HashMap<String, CapsuleTypeConfig> {
        entries.iter().map(|(k, v)| (k.to_string(), v.clone())).collect()
    }

    const IMPL_BLOCK: &str = r#"impl PyParser {
    pub fn make_parser(&self, py: pyo3::Python<'_>) -> PyResult<Parser> {
        todo!()
    }
}
"#;

    /// When a `ConstructFrom` capsule method has no `TypeRef::String` parameter to look its
    /// dependency up through, `rewrite_capsule_methods` must emit a catchable `PyRuntimeError`
    /// rather than splicing an `unreachable!()` into the generated Python-extension source.
    #[test]
    fn rewrite_capsule_methods_construct_from_no_dependency_param_avoids_unreachable() {
        let method = MethodDef {
            name: "make_parser".to_string(),
            params: vec![],
            return_type: TypeRef::Named("Parser".to_string()),
            ..Default::default()
        };
        let typ = TypeDef {
            name: "Parser".to_string(),
            methods: vec![method],
            ..Default::default()
        };
        let capsules = capsule_map(&[(
            "Parser",
            CapsuleTypeConfig::ConstructFrom {
                python_type: "tree_sitter.Parser".to_string(),
                construct_from: "Language".to_string(),
            },
        )]);

        let out = rewrite_capsule_methods(IMPL_BLOCK.to_string(), &typ, &capsules, &[]);

        assert!(
            !out.contains("unreachable!()"),
            "generated source must not contain unreachable!(): {out}"
        );
        assert!(
            out.contains("no parameter available to obtain Language capsule"),
            "expected loud-failure marker in output: {out}"
        );
    }

    /// Sibling positive control: when a `TypeRef::String` parameter IS present, codegen still
    /// emits the real `get_{dep_snake}(py, ...)` lookup call, unchanged by the fix above.
    #[test]
    fn rewrite_capsule_methods_construct_from_uses_string_param_dependency_lookup() {
        let method = MethodDef {
            name: "make_parser".to_string(),
            params: vec![ParamDef {
                name: "source".to_string(),
                ty: TypeRef::String,
                ..Default::default()
            }],
            return_type: TypeRef::Named("Parser".to_string()),
            ..Default::default()
        };
        let typ = TypeDef {
            name: "Parser".to_string(),
            methods: vec![method],
            ..Default::default()
        };
        let capsules = capsule_map(&[(
            "Parser",
            CapsuleTypeConfig::ConstructFrom {
                python_type: "tree_sitter.Parser".to_string(),
                construct_from: "Language".to_string(),
            },
        )]);

        let out = rewrite_capsule_methods(IMPL_BLOCK.to_string(), &typ, &capsules, &[]);

        assert!(
            !out.contains("unreachable!()"),
            "generated source must not contain unreachable!(): {out}"
        );
        assert!(
            out.contains("get_language(py, source.clone())?.bind(py).clone()"),
            "expected real dependency lookup call in output: {out}"
        );
        assert!(
            out.contains("py.import(\"tree_sitter\")?"),
            "expected python-side factory import in output: {out}"
        );
    }
}