用 Rust 实现增量 Python 类型检查:Pyrefly 源码深度解析

1. 为什么需要它？

当团队的代码库从几千行膨胀到数十万行，传统的 Python 动态检查往往只能在运行时抛出错误，开发者只能在 CI 或手动运行 mypy 后才发现问题。此时两大痛点凸显：  

• 检查延迟：每次完整检查都要遍历全部文件，耗时从几秒飙到几分钟，阻塞编辑器的实时诊断。  
• 跨文件依赖不易增量：对一个模块的微小改动会迫使整个项目重新解析，导致 CI 资源浪费。

Pyrefly 把这两个痛点当作“漏水的屋顶”和“水管交叉的水库”。它用 增量图 记录每块代码的水位（类型信息），只在水位变化的管道上补漏，而不是把整栋楼的屋顶都重新铺。核心技术栈围绕 Rust 并发 + union‑find，在保证内存安全的前提下实现“闪电般”的增量检查和完整的 LSP 功能。

关键组件概览：  

2. 核心架构

系统的“骨架”是一条由 解析 → 构图 → 求解 → 服务 的流水线。可以把它想象成 一个自动化的邮局：  

1. 投递（解析）把 Python 源码切成信封（AST）。  
2. 分拣中心（增量图）记录每封信的目的地以及与其他信件的关联。  
3. 分拣机器人（求解器）依据地址合并相同收件人（类型变量统一），并在必要时把信件转寄（流类型细化）。  
4. 投递员（LSP）把最终的诊断结果送到开发者的 IDE。

关键类比：union‑find 如城市地下排水系统

想象每个类型变量是城市里的一口井，初始时井之间互不相连。union‑find 就像在地下铺设管道，把两口井连通以共享水位（统一类型）。路径压缩相当于在管道交叉点安装阀门，让后续的水流直接走最短路径，避免重复遍历。正因为这种“管道+阀门”的结构，Pyrefly 能在数十万行代码之间几乎瞬时完成类型统一。

数据流向（一个检查请求的旅程）

1. LSP 收到编辑器的“打开文件”事件，触发 pyrefly_lsp 创建一个 **Session**。  
2. Session 调用 pyrefly_python 把文件解析成 AST，并把每个模块的 Symbol 注册进 **pyrefly_graph**。  
3. 变更导致的节点标记为 脏，图的调度器只把这些脏节点和它们的依赖推送到 **pyrefly_solver**。  
4. solver 读取 **Variables**（union‑find 实例），在内部的 VariableNode 上执行路径压缩和合并，产生统一的 **Type**。  
5. 求解完成后，Session 把诊断信息、补全候选等序列化为 JSON，交给 lsp-server 按 LSP 协议回传给 IDE。

3. 源码级复盘

下面挑选 求解器（solver）、增量图（graph）、LSP 实现（lsp） 三个核心模块，分别从设计哲学、代码拆解、决策分析和潜在风险四个维度展开显微镜式观察。

3.1 pyrefly_solver：核心求解器

功能一句话总结
使用 union‑find 与气体（gas）控制的递归深度，实现跨模块的增量类型约束求解。

设计哲学

求解器的目标是 在保持极致吞吐的同时防止递归失控。Rust 的所有权模型天然防止数据竞争，项目在此基础上额外采用 RefCell 为每个变量提供内部可变性，使得并行路径压缩成为可能。

代码逻辑拆解

// 变量集合：每个 Var 对应一个内部可变的节点
#[derive(Debug, Default)]
struct Variables(SmallMap<Var, RefCell<VariableNode>>);

• 步骤 1：创建映射
  SmallMap 按需分配，避免全局锁；RefCell 让求解器在共享 Arc<Variables> 的情况下仍能就地修改节点。

/// 统一两个变量的根节点，返回合并后的根
fn union(&self, a: Var, b: Var) {
    let root_a = self.find(a);
    let root_b = self.find(b);
    if root_a != root_b {
        // 路径压缩：把较浅的根直接指向更深的根
        self.nodes[root_a].borrow_mut().parent = root_b;
    }
}

• 步骤 2：路径压缩
  每次 find 都递归抓取父节点并写回，等价于在地下管道里安装短路阀门，让后续水流直接流向最终汇合点。

/// 查找变量的根节点并进行路径压缩
fn find(&self, var: Var) -> Var {
    let node = self.nodes[&var].borrow();
    if node.parent == var { return var; }
    let root = self.find(node.parent);
    drop(node); // 释放只读借用
    self.nodes[&var].borrow_mut().parent = root; // 写回压缩路径
    root
}

• 步骤 3：气体控制递归深度
  为防止在极端递归（如深度嵌套的泛型）中栈溢出，求解器在每一次递归入口消耗 gas，耗尽即提前返回错误。

fn solve_type(&self, ty: &Type, gas: &mut Gas) -> Result<Type, SolveError> {
    gas.consume()?;               // 每进入一次递归消耗 1 单位
    match ty {
        Type::Union(inner) => self.solve_union(inner, gas),
        // …其他分支…
    }
}

决策分析

