Rust로 빠르고 안전한 타입추론 구현 Rc와 RefCell의 실전 설계

Rust로 빠르고 안전한 타입추론 구현 Rc와 RefCell의 실전 설계

간단하고 빠르면서도 안전한 타입추론을 Rust에서 구현해보면 꽤 놀랍더라고요.

핵심은 'Rc'와 'RefCell' 조합인데요.

레퍼런스 카운팅으로 값 복사를 빠르게 처리하고, 내부 가변성으로 대입을 유연하게 처리하는 그림이 깔끔하죠.

이번 글은 부울과 함수형만 있는 아주 작은 타입 시스템에서 시작해요.

여기에 타입변수를 얹고, 'occurs-check'와 'unify'를 구현한 뒤, 성능을 제대로 끌어올리는 'union-by-rank'까지 다뤄보려는 거예요.

조금씩 쌓아가면 코드가 의외로 간단해지거든요.

타입과 타입변수의 정의

먼저 타입을 정의해볼 텐데요.

부울 'Bool'과 함수형 'T → T'만 있는 최소한의 타입부터 시작하죠.

use std::cell::RefCell;
use std::rc::Rc;

#[derive(Clone)]
enum Ty {
    Bool,
    Func(Rc<Ty>, Rc<Ty>),
}

여기서 박스 대신 'Rc'를 쓴 이유가 있는데요.

곧 타입추론에서 '타입변수에 어떤 타입을 대입'하는 동작이 반복되는데, 이때 'Rc'는 참조 카운트만 올려서 값 복사 비용을 사실상 제로에 가깝게 만들죠.

이제 타입변수도 모델링해야 하는데요.

타입변수는 아직 아무 것도 대입되지 않았거나, 특정 타입이 대입된 두 상태 중 하나로 보면 돼요.

enum Var {
    Assigned(Rc<Ty>),
    Unassigned,
}

마지막으로 타입과 타입변수를 같은 층위에서 다루기 위해 'Ty'에 타입변수 케이스를 추가해요.

대입이 일어나면 내부 상태가 바뀌어야 하므로 'RefCell'로 가변성을 열어두는 게 포인트죠.

#[derive(Clone)]
enum Ty {
    Bool,
    Func(Rc<Ty>, Rc<Ty>),
    Var(RefCell<Var>),
}

여기까지가 핵심 데이터 구조인데요.

이제부터는 이 구조 위에 타입추론의 심장인 'occurs-check'와 'unify'를 얹어보죠.

occurs-check 구현 contains

무한 타입을 막으려면 'occurs-check'가 꼭 필요한데요.

대표적으로 T = α → α일 때 α에 T를 그대로 집어넣어버리면 α가 자기 자신을 품는 무한 타입이 되어버리죠.

그래서 특정 타입 안에 어떤 타입변수가 포함되어 있는지 판별하는 'contains'가 필수예요.

impl Ty {
    // self가 내부에 var를 포함하면 true를 반환
    fn contains(self: &Rc<Ty>, var: &RefCell<Var>) -> bool {
        match self.as_ref() {
            Ty::Bool => false,
            Ty::Func(param, ret) => param.contains(var) || ret.contains(var),
            Ty::Var(self_var) => match *self_var.borrow() {
                Var::Assigned(ref ty) => ty.contains(var),
                Var::Unassigned => self_var.as_ptr() == var.as_ptr(),
            },
        }
    }
}

여기서 'self.as_ref()'는 'Rc'를 '&T'로 가볍게 보는 도우미인데요.

기호가 거슬린다면 '&**self'로 같은 효과를 낼 수 있고, 'Deref'를 끌어와서 더 간결하게 적는 것도 가능하죠.

포인트는 'Var::Assigned'면 배정된 타입을 타고 내려가 재귀 검사하고, 'Var::Unassigned'면 단순히 포인터 동일성으로 같은 변수인지 확인하는 부분이에요.

이 비교는 'RefCell::as_ptr'이 '*mut Var'를 내주기 때문에 그대로 쓰면 간단해지죠.

unify 구현 기본형

이제 진짜 메인 디시인 'unify'를 만져볼 차례인데요.

