hakorune/plugins/nyash-python-parser-plugin/src/lib.rs

use pyo3::prelude::*;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug)]
struct ParseResult {
    success: bool,
    dump: Option<String>,
    counts: ParseCounts,
    unsupported: Vec<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct ParseCounts {
    total_nodes: usize,
    functions: usize,
    classes: usize,
    supported: usize,
    unsupported: usize,
}

/// FFI: プラグインABIバージョン
#[no_mangle]
pub extern "C" fn nyash_plugin_abi_version() -> u32 {
    1
}

/// FFI: プラグイン初期化
#[no_mangle]
pub extern "C" fn nyash_plugin_init() -> i32 {
    // Python初期化は pyo3 の auto-initialize が処理
    0
}

/// FFI: プラグインメソッド呼び出し（BID形式）
#[no_mangle]
pub extern "C" fn nyash_plugin_invoke(
    type_id: u32,
    method_id: u32,
    _instance_id: u32,
    args: *const u8,
    args_len: usize,
    result: *mut u8,
    result_len: *mut usize,
) -> i32 {
    const TYPE_ID_PARSER: u32 = 60;
    const METHOD_BIRTH: u32 = 0;
    const METHOD_PARSE: u32 = 1;
    const METHOD_FINI: u32 = u32::MAX;

    if type_id != TYPE_ID_PARSER {
        return -3; // NYB_E_INVALID_METHOD
    }

    match method_id {
        METHOD_BIRTH => {
            // インスタンスIDを返す
            unsafe {
                let instance_id = 1u32; // 簡易実装
                if *result_len < 4 {
                    *result_len = 4;
                    return -1; // NYB_E_SHORT_BUFFER
                }
                let out = std::slice::from_raw_parts_mut(result, *result_len);
                out[0..4].copy_from_slice(&instance_id.to_le_bytes());
                *result_len = 4;
            }
            0
        }
        METHOD_PARSE => {
            // 引数からコードを取得（TLV形式の文字列を期待）
            let code = unsafe {
                if args.is_null() || args_len < 4 {
                    // 引数なしの場合は環境変数から取得
                    std::env::var("NYASH_PY_CODE")
                        .unwrap_or_else(|_| "def main():\n    return 0".to_string())
                } else {
                    // TLVデコード（簡易版）
                    let buf = std::slice::from_raw_parts(args, args_len);
                    if args_len >= 8 {
                        let tag = u16::from_le_bytes([buf[0], buf[1]]);
                        let len = u16::from_le_bytes([buf[2], buf[3]]) as usize;
                        if tag == 6 && 4 + len <= args_len {
                            match std::str::from_utf8(&buf[4..4 + len]) {
                                Ok(s) => s.to_string(),
                                Err(_) => std::env::var("NYASH_PY_CODE")
                                    .unwrap_or_else(|_| "def main():\n    return 0".to_string()),
                            }
                        } else {
                            std::env::var("NYASH_PY_CODE")
                                .unwrap_or_else(|_| "def main():\n    return 0".to_string())
                        }
                    } else {
                        std::env::var("NYASH_PY_CODE")
                            .unwrap_or_else(|_| "def main():\n    return 0".to_string())
                    }
                }
            };

            // パース実行
            let parse_result = Python::with_gil(|py| parse_python_code(py, &code));

            // JSONにシリアライズ
            match serde_json::to_string(&parse_result) {
                Ok(json) => {
                    unsafe {
                        let bytes = json.as_bytes();
                        let need = 4 + bytes.len();
                        if *result_len < need {
                            *result_len = need;
                            return -1; // NYB_E_SHORT_BUFFER
                        }
                        let out = std::slice::from_raw_parts_mut(result, *result_len);
                        // TLVエンコード（tag=6:string）
                        out[0..2].copy_from_slice(&6u16.to_le_bytes());
                        out[2..4].copy_from_slice(&(bytes.len() as u16).to_le_bytes());
                        out[4..4 + bytes.len()].copy_from_slice(bytes);
                        *result_len = need;
                    }
                    0
                }
                Err(_) => -4, // エラー
            }
        }
        METHOD_FINI => 0,
        _ => -3, // NYB_E_INVALID_METHOD
    }
}

/// FFI: Pythonコードをパース
#[no_mangle]
pub extern "C" fn nyash_python_parse(
    code: *const std::os::raw::c_char,
) -> *mut std::os::raw::c_char {
    let code = unsafe {
        if code.is_null() {
            return std::ptr::null_mut();
        }
        match std::ffi::CStr::from_ptr(code).to_str() {
            Ok(s) => s,
            Err(_) => return std::ptr::null_mut(),
        }
    };

    let result = Python::with_gil(|py| parse_python_code(py, code));

    match serde_json::to_string(&result) {
        Ok(json) => {
            let c_str = std::ffi::CString::new(json).unwrap();
            c_str.into_raw()
        }
        Err(_) => std::ptr::null_mut(),
    }
}

/// FFI: 文字列解放
#[no_mangle]
pub extern "C" fn nyash_python_free_string(ptr: *mut std::os::raw::c_char) {
    if !ptr.is_null() {
        unsafe {
            let _ = std::ffi::CString::from_raw(ptr);
        }
    }
}

fn parse_python_code(py: Python, code: &str) -> ParseResult {
    let mut result = ParseResult {
        success: false,
        dump: None,
        counts: ParseCounts {
            total_nodes: 0,
            functions: 0,
            classes: 0,
            supported: 0,
            unsupported: 0,
        },
        unsupported: Vec::new(),
    };

    // Pythonのastモジュールをインポート
    let ast_module = match py.import_bound("ast") {
        Ok(m) => m,
        Err(e) => {
            result.dump = Some(format!("Failed to import ast module: {}", e));
            return result;
        }
    };

    // コードをパース
    let tree = match ast_module.call_method1("parse", (code,)) {
        Ok(t) => t,
        Err(e) => {
            result.dump = Some(format!("Parse error: {}", e));
            return result;
        }
    };

    // ASTをダンプ（文字列表現）
    if let Ok(dump_str) = ast_module.call_method1("dump", (&tree,)) {
        if let Ok(s) = dump_str.extract::<String>() {
            result.dump = Some(s);
        }
    }

    // ASTを解析してカウント
    analyze_ast(py, &tree, &mut result);

    result.success = true;
    result
}

fn analyze_ast(_py: Python, node: &Bound<'_, PyAny>, result: &mut ParseResult) {
    result.counts.total_nodes += 1;

    // ノードタイプを取得 - __class__.__name__ を使用
    if let Ok(class_obj) = node.getattr("__class__") {
        if let Ok(name_obj) = class_obj.getattr("__name__") {
            if let Ok(type_name) = name_obj.extract::<String>() {
                match type_name.as_str() {
                    "FunctionDef" => {
                        result.counts.functions += 1;
                        result.counts.supported += 1;
                    }
                    "AsyncFunctionDef" => {
                        result.counts.functions += 1;
                        result.counts.unsupported += 1;
                        result.unsupported.push("async function".to_string());
                    }
                    "ClassDef" => {
                        result.counts.classes += 1;
                        result.counts.unsupported += 1;
                        result.unsupported.push("class definition".to_string());
                    }
                    "For" | "While" | "If" => {
                        result.counts.supported += 1;
                    }
                    "Yield" | "YieldFrom" => {
                        result.counts.unsupported += 1;
                        result.unsupported.push("generator".to_string());
                    }
                    _ => {}
                }
            }
        }
    }

    // 子ノードを再帰的に解析
    // ast.walk() を使って全ノードを取得
    if let Ok(ast_module) = node.py().import_bound("ast") {
        if let Ok(walk_iter) = ast_module.call_method1("walk", (node,)) {
            if let Ok(nodes) = walk_iter.iter() {
                for child_result in nodes {
                    if let Ok(child) = child_result {
                        // 自分自身はスキップ（すでにカウント済み）
                        if !child.is(node) {
                            // 再帰的に解析（ただし walk は全ノードを返すので、
                            // 実際には再帰なしでフラットに処理される）
                            result.counts.total_nodes += 1;

                            if let Ok(class_obj) = child.getattr("__class__") {
                                if let Ok(name_obj) = class_obj.getattr("__name__") {
                                    if let Ok(type_name) = name_obj.extract::<String>() {
                                        match type_name.as_str() {
                                            "FunctionDef" => {
                                                result.counts.functions += 1;
                                                result.counts.supported += 1;
                                            }
                                            "AsyncFunctionDef" => {
                                                result.counts.functions += 1;
                                                result.counts.unsupported += 1;
                                                if !result
                                                    .unsupported
                                                    .contains(&"async function".to_string())
                                                {
                                                    result
                                                        .unsupported
                                                        .push("async function".to_string());
                                                }
                                            }
                                            "ClassDef" => {
                                                result.counts.classes += 1;
                                                result.counts.unsupported += 1;
                                                if !result
                                                    .unsupported
                                                    .contains(&"class definition".to_string())
                                                {
                                                    result
                                                        .unsupported
                                                        .push("class definition".to_string());
                                                }
                                            }
                                            "For" | "While" | "If" => {
                                                result.counts.supported += 1;
                                            }
                                            "Yield" | "YieldFrom" => {
                                                result.counts.unsupported += 1;
                                                if !result
                                                    .unsupported
                                                    .contains(&"generator".to_string())
                                                {
                                                    result
                                                        .unsupported
                                                        .push("generator".to_string());
                                                }
                                            }
                                            _ => {}
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_parse() {
        pyo3::prepare_freethreaded_python();

        Python::with_gil(|py| {
            let code = "def main():\n    return 0";
            let result = parse_python_code(py, code);

            assert!(result.success);
            assert_eq!(result.counts.functions, 1);
            assert_eq!(result.counts.supported, 1);
        });
    }
}