潜在风险

• **跨线程使用 RefCell**：如果未来将求解器迁移到多线程共享 Variables，未加 Sync 的 RefCell 会导致未定义行为。需在改造时引入 parking_lot::RwLock 或 crossbeam::atomic::AtomicCell。  
• 气体耗尽：在极端的递归类型（如无限展开的 TypeVar）上，gas 可能过早触发错误，导致误报。生产环境应将 Gas::new 参数调高或在配置中提供伸缩阈值。

3.2 pyrefly_graph：增量依赖图

功能一句话总结
维护模块之间的类型计算依赖，并在文件改动时精准定位需要重新求解的子图。

设计哲学

增量图的职责类似 “城市的供水调度中心”：每条管线记录哪块水库（模块）供给哪条支线（类型计算），当某座水库泄漏（代码改动）时，只关闭受影响的支线，不必关闭整座城市的供水。

代码逻辑拆解

/// Graph 节点，保存模块的计算结果和依赖集合
pub struct Node {
    /// 上游依赖的 Var 集合
    pub inputs: FxHashSet<Var>,
    /// 下游受影响的 Node 列表
    pub dependents: FxHashSet<NodeId>,
    /// 缓存的求解结果，使用 Arc 共享避免拷贝
    pub cached: Arc<Type>,
}

• 步骤 1：记录依赖
  当模块 A 被解析，A 的每个 Var 都会加入 inputs，并把 A 加入被它使用的模块的 dependents。

fn add_edge(&mut self, from: NodeId, to: NodeId, var: Var) {
    self.nodes[from].dependents.insert(to);
    self.nodes[to].inputs.insert(var);
}

• 步骤 2：脏标记传播
  文件修改触发 mark_dirty(node_id)，递归向下标记所有受影响的 dependents。

fn mark_dirty(&mut self, id: NodeId) {
    if self.nodes[id].dirty { return; }
    self.nodes[id].dirty = true;
    for &child in &self.nodes[id].dependents {
        self.mark_dirty(child);
    }
}

• 步骤 3：增量求解调度
  调度器遍历 dirty 节点，按拓扑顺序把它们送入求解器。每完成一次求解，cached 被更新，dirty 复位。

fn schedule(&mut self) -> impl Iterator<Item = NodeId> {
    self.nodes
        .iter()
        .enumerate()
        .filter(|(_, n)| n.dirty)
        .map(|(i, _)| NodeId(i))
}

决策分析

潜在风险

• 图循环检测：如果在 add_edge 时引入循环依赖，调度器的拓扑遍历会死循环。当前实现依赖 mark_dirty 的递归深度，极端情况下可能导致栈溢出。建议在 add_edge 前加入环检测或使用 Kahn 算法生成拓扑序。  
• 内存泄漏：Arc<Type> 若在大量短生命周期的模块之间频繁切换，可能导致旧的 Type 无法及时回收。生产中应监控 Arc::strong_count，必要时对 cached 使用 Weak 引用做惰性回收。

3.3 pyrefly_lsp：语言服务器桥梁

功能一句话总结
把求解器的诊断、补全等结果封装为 LSP 消息，实时推送给编辑器。

设计哲学

LSP 就像 “快递员的调度中心”：所有快递（诊断、补全）先进入中心排队，中心根据线路（文件路径）决定哪位快递员（IDE）先收货。这里的核心是 异步消息队列 与 并发状态管理，确保编辑器在高频键入时仍能得到低延迟响应。

代码逻辑拆解

// LSP 主入口，使用 Tokio 运行时
#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let (connection, io_threads) = lsp_server::Connection::stdio();
    let mut server = Server::new(connection);
    server.run().await?;
    io_threads.join()?;
    Ok(())
}

• 步骤 1：建立 Tokio 运行时
  #[tokio::main] 创建多线程调度器，保证网络 I/O 与内部求解任务可以并行。

impl Server {
    async fn run(&mut self) -> Result<(), anyhow::Error> {
        while let Some(msg) = self.conn.receiver.next().await {
            self.handle_message(msg).await?;
        }
        Ok(())
    }
}

• 步骤 2：消息分派
  handle_message 根据 LSP 方法把请求路由到对应的业务模块。对 textDocument/didChange，它会把改动写入 Session，触发增量图的脏标记。

async fn handle_message(&mut self, msg: Message) -> Result<(), anyhow::Error> {
    match msg.method.as_str() {
        "textDocument/didChange" => self.session.on_change(msg.params).await,
        "textDocument/completion" => self.session.on_completion(msg.params).await,
        // … 其它协议 …
        _ => {}
    }
    Ok(())
}

• 步骤 3：异步返回
  完成求解后，用 self.conn.sender.send(Notification::new(...)) 把诊断结果推送给 IDE，整个过程不阻塞主线程。

决策分析

潜在风险

• 任务饥饿：如果求解器的 Rayon 线程池被大量长时间的 solve_type 占满，LSP 的网络任务可能被延迟。建议在 tokio 中使用 spawn_blocking

项目地址: https://github.com/facebook/pyrefly