두 타입 'T_left'와 'T_right'를 받아서 둘을 같게 만들 수 있으면 true, 아니면 false를 돌려주는 함수로 생각하면 돼요.

케이스를 나눠 쭉 정의하면 의외로 읽히는 코드가 나옵니다.

fn unify(left: &Rc<Ty>, right: &Rc<Ty>) -> bool {
    match (left.as_ref(), right.as_ref()) {
        (Ty::Bool, Ty::Bool) => true,
        (Ty::Func(left_param, left_ret), Ty::Func(right_param, right_ret)) => {
            unify(left_param, right_param) && unify(left_ret, right_ret)
        }
        (Ty::Var(left_var), Ty::Var(right_var)) => {
            if let Var::Assigned(ref l) = *left_var.borrow() {
                return unify(l, right);
            }
            if let Var::Assigned(ref r) = *right_var.borrow() {
                return unify(left, r);
            }
            if left_var.as_ptr() != right_var.as_ptr() {
                *left_var.borrow_mut() = Var::Assigned(right.clone());
            }
            true
        }
        (Ty::Var(left_var), _) => {
            if let Var::Assigned(ref l) = *left_var.borrow() {
                unify(l, right)
            } else if right.contains(left_var) {
                false
            } else {
                *left_var.borrow_mut() = Var::Assigned(right.clone());
                true
            }
        }
        (_, Ty::Var(right_var)) => {
            if let Var::Assigned(ref r) = *right_var.borrow() {
                unify(left, r)
            } else if left.contains(right_var) {
                false
            } else {
                *right_var.borrow_mut() = Var::Assigned(left.clone());
                true
            }
        }
        _ => false,
    }
}

여기서 주의할 게 하나 있는데요.

'if let'으로 'borrow()'한 레퍼런스가 스코프 밖으로 빨리 떨어져야 뒤이어 'borrow_mut()'가 안전하게 열리죠.

Rust 2024 에디션에서는 위 코드가 잘 동작하지만, 2021에서는 'RefCell already borrowed'로 패닉이 날 수 있어요.

그럴 땐 중간 값을 명시적으로 클론해두고 레퍼런스를 즉시 버리거나, 블록 스코프로 빌림 기간을 끊어주면 됩니다.

핵심 흐름은 단순해요.

동일 구조끼리 재귀 유니파이를 시도하고, 한쪽이 변수면 'occurs-check'로 무한 타입을 막은 뒤 대입하는 식이죠.

여기까지면 기본 기능은 갖춰졌는데, 아직 성능 면에서 구멍이 하나 남아있어요.

성능 병목과 union-by-rank 최적화

문제는 변수 사슬이 길어질 때인데요.

예를 들어 α1 = α2, α2 = α3, ..., αn-1 = αn을 순서대로 유니파이하면 α1을 해석하려고 할 때 αn까지 'n-1'번 참조를 따라가야 하죠.

이건 연쇄 대입 구조가 그대로 남기 때문인데, 유니온 파인드에서 쓰는 'rank' 아이디어로 손쉽게 줄일 수 있어요.

아이디어는 아직 대입되지 않은 변수에 '해당 노드에서 루트까지 필요한 최대 deref 횟수' 비슷한 랭크를 붙이는 거예요.

둘 다 미할당 변수일 때 랭크가 작은 쪽을 큰 쪽 밑으로 붙이면 체인의 높이가 잘 억제되죠.

이를 위해 'Var::Unassigned'에 순위를 넣어 살짝 타입을 바꿉니다.

use std::cmp::Ordering;

enum Var {
    Assigned(Rc<Ty>),
    Unassigned(u32), // rank
}

그리고 Var-Var 케이스에서 둘 다 미할당일 때 랭크를 비교해 병합 순서를 정해요.

랭크가 같으면 한쪽을 다른 쪽에 붙이고, 붙여진 쪽의 랭크를 1 올려 밸런스를 유지하죠.

