大模型辅助的Rust代码生成:从Prompt设计到安全代码的智能推导

cover

一、Rust代码生成的"安全鸿沟":为什么AI写的Rust代码十有八九编译不过

用大模型生成 Python 或 JavaScript 代码,通常能直接运行。但生成 Rust 代码?编译器大概率会拒绝——缺少生命周期标注、借用规则冲突、trait 约束不完整、错误处理缺失。这不是大模型不够聪明,而是 Rust 的类型系统和所有权模型对代码正确性的要求远超其他语言。一段"逻辑正确"的 Rust 代码,如果缺少 <'a> 标注或 + Send + 'static 约束,编译器会毫不留情地报错。

核心挑战在于:大模型生成代码时,倾向于产出"语义正确但类型不完整"的代码。解决思路不是让模型"更聪明",而是在 Prompt 中注入 Rust 的类型约束规则,并在生成后通过编译器反馈进行迭代修正。

二、Rust代码生成的约束注入与迭代修正架构

flowchart TB
    A[用户需求描述] --> B[约束注入 Prompt]
    B --> C[大模型代码生成]
    C --> D[编译器校验]
    D -->|编译通过| E[Clippy 检查]
    E -->|无警告| F[输出最终代码]
    D -->|编译失败| G[错误信息提取]
    E -->|有警告| G
    G --> H[错误上下文注入]
    H --> B

    subgraph 约束注入
        B1[所有权规则提示] --> B
        B2[生命周期标注模板] --> B
        B3[错误处理模式] --> B
        B4[trait 约束清单] --> B
    end

    subgraph 迭代修正
        G1[错误类型分类] --> G
        G2[相关代码定位] --> G
        G3[修正建议生成] --> H
    end

架构的核心是"约束注入 + 编译器反馈循环"。约束注入在 Prompt 中预置 Rust 的关键规则(所有权、生命周期、错误处理),减少模型生成不合规代码的概率。编译器反馈循环将编译错误信息注入下一轮 Prompt,让模型基于具体错误进行修正,而非盲目重试。通常 2-3 轮迭代即可通过编译。

三、Rust代码生成的工程实现

3.1 约束注入的Prompt模板

use serde::{Deserialize, Serialize};

/// Rust 代码生成的约束配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RustGenConfig {
    /// 是否要求所有函数包含错误处理
    pub require_error_handling: bool,
    /// 是否要求生命周期标注
    pub require_lifetime_annotations: bool,
    /// 是否要求 Send + Sync 约束
    pub require_thread_safe: bool,
    /// 最大迭代修正次数
    pub max_iterations: usize,
    /// 目标 Rust edition
    pub edition: String,
}

impl Default for RustGenConfig {
    fn default() -> Self {
        Self {
            require_error_handling: true,
            require_lifetime_annotations: true,
            require_thread_safe: true,
            max_iterations: 3,
            edition: "2021".to_string(),
        }
    }
}

/// 构造带约束的 Prompt
pub fn build_rust_prompt(requirement: &str, config: &RustGenConfig) -> String {
    let mut constraints = vec![
        "1. 所有代码必须通过 rustc 编译,不得使用 unsafe 块".to_string(),
        "2. 使用 Result<T, E> 处理所有可能失败的操作".to_string(),
        "3. 所有公共 API 必须有文档注释".to_string(),
    ];

    if config.require_error_handling {
        constraints.push(
            "4. 禁止使用 .unwrap() 或 .expect(),使用 ? 操作符传播错误".to_string()
        );
    }

    if config.require_lifetime_annotations {
        constraints.push(
            "5. 函数签名中的引用必须显式标注生命周期(不要依赖省略规则)".to_string()
        );
    }

    if config.require_thread_safe {
        constraints.push(
            "6. 所有跨线程传递的类型必须满足 Send + Sync".to_string()
        );
    }

    format!(
        "请根据以下需求生成 Rust 代码。\n\n\
         ## 需求\n{}\n\n\
         ## 约束\n{}\n\n\
         ## 输出格式\n\
         - 只输出可编译的完整代码\n\
         - 包含必要的 use 声明\n\
         - 包含 #[cfg(test)] 模块和至少一个测试用例",
        requirement,
        constraints.join("\n")
    )
}

3.2 编译器反馈循环

use std::process::Command;

/// 编译器反馈结果
#[derive(Debug, Clone)]
pub struct CompileFeedback {
    pub success: bool,
    pub errors: Vec<CompilerMessage>,
    pub warnings: Vec<CompilerMessage>,
}

#[derive(Debug, Clone)]
pub struct CompilerMessage {
    pub line: usize,
    pub column: usize,
    pub message: String,
    pub code: Option<String>,  // 错误码,如 E0499
    pub level: String,         // error / warning
}

