lysp: rework upvalue handling

userspace/lib/lysp/examples/upvalue.lysp

@@ -0,0 +1,55 @@
+;; open upvalue get
+(setq res-1
+      (let (a 123)
+        ((lambda (b) (+ a b)) 321)
+        ))
+;; shared open upvalue get
+(setq res-2
+      (let (a 123)
+        (let
+          (
+           x ((lambda (b) (+ a b) (+ a b)) 321)
+           y ((lambda (b) (+ a b) (* a b)) 2)
+           )
+          (+ x y)
+          )
+        )
+      )
+;; closed upvalue get
+(setq func-0 (let (a 123) (lambda (b) (+ a b))))
+
+;; open upvalue set
+(setq res-3
+      (let (a 123)
+        ((lambda (b) (setq a b)) 321)
+        a
+        )
+      )
+;; shared open upvalue set
+(setq res-4
+      (let (a 123)
+        (let (
+              funcs (list
+                      (lambda () (setq a 1))
+                      (lambda () (setq a (+ a 1)))
+                      (lambda () (setq a (+ a 2)))
+                      )
+              )
+          (while (not (nil? funcs))
+            ((car funcs))
+            (setq funcs (cdr funcs))
+            )
+
+          a
+          )
+        )
+      )
+;; closed upvalue set
+(setq func-1 (let (a 123) (lambda (b) (setq a (+ a b)) a)))
+
+(assert (= 444 res-1))
+(assert (= 690 res-2))
+(assert (= 444 (func-0 321)))
+(assert (= 321 res-3))
+(assert (= 4 res-4))
+(assert (= 444 (func-1 321)))

userspace/lib/lysp/src/vm/machine.rs

@@ -22,7 +22,9 @@ use crate::{
         macros::MacroExpand,
         prelude,
         stack::Stack,
-        value::{BytecodeFunction, ClosureValue, IdentifierValue, NumberValue, UpvalueValue},
+        value::{
+            BytecodeFunction, ClosureValue, IdentifierValue, NumberValue, Upvalue, UpvalueRef,
+        },
     },
 };
 