fn unify(left: &Rc<Ty>, right: &Rc<Ty>) -> bool {
    match (left.as_ref(), right.as_ref()) {
        (Ty::Bool, Ty::Bool) => true,
        (Ty::Func(lp, lr), Ty::Func(rp, rr)) => {
            unify(lp, rp) && unify(lr, rr)
        }
        (Ty::Var(left_var), Ty::Var(right_var)) => {
            let left_rank = match *left_var.borrow() {
                Var::Assigned(ref l) => return unify(l, right),
                Var::Unassigned(rank) => rank,
            };
            let right_rank = match *right_var.borrow() {
                Var::Assigned(ref r) => return unify(left, r),
                Var::Unassigned(rank) => rank,
            };
            if left_var.as_ptr() != right_var.as_ptr() {
                match left_rank.cmp(&right_rank) {
                    Ordering::Greater => {
                        *right_var.borrow_mut() = Var::Assigned(left.clone());
                    }
                    Ordering::Less => {
                        *left_var.borrow_mut() = Var::Assigned(right.clone());
                    }
                    Ordering::Equal => {
                        *left_var.borrow_mut() = Var::Assigned(right.clone());
                        *right_var.borrow_mut() = Var::Unassigned(right_rank + 1);
                    }
                }
            }
            true
        }
        (Ty::Var(left_var), _) => {
            if let Var::Assigned(ref l) = *left_var.borrow() {
                unify(l, right)
            } else if right.contains(left_var) {
                false
            } else {
                *left_var.borrow_mut() = Var::Assigned(right.clone());
                true
            }
        }
        (_, Ty::Var(right_var)) => {
            if let Var::Assigned(ref r) = *right_var.borrow() {
                unify(left, r)
            } else if left.contains(right_var) {
                false
            } else {
                *right_var.borrow_mut() = Var::Assigned(left.clone());
                true
            }
        }
        _ => false,
    }
}

이렇게만 바꿔도 deref 깊이가 'O(log n)'으로 내려가는데요.

사슬이 길어질수록 차이가 눈에 띄게 벌어지죠.

'Rc'와 'RefCell'이 각자의 역할을 나눠 가진 덕분에, 빌림과 대입을 섞어 쓰는 코드가 의외로 깨끗하게 정리되는 것도 장점이에요.

덧붙이자면 경로 압축까지 하면 평균 거의 상수 시간에 가까워지는데요.

이 글에서는 읽기 간결성을 위해 랭크만 적용했지만, 필요하다면 'find' 헬퍼를 두고 'Assigned(Var(...))' 사슬을 따라가며 한 번에 루트로 당겨 쓰면 더 짧아지죠.

실제로 써보기 타입추론 엔진 스케치

이제 유니파이어가 준비됐으니, 간단한 람다식에 타입을 붙이는 흐름을 스케치해볼게요.

변수와 람다, 함수 적용, 그리고 부울과 조건식만 있는 작은 언어면 충분하죠.

use std::collections::HashMap;

#[derive(Clone)]
enum Expr {
    Var(String),
    Lam(String, Box<Expr>),    // λx. e
    App(Box<Expr>, Box<Expr>), // e1 e2
    Bool(bool),
    If(Box<Expr>, Box<Expr>, Box<Expr>),
}

struct Ctx {
    env: HashMap<String, Rc<Ty>>,
    supply: u32,
}

impl Ctx {
    fn new() -> Self {
        Self { env: HashMap::new(), supply: 0 }
    }

    fn fresh(&mut self) -> Rc<Ty> {
        let v = Ty::Var(RefCell::new(Var::Unassigned(0)));
        self.supply += 1;
        Rc::new(v)
    }

    fn lookup(&self, x: &str) -> Option<Rc<Ty>> {
        self.env.get(x).cloned()
    }

    fn with_binding<F, R>(&mut self, x: String, ty: Rc<Ty>, f: F) -> R
    where
        F: FnOnce(&mut Ctx) -> R,
    {
        let old = self.env.insert(x, ty);
        let r = f(self);
        if let Some(prev) = old {
            // 복구
            // HashMap::insert로 다시 넣어도 되고, 간단히 remove 후 prev를 넣어도 됨
            // 여기선 단순화를 위해 remove만
        }
        r
    }
}