/// 编译 Rust 代码并提取反馈
pub fn compile_and_extract_feedback(
    source_code: &str,
    project_dir: &str,
) -> CompileFeedback {
    // 将代码写入临时文件
    let src_path = format!("{}/src/main.rs", project_dir);
    std::fs::write(&src_path, source_code).expect("写入源文件失败");

    // 执行编译,捕获 JSON 格式的错误输出
    let output = Command::new("cargo")
        .args(["check", "--message-format=json"])
        .current_dir(project_dir)
        .output()
        .expect("执行 cargo check 失败");

    let mut errors = Vec::new();
    let mut warnings = Vec::new();

    // 解析编译器 JSON 输出
    for line in String::from_utf8_lossy(&output.stdout).lines() {
        if let Ok(msg) = serde_json::from_str::<serde_json::Value>(line) {
            if msg["reason"] == "compiler-message" {
                let level = msg["message"]["level"].as_str().unwrap_or("");
                let message_text = msg["message"]["message"].as_str().unwrap_or("");
                let code = msg["message"]["code"]["code"].as_str().map(String::from);

                // 提取行号和列号
                let (line, column) = msg["message"]["spans"]
                    .get(0)
                    .map(|span| {
                        (
                            span["line_start"].as_u64().unwrap_or(0) as usize,
                            span["column_start"].as_u64().unwrap_or(0) as usize,
                        )
                    })
                    .unwrap_or((0, 0));

                let cm = CompilerMessage {
                    line,
                    column,
                    message: message_text.to_string(),
                    code,
                    level: level.to_string(),
                };

                match level {
                    "error" => errors.push(cm),
                    "warning" => warnings.push(cm),
                    _ => {}
                }
            }
        }
    }

    CompileFeedback {
        success: errors.is_empty(),
        errors,
        warnings,
    }
}

/// 构造修正 Prompt:将编译错误注入下一轮生成
pub fn build_correction_prompt(
    original_code: &str,
    feedback: &CompileFeedback,
) -> String {
    let error_descriptions: Vec<String> = feedback.errors.iter()
        .map(|e| {
            format!(
                "行 {} 列 {}: [{}] {}",
                e.line, e.column,
                e.code.as_deref().unwrap_or("?"),
                e.message
            )
        })
        .collect();

    format!(
        "以下 Rust 代码存在编译错误,请修正:\n\n\
         ## 原始代码\n```rust\n{}\n```\n\n\
         ## 编译错误\n{}\n\n\
         ## 修正要求\n\
         - 只修改导致编译错误的代码\n\
         - 保持原有的逻辑意图\n\
         - 输出完整的修正后代码",
        original_code,
        error_descriptions.join("\n")
    )
}

3.3 完整的迭代生成流程

pub struct RustCodeGenerator {
    llm_client: LlmClient,
    config: RustGenConfig,
}

impl RustCodeGenerator {
    pub fn new(llm_client: LlmClient, config: RustGenConfig) -> Self {
        Self { llm_client, config }
    }

    /// 迭代生成 Rust 代码,直到编译通过或达到最大迭代次数
    pub async fn generate(&self, requirement: &str, project_dir: &str)
        -> Result<String>
    {
        // 第一轮:带约束的初始生成
        let prompt = build_rust_prompt(requirement, &self.config);
        let mut code = self.llm_client.generate_code(&prompt).await?;

        for iteration in 0..self.config.max_iterations {
            // 编译校验
            let feedback = compile_and_extract_feedback(&code, project_dir);

            if feedback.success && feedback.warnings.is_empty() {
                return Ok(code);
            }

            if feedback.success {
                // 编译通过但有警告,运行 Clippy 进一步检查
                let clippy_result = run_clippy(project_dir);
                if clippy_result.success {
                    return Ok(code);
                }
            }

            // 构造修正 Prompt
            let correction_prompt = build_correction_prompt(&code, &feedback);
            code = self.llm_client.generate_code(&correction_prompt).await?;
        }

        anyhow::bail!(
            "经过 {} 轮迭代仍未通过编译",
            self.config.max_iterations
        )
    }
}

四、AI代码生成的局限性与工程权衡

编译通过 ≠ 逻辑正确:编译器只检查类型安全和内存安全,不检查业务逻辑。AI 生成的代码可能编译通过但行为与需求不符——比如排序方向反了、边界条件遗漏、并发场景下的竞态条件。编译器反馈循环只能修复类型错误,无法修复逻辑错误。

迭代修正的 token 成本:每轮迭代需要将完整代码和错误信息发送给 LLM,一个 200 行的代码文件加错误信息约 3000-5000 tokens。3 轮迭代约 10000-15000 tokens,成本是单次生成的 3-5 倍。对于简单函数(< 30 行),手动编写可能比 AI 生成 + 迭代修正更高效。

生命周期标注的生成质量:大模型对复杂生命周期标注的生成准确率较低,尤其是涉及多个生命周期参数和约束关系时。约束注入可以提醒模型添加标注,但具体标注哪个生命周期参数仍依赖模型的理解能力。对于复杂泛型代码,建议先用 AI 生成骨架,再手动补全生命周期标注。

unsafe 代码的安全审查:AI 可能生成 unsafe 块来绕过借用检查器,这破坏了 Rust 的安全保证。约束注入中禁止 unsafe 可以减少这种情况,但模型可能通过 FFI 调用间接引入 unsafe。生成代码的安全审查不可省略。

五、总结

大模型辅助的 Rust 代码生成通过"约束注入 + 编译器反馈循环"的架构,将"生成→编译失败→手动修复"的流程自动化。约束注入在 Prompt 中预置所有权规则和错误处理要求,编译器反馈循环将具体错误注入下一轮 Prompt 进行定向修正。但 AI 生成代码的局限明显:编译通过不等于逻辑正确、迭代修正增加 token 成本、复杂生命周期标注质量低、unsafe 代码需要安全审查。落地建议:简单函数手动编写更高效,AI 适合生成模板代码和样板代码;生成后必须进行逻辑审查和安全审查;复杂生命周期标注建议手动补全;设置最大迭代次数(3 次)避免无限循环。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