@@ -31,6 +33,7 @@ pub struct CallFrame {
     pub closure: ClosureValue,
     pub ip: usize,
     pub base_pointer: usize,
+    open_upvalue_trackers: Vec<OpenUpvalueTracker>,
 }
 
 pub struct Machine {
@@ -38,8 +41,7 @@ pub struct Machine {
     call_stack: Stack<CallFrame>,
     tmp: Option<Value>,
 
-    upvalue_arena: Vec<UpvalueValue>,
+    // upvalue_arena: Vec<UpvalueValue>,
-
     pub trace_instructions: bool,
     pub trace_returns: bool,
     pub trace_stack: bool,
@@ -47,14 +49,28 @@ pub struct Machine {
     pub trace_macros: bool,
 }
 
+#[derive(Debug)]
+struct OpenUpvalueTracker {
+    sp: usize,
+    upvalue_ref: UpvalueRef,
+}
+
+impl From<usize> for OpenUpvalueTracker {
+    fn from(value: usize) -> Self {
+        Self {
+            sp: value,
+            upvalue_ref: UpvalueRef::open(value),
+        }
+    }
+}
+
 impl Default for Machine {
     fn default() -> Self {
         Self {
             data_stack: Stack::new(1024),
             call_stack: Stack::new(64),
             tmp: None,
-            upvalue_arena: vec![],
+            // upvalue_arena: vec![],
-
             trace_calls: false,
             trace_stack: false,
             trace_returns: false,
@@ -118,30 +134,30 @@ impl Machine {
             .ok_or(MachineError::UndefinedLocalReference)
     }
 
-    fn upvalue_slot(&mut self, id: LocalId) -> Result<&mut Value, MachineError> {
+    // fn upvalue_slot(&mut self, id: LocalId) -> Result<&mut Value, MachineError> {
-        let frame = self
+    //     let frame = self
-            .call_stack
+    //         .call_stack
-            .head()
+    //         .head()
-            .ok_or(MachineError::CallStackUnderflow)?;
+    //         .ok_or(MachineError::CallStackUnderflow)?;
-        let upvalue_arena_index = frame
+    //     let upvalue_arena_index = frame
-            .closure
+    //         .closure
-            .upvalues
+    //         .upvalues
-            .get(usize::from(id))
+    //         .get(usize::from(id))
-            .copied()
+    //         .copied()
-            .ok_or(MachineError::UndefinedUpvalueReference)?;
+    //         .ok_or(MachineError::UndefinedUpvalueReference)?;
-        let upvalue_value = self
+    //     let upvalue_value = self
-            .upvalue_arena
+    //         .upvalue_arena
-            .get_mut(upvalue_arena_index)
+    //         .get_mut(upvalue_arena_index)
-            .ok_or(MachineError::UndefinedUpvalueReference)?;
+    //         .ok_or(MachineError::UndefinedUpvalueReference)?;
 
-        match upvalue_value {
+    //     match upvalue_value {
-            UpvalueValue::Open(sp) => self
+    //         UpvalueValue::Open(sp) => self
-                .data_stack
+    //             .data_stack
-                .get_mut(*sp)
+    //             .get_mut(*sp)
-                .ok_or(MachineError::UndefinedUpvalueReference),
+    //             .ok_or(MachineError::UndefinedUpvalueReference),
-            UpvalueValue::Closed(boxed) => Ok(boxed.as_mut()),
+    //         UpvalueValue::Closed(boxed) => Ok(boxed.as_mut()),
-        }
+    //     }
-    }
+    // }
 
     fn execute_get_local(&mut self, id: LocalId) -> Result<(), MachineError> {
         let value = self.local_slot(id)?.clone();
@@ -155,14 +171,22 @@ impl Machine {
     }
 
     fn execute_get_upvalue(&mut self, id: LocalId) -> Result<(), MachineError> {
-        let value = self.upvalue_slot(id)?.clone();
+        let frame = self
+            .call_stack
+            .head()
+            .ok_or(MachineError::CallStackUnderflow)?;
+        let value = frame.closure.upvalues[usize::from(id)].read(&self.data_stack)?;
         self.push(value)?;
         Ok(())
     }
 
     fn execute_set_upvalue(&mut self, id: LocalId) -> Result<(), MachineError> {
         let value = self.pop()?;
-        *self.upvalue_slot(id)? = value;
+        let frame = self
+            .call_stack
+            .head()
+            .ok_or(MachineError::CallStackUnderflow)?;
+        frame.closure.upvalues[usize::from(id)].write(&mut self.data_stack, value)?;
         Ok(())
     }
 
@@ -329,6 +353,7 @@ impl Machine {
             closure,
             base_pointer,
             ip: 0,
+            open_upvalue_trackers: vec![],
         };
         self.call_stack
             .push(frame)
@@ -337,33 +362,6 @@ impl Machine {
         Ok(())
     }
 
-    fn capture_stack_upvalue(&mut self, sp: usize) -> Result<usize, MachineError> {
-        for arena_index in (0..self.upvalue_arena.len()).rev() {
-            let upvalue = &self.upvalue_arena[arena_index];
-            if let UpvalueValue::Open(target_sp) = upvalue
-                && *target_sp == sp
-            {
-                return Ok(arena_index);
-            }
-        }
-
-        let arena_index = self.upvalue_arena.len();
-        self.upvalue_arena.push(UpvalueValue::Open(sp));
-        Ok(arena_index)
-    }
-
-    fn close_upvalues(&mut self, sp: usize) {
-        // TODO this is inefficient
-        for uv in self.upvalue_arena.iter_mut() {
-            if let UpvalueValue::Open(target_sp) = uv
-                && *target_sp >= sp
-            {
-                let value = self.data_stack[*target_sp].clone();
-                *uv = UpvalueValue::Closed(Box::new(value));
-            }
-        }
-    }
-
     fn execute_make_closure(&mut self) -> Result<(), MachineError> {
         let value = self.pop()?;
         let Value::Function(function) = value else {
@@ -379,16 +377,45 @@ impl Machine {
             function,
             upvalues: vec![],
         };
+        let frame = self
+            .call_stack
+            .head_mut()
+            .ok_or(MachineError::CallStackUnderflow)?;
+        // eprintln!("BEGIN MAKE CLOSURE");
         for upvalue_def in closure.function.upvalues.iter() {
-            if upvalue_def.is_local {
+            let upvalue_ref = if upvalue_def.is_local {
-                let frame = self.call_stack.head().unwrap();
                 let sp = frame.base_pointer + 1 + usize::from(upvalue_def.index);
-                let arena_index = self.capture_stack_upvalue(sp)?;
+                let tracker_index = frame
-                closure.upvalues.push(arena_index);
+                    .open_upvalue_trackers
+                    .binary_search_by_key(&sp, |t| t.sp)
+                    // .inspect(|i| {
+                    //     eprintln!("UPVALUE: found open tracker i={i}, sp={sp}");
+                    // })
+                    .unwrap_or_else(|i| {
+                        // eprintln!("UPVALUE: create new open tracker i={i}, sp={sp}");
+                        frame
+                            .open_upvalue_trackers
+                            .insert(i, OpenUpvalueTracker::from(sp));
+                        i
+                    });
+                frame.open_upvalue_trackers[tracker_index]
+                    .upvalue_ref
+                    .clone()
             } else {
-                todo!();
+                todo!()
-            }
+            };
+
+            closure.upvalues.push(upvalue_ref);
         }
+
+        // for (i, upvalue) in closure.upvalues.iter().enumerate() {
+        //     eprintln!(
+        //         "UPVALUE #{i}: {upvalue:?} => {}",
+        //         upvalue.read(&self.data_stack).unwrap()
+        //     );
+        // }
+        // eprintln!("END MAKE CLOSURE");
+
         self.push(Value::Closure(closure))?;
         Ok(())
     }
@@ -589,7 +616,25 @@ impl Machine {
                 self.push(tmp)?;
             }
             Instruction::CloseUpvalue => {
-                self.close_upvalues(self.data_stack.pointer() - 1);
+                let frame = self
+                    .call_stack
+                    .head_mut()
+                    .ok_or(MachineError::CallStackUnderflow)?;
+                let upvalue_sp = self.data_stack.pointer() - 1;
+                let start = frame
+                    .open_upvalue_trackers
+                    .binary_search_by_key(&upvalue_sp, |t| t.sp)
+                    .unwrap_or_else(|i| i);
+                for tracker in frame.open_upvalue_trackers.drain(start..) {
+                    let value = self
+                        .data_stack
+                        .get(tracker.sp)
+                        .cloned()
+                        .ok_or(MachineError::UndefinedUpvalueReference)?;
+                    let old = tracker.upvalue_ref.close(value);
+                    // eprintln!("CLOSE UPVALUE sp={} -> {}", tracker.sp, value);
+                    assert_eq!(old, Upvalue::Open(tracker.sp));
+                }
                 self.pop()?;
             }
             Instruction::Branch => {

userspace/lib/lysp/src/vm/value/closure.rs

@@ -1,17 +1,62 @@
-use std::{fmt, rc::Rc};
+use std::{cell::RefCell, fmt, rc::Rc};
 
-use crate::vm::{Value, value::BytecodeFunction};
+use crate::{
+    error::MachineError,
+    vm::{Value, stack::Stack, value::BytecodeFunction},
+};
 
-#[derive(Debug, Clone, PartialEq)]
+// #[derive(Debug, Clone, PartialEq)]
-pub enum UpvalueValue {
+// pub enum UpvalueValue {
+//     Open(usize),
+//     Closed(Box<Value>),
+// }
+
+#[derive(Debug, PartialEq)]
+pub enum Upvalue {
     Open(usize),
-    Closed(Box<Value>),
+    Closed(Value),
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct UpvalueRef(Rc<RefCell<Upvalue>>);
+
 #[derive(Clone, PartialEq)]
 pub struct ClosureValue {
     pub function: Rc<BytecodeFunction>,
-    pub upvalues: Vec<usize>,
+    pub upvalues: Vec<UpvalueRef>,
+}
+
+impl UpvalueRef {
+    pub fn open(sp: usize) -> Self {
+        Self(Rc::new(RefCell::new(Upvalue::Open(sp))))
+    }
+
+    pub fn read(&self, stack: &Stack<Value>) -> Result<Value, MachineError> {
+        match &*self.0.borrow() {
+            &Upvalue::Open(sp) => stack
+                .get(sp)
+                .cloned()
+                .ok_or(MachineError::UndefinedUpvalueReference),
+            Upvalue::Closed(value) => Ok(value.clone()),
+        }
+    }
+
+    pub fn write(&self, stack: &mut Stack<Value>, value: Value) -> Result<(), MachineError> {
+        match &mut *self.0.borrow_mut() {
+            &mut Upvalue::Open(sp) => {
+                let slot = stack
+                    .get_mut(sp)
+                    .ok_or(MachineError::UndefinedUpvalueReference)?;
+                *slot = value;
+            }
+            Upvalue::Closed(slot) => *slot = value,
+        }
+        Ok(())
+    }
+
+    pub fn close(&self, value: Value) -> Upvalue {
+        self.0.replace(Upvalue::Closed(value))
+    }
 }
 
 impl ClosureValue {

userspace/lib/lysp/src/vm/value/mod.rs

@@ -20,7 +20,7 @@ mod vector;
 pub mod convert;
 
 pub use boolean::BooleanValue;
-pub use closure::{ClosureValue, UpvalueValue};
+pub use closure::{ClosureValue, Upvalue, UpvalueRef};
 pub use cons::ConsCell;
 pub use function::BytecodeFunction;
 pub use hashtable::{HashTable, HashTableData};