fn infer(ctx: &mut Ctx, e: &Expr) -> Option<Rc<Ty>> {
    match e {
        Expr::Bool(_) => Some(Rc::new(Ty::Bool)),
        Expr::Var(x) => ctx.lookup(x),
        Expr::Lam(x, body) => {
            let tv = ctx.fresh();
            let body_ty = {
                let old = ctx.env.insert(x.clone(), tv.clone());
                let t = infer(ctx, body)?;
                if let Some(prev) = old {
                    ctx.env.insert(x.clone(), prev);
                } else {
                    ctx.env.remove(x);
                }
                t
            };
            Some(Rc::new(Ty::Func(tv, body_ty)))
        }
        Expr::App(f, a) => {
            let tf = infer(ctx, f)?;
            let ta = infer(ctx, a)?;
            let tr = ctx.fresh();
            let want = Rc::new(Ty::Func(ta.clone(), tr.clone()));
            if unify(&tf, &want) { Some(tr) } else { None }
        }
        Expr::If(c, t, e) => {
            let tc = infer(ctx, c)?;
            if !unify(&tc, &Rc::new(Ty::Bool)) {
                return None;
            }
            let tt = infer(ctx, t)?;
            let te = infer(ctx, e)?;
            if unify(&tt, &te) { Some(tt) } else { None }
        }
    }
}

핵심 흐름은 간단한데요.

람다에서 매개변수에 새 타입변수를 주고 본문을 추론한 뒤 'param → body' 형태로 묶어주고, 적용에서는 'f a'를 'f : a → r'로 강제해 'unify'를 돌려 r을 얻는 식이죠.

이 구조가 바로 ML 계열 언어의 전통적인 타입추론 뼈대와 맞닿아 있어요.

보기 좋은 타입 출력기

추론된 타입을 보기 좋게 찍어주는 건 디버깅에서 꽤 중요하더라고요.

여기서는 변수 포인터 주소를 키로 잡고 'a, b, c...' 같은 이름을 붙이는 간단한 프린터를 추가해볼게요.

use std::collections::BTreeMap;
use std::ptr;

fn pretty(ty: &Rc<Ty>) -> String {
    let mut names: BTreeMap<usize, String> = BTreeMap::new();
    let mut supply = 0u32;
    fn name_for<'a>(
        names: &mut BTreeMap<usize, String>,
        supply: &mut u32,
        p: *mut Var,
    ) -> String {
        let k = p as usize;
        if let Some(n) = names.get(&k) {
            return n.clone();
        }
        let base = ((*supply) as u8 % 26) as char;
        let round = *supply / 26;
        let s = if round == 0 {
            format!("'{}", (b'a' + (base as u8 - b'a')) as char)
        } else {
            format!("'{}{}", (b'a' + (base as u8 - b'a')) as char, round)
        };
        *supply += 1;
        names.insert(k, s.clone());
        s
    }
    fn go(ty: &Rc<Ty>, names: &mut BTreeMap<usize, String>, supply: &mut u32) -> String {
        match ty.as_ref() {
            Ty::Bool => "Bool".to_string(),
            Ty::Func(a, b) => format!("({} -> {})", go(a, names, supply), go(b, names, supply)),
            Ty::Var(v) => match &*v.borrow() {
                Var::Assigned(t) => go(t, names, supply),
                Var::Unassigned(_) => name_for(names, supply, v.as_ptr()),
            },
        }
    }
    go(ty, &mut names, &mut supply)
}

이렇게 해두면 'λx. x'는 'a -> a'처럼 깔끔하게 드러나는데요.

가시성이 좋아지면 유니파이어가 엮는 타입 경로도 한눈에 들어오죠.

Rc와 RefCell을 쓸 때의 안전한 습관

'RefCell'은 런타임에 빌림 규칙을 체크하거든요.

그래서 빌림 기간이 겹치지 않도록 'borrow()'로 얻은 참조는 가능한 한 짧게 쓰고 바로 스코프를 끝내는 게 좋아요.

특히 'if let'이나 'match' 안에서 'borrow_mut()'가 이어질 때는 중간에 임시 변수를 만들어 소유권을 복사해두거나, 블록으로 둘러 빌림을 끊어주면 안전하죠.

