クロージャを実装する——upvalueをOpen/Closedの2状態で表現しスタックとヒープを橋渡しする

fn outer(x: int): int {
    fn inner(): int { return x; }
    return inner();
}

struct FunctionScope {
    chunk: Chunk,
    locals: Vec<Local>,
    scope_depth: usize,
    loops: Vec<LoopContext>,
    pops_return_value: bool,
    upvalues: Vec<UpvalueDesc>,
}

pub struct Compiler {
    frames: Vec<FunctionScope>,
}

Expr::Ident(name) => {
    let frame_index = self.frames.len() - 1;
    if let Some(slot) = self.resolve_local_in(frame_index, name) {
        self.emit_get_local(slot as u8, line);
        return Ok(());
    }
    if let Some(index) = self.resolve_upvalue(frame_index, name) {
        self.emit_get_upvalue(index, line);
        return Ok(());
    }
    Err(self.error(line, column, format!("undefined variable '{name}'")))
}

fn resolve_upvalue(&mut self, frame_index: usize, name: &str) -> Option<u8> {
    if frame_index == 0 {
        return None; // トップレベルには外側が無い
    }
    let outer = frame_index - 1;

    if let Some(slot) = self.resolve_local_in(outer, name) {
        // 一つ外側のフレームのローカル変数が見つかった。
        // 自分のスコープより長生きする必要があるとマークする。
        self.frames[outer].locals[slot].is_captured = true;
        return Some(self.add_upvalue(frame_index, slot as u8, true));
    }
    if let Some(index) = self.resolve_upvalue(outer, name) {
        // 一つ外側にも無かった。さらに外側にあるはずなので、
        // 一つ外側の関数自身のupvalueとして間接的に参照する。
        return Some(self.add_upvalue(frame_index, index, false));
    }
    None
}

fn a(x: int): int {
    fn b(): int {
        fn c(): int { return x; }
        return c();
    }
    return b();
}

fn emit_closure(&mut self, function: Value, upvalues: &[UpvalueDesc], line: usize) {
    let index = self.frame().chunk.add_constant(function);
    self.frame().chunk.write_op(OpCode::Closure, line);
    self.frame().chunk.write_u8(index, line);
    for upvalue in upvalues {
        self.frame().chunk.write_u8(upvalue.is_local as u8, line);
        self.frame().chunk.write_u8(upvalue.index, line);
    }
}

OP_CLOSURE  <inner関数の定数インデックス>
  is_local=1  index=1     ; outerのスロット1(パラメータx)を直接捕まえる

pub struct FunctionObj {
    pub name: String,
    pub arity: u8,
    pub chunk: Chunk,
    pub upvalue_count: u8,
}

OpCode::Closure => {
    let Value::Function(function) = self.read_constant() else {
        unreachable!("compiler always emits a function constant for OP_CLOSURE");
    };
    let mut upvalues = Vec::with_capacity(function.upvalue_count as usize);
    for _ in 0..function.upvalue_count {
        let is_local = self.read_u8() != 0;
        let index = self.read_u8() as usize;
        let upvalue = if is_local {
            // 今実行中のフレーム(=直接の外側の関数)のローカルスロットを直接捕まえる
            self.capture_upvalue(self.locals_base() + index)
        } else {
            // 今実行中の関数自身がすでに持っているupvalueを共有する
            Rc::clone(&self.current_closure().upvalues[index])
        };
        upvalues.push(upvalue);
    }
    self.push(Value::Closure(Rc::new(ClosureObj { function, upvalues })));
}

pub enum Upvalue {
    Open(usize),   // スタックのこの位置を指している
    Closed(Value), // 値をヒープ側へコピー済み
}

OpCode::GetUpvalue => {
    let index = self.read_u8() as usize;
    let upvalue = Rc::clone(&self.current_closure().upvalues[index]);
    let value = match &*upvalue.borrow() {
        Upvalue::Open(stack_index) => self.stack[*stack_index].clone(),
        Upvalue::Closed(value) => value.clone(),
    };
    self.push(value);
}

fn capture_upvalue(&mut self, stack_index: usize) -> Rc<RefCell<Upvalue>> {
    if let Some(existing) = self.open_upvalues.iter().find(|upvalue| open_slot(upvalue) == Some(stack_index)) {
        return Rc::clone(existing);
    }
    let upvalue = Rc::new(RefCell::new(Upvalue::Open(stack_index)));
    self.open_upvalues.push(Rc::clone(&upvalue));
    upvalue
}

fn end_scope(&mut self, line: usize) {
    self.frame().scope_depth -= 1;
    let scope_depth = self.frame().scope_depth;
    while let Some(local) = self.frame().locals.last() {
        if local.depth <= scope_depth {
            break;
        }
        let is_captured = local.is_captured;
        self.frame().locals.pop();
        if is_captured {
            self.frame().chunk.write_op(OpCode::CloseUpvalue, line);
        } else {
            self.frame().chunk.write_op(OpCode::Pop, line);
        }
    }
}

OpCode::CloseUpvalue => {
    let stack_index = self.stack.len() - 1;
    self.close_upvalue_at(stack_index);
    self.pop();
}

OpCode::Return => {
    let Some(frame) = self.frames.pop() else {
        return Ok(());
    };
    let result = self.pop();
    // このフレーム内で作られたクロージャがまだopenなupvalueを
    // 持っているかもしれない。スタック領域を捨てる前にヒープへ退避する。
    self.close_upvalues_from(frame.stack_base);
    self.stack.truncate(frame.stack_base);
    self.push(result);
}

fn make(x: int): int {
    let y = x + 1;
    fn get(): int { return x + y; }
    return get();
}

fn outer(x: int): int {
    fn a(): int {
        fn b(): int {
            fn c(): int { return x; }
            return c();
        }
        return b();
    }
    return a();
}

let m = make(10);
let o = outer(7);

$ cargo run -- examples/closures.na
stack after execution: [Int(21), Int(7)]