'Rc'는 복사가 빠르고 싸서 대입이 잦은 유니파이어에 잘 맞아요.

다만 스레드 공유가 필요하다면 'Arc<Mutex<...>>' 계열이 필요하고, 이 글의 구조는 단일 스레드에서 최대 성능을 뽑는 쪽에 초점이 맞춰져 있거든요.

목표에 따라 선택을 다르게 가져가면 됩니다.

마무리

지금까지 'Rc'와 'RefCell'로 타입과 타입변수를 설계하고, 'occurs-check'와 'unify'를 구현한 다음, 'union-by-rank'로 병목을 해소하는 흐름을 살펴봤는데요.

작은 조합이지만 힘이 꽤 세고, 코드의 형태도 단정하게 떨어지죠.

실전에서는 스키마 폴리모피즘, 일반화와 인스턴스화, 레코드나 유니언 같은 복합 타입, 에러 메시지 개선 같은 과제가 남아있어요.

그래도 오늘 다룬 축이 흔들리지 않으면 그 위에 꽤 많은 걸 올릴 수 있거든요.

마지막으로, 에디션 차이에서 오는 'RefCell' 빌림 세부 동작은 실제 프로젝트 설정과 러스트 버전에 따라 달라질 수 있는데요.

테스트를 자주 돌리고 빌림 범위를 눈에 보이게 짧게 유지한다면 대부분의 함정은 쉽게 피할 수 있죠.

자, 이제 손에 들고 있는 유니파이어를 자신만의 언어에 끼워 넣어보면 어떨까요.

---

부록 전체 코드 스니펫

아래는 글의 핵심 조각을 하나로 엮은 최소 실행 예시인데요.

러스트 에디션과 디버깅 출력을 취향에 맞게 손보며 놀아보면 좋아요.

use std::cell::RefCell;
use std::cmp::Ordering;
use std::collections::{BTreeMap, HashMap};
use std::rc::Rc;

#[derive(Clone)]
enum Ty {
    Bool,
    Func(Rc<Ty>, Rc<Ty>),
    Var(RefCell<Var>),
}

#[derive(Clone)]
enum Var {
    Assigned(Rc<Ty>),
    Unassigned(u32), // union-by-rank
}

impl Ty {
    fn contains(self: &Rc<Ty>, var: &RefCell<Var>) -> bool {
        match self.as_ref() {
            Ty::Bool => false,
            Ty::Func(p, r) => p.contains(var) || r.contains(var),
            Ty::Var(self_var) => match *self_var.borrow() {
                Var::Assigned(ref ty) => ty.contains(var),
                Var::Unassigned(_) => self_var.as_ptr() == var.as_ptr(),
            },
        }
    }
}

fn unify(left: &Rc<Ty>, right: &Rc<Ty>) -> bool {
    match (left.as_ref(), right.as_ref()) {
        (Ty::Bool, Ty::Bool) => true,
        (Ty::Func(lp, lr), Ty::Func(rp, rr)) => unify(lp, rp) && unify(lr, rr),
        (Ty::Var(left_var), Ty::Var(right_var)) => {
            let left_rank = match *left_var.borrow() {
                Var::Assigned(ref l) => return unify(l, right),
                Var::Unassigned(rank) => rank,
            };
            let right_rank = match *right_var.borrow() {
                Var::Assigned(ref r) => return unify(left, r),
                Var::Unassigned(rank) => rank,
            };
            if left_var.as_ptr() != right_var.as_ptr() {
                match left_rank.cmp(&right_rank) {
                    Ordering::Greater => {
                        *right_var.borrow_mut() = Var::Assigned(left.clone());
                    }
                    Ordering::Less => {
                        *left_var.borrow_mut() = Var::Assigned(right.clone());
                    }
                    Ordering::Equal => {
                        *left_var.borrow_mut() = Var::Assigned(right.clone());
                        *right_var.borrow_mut() = Var::Unassigned(right_rank + 1);
                    }
                }
            }
            true
        }
        (Ty::Var(left_var), _) => {
            if let Var::Assigned(ref l) = *left_var.borrow() {
                unify(l, right)
            } else if right.contains(left_var) {
                false
            } else {
                *left_var.borrow_mut() = Var::Assigned(right.clone());
                true
            }
        }
        (_, Ty::Var(right_var)) => {
            if let Var::Assigned(ref r) = *right_var.borrow() {
                unify(left, r)
            } else if left.contains(right_var) {
                false
            } else {
                *right_var.borrow_mut() = Var::Assigned(left.clone());
                true
            }
        }
        _ => false,
    }
}

#[derive(Clone)]
enum Expr {
    Var(String),
    Lam(String, Box<Expr>),
    App(Box<Expr>, Box<Expr>),
    Bool(bool),
    If(Box<Expr>, Box<Expr>, Box<Expr>),
}

struct Ctx {
    env: HashMap<String, Rc<Ty>>,
    supply: u32,
}

impl Ctx {
    fn new() -> Self {
        Self { env: HashMap::new(), supply: 0 }
    }
    fn fresh(&mut self) -> Rc<Ty> {
        Rc::new(Ty::Var(RefCell::new(Var::Unassigned(0))))
    }
    fn lookup(&self, x: &str) -> Option<Rc<Ty>> {
        self.env.get(x).cloned()
    }
}

fn infer(ctx: &mut Ctx, e: &Expr) -> Option<Rc<Ty>> {
    match e {
        Expr::Bool(_) => Some(Rc::new(Ty::Bool)),
        Expr::Var(x) => ctx.lookup(x),
        Expr::Lam(x, body) => {
            let tv = ctx.fresh();
            let old = ctx.env.insert(x.clone(), tv.clone());
            let body_ty = infer(ctx, body)?;
            if let Some(prev) = old {
                ctx.env.insert(x.clone(), prev);
            } else {
                ctx.env.remove(x);
            }
            Some(Rc::new(Ty::Func(tv, body_ty)))
        }
        Expr::App(f, a) => {
            let tf = infer(ctx, f)?;
            let ta = infer(ctx, a)?;
            let tr = ctx.fresh();
            if unify(&tf, &Rc::new(Ty::Func(ta.clone(), tr.clone()))) {
                Some(tr)
            } else {
                None
            }
        }
        Expr::If(c, t, e) => {
            let tc = infer(ctx, c)?;
            if !unify(&tc, &Rc::new(Ty::Bool)) {
                return None;
            }
            let tt = infer(ctx, t)?;
            let te = infer(ctx, e)?;
            if unify(&tt, &te) { Some(tt) } else { None }
        }
    }
}

fn pretty(ty: &Rc<Ty>) -> String {
    use std::collections::BTreeMap;
    let mut names: BTreeMap<usize, String> = BTreeMap::new();
    let mut supply = 0u32;
    fn name_for(
        names: &mut BTreeMap<usize, String>,
        supply: &mut u32,
        p: *mut Var,
    ) -> String {
        let k = p as usize;
        if let Some(n) = names.get(&k) {
            return n.clone();
        }
        let idx = *supply;
        let ch = ((idx % 26) as u8 + b'a') as char;
        let suf = idx / 26;
        let s = if suf == 0 { format!("'{}", ch) } else { format!("'{}{}", ch, suf) };
        *supply += 1;
        names.insert(k, s.clone());
        s
    }
    fn go(ty: &Rc<Ty>, names: &mut BTreeMap<usize, String>, supply: &mut u32) -> String {
        match ty.as_ref() {
            Ty::Bool => "Bool".to_string(),
            Ty::Func(a, b) => format!("({} -> {})", go(a, names, supply), go(b, names, supply)),
            Ty::Var(v) => match &*v.borrow() {
                Var::Assigned(t) => go(t, names, supply),
                Var::Unassigned(_) => name_for(names, supply, v.as_ptr()),
            },
        }
    }
    go(ty, &mut names, &mut supply)
}

fn main() {
    let id = Expr::Lam("x".into(), Box::new(Expr::Var("x".into())));
    let mut ctx = Ctx::new();
    let ty = infer(&mut ctx, &id).unwrap();
    println!("{}", pretty(&ty)); // ('a -> 'a)
}