update repeat w/ new jumps

src/compiler.rs

@@ -410,7 +410,7 @@ impl<'a> Compiler<'a> {
         self.chunk.bytecode.push(low);
     }
 
-    fn jump_stub(&mut self, op: Op) -> usize {
+    fn stub_jump(&mut self, op: Op) -> usize {
         let out = self.chunk.bytecode.len();
         self.emit_op(op);
         self.emit_byte(0xff);
@@ -554,10 +554,16 @@ impl<'a> Compiler<'a> {
     }
 
     fn pop_n(&mut self, n: usize) {
+        match n {
+            0 => (),
+            1 => self.pop(),
+            n => {
                 self.emit_op(Op::PopN);
                 self.emit_byte(n);
                 self.stack_depth -= n;
             }
+        }
+    }
 
     fn enter_loop(&mut self) {
         self.loop_info
@@ -657,10 +663,10 @@ impl<'a> Compiler<'a> {
             }
             If(cond, then, r#else) => {
                 self.visit(cond);
-                let jif_idx = self.jump_stub(Op::JumpIfFalse);
+                let jif_idx = self.stub_jump(Op::JumpIfFalse);
                 self.stack_depth -= 1;
                 self.visit(then);
-                let jump_idx = self.jump_stub(Op::Jump);
+                let jump_idx = self.stub_jump(Op::Jump);
                 self.visit(r#else);
                 self.stack_depth -= 1;
                 let end_idx = self.len();
@@ -722,7 +728,7 @@ impl<'a> Compiler<'a> {
                 self.emit_op(Op::MatchTuple);
                 self.emit_byte(members.len());
                 // skip everything if tuple lengths don't match
-                let before_load_tup_idx = self.jump_stub(Op::JumpIfNoMatch);
+                let before_load_tup_idx = self.stub_jump(Op::JumpIfNoMatch);
 
                 // set up the per-member conditional logic
                 let mut jump_idxes = vec![];
@@ -743,12 +749,11 @@ impl<'a> Compiler<'a> {
                     // visit the pattern member
                     self.visit(member);
                     // and jump if there's no match
-                    jump_idxes.push(self.jump_stub(Op::JumpIfNoMatch));
+                    jump_idxes.push(self.stub_jump(Op::JumpIfNoMatch));
                 }
 
                 // if we get here--not having jumped on no match--we're matched; jump the "no match" code
-                self.match_depth = match_depth + members.len();
-                let jump_idx = self.jump_stub(Op::Jump);
+                let jump_idx = self.stub_jump(Op::Jump);
 
                 // patch up the previous no match jumps to jump to clean-up code
                 for idx in jump_idxes {
@@ -766,45 +771,52 @@ impl<'a> Compiler<'a> {
 
                 // patch up the yes-matches unconditional jump
                 self.patch_jump(jump_idx, self.len() - jump_idx - 3);
+
+                // finally, for any further matches (e.g. nested lists/tuples)
+                // add increase the match depth, since we've added a bunch
+                // of bindings to the stack
+                self.match_depth = match_depth + members.len();
             }
             ListPattern(members) => {
                 self.emit_op(Op::MatchList);
                 self.emit_byte(members.len());
-                self.emit_op(Op::JumpIfNoMatch);
-                let before_load_list_idx = self.len();
-                self.emit_byte(0xff);
+                let before_load_tup_idx = self.stub_jump(Op::JumpIfNoMatch);
+
                 let mut jump_idxes = vec![];
-                self.match_depth += members.len();
+                let match_depth = self.match_depth;
+                self.match_depth = members.len();
+
                 self.emit_op(Op::LoadList);
                 self.stack_depth += members.len();
+
                 for member in members {
                     self.match_depth -= 1;
                     self.emit_op(Op::MatchDepth);
                     self.emit_byte(self.match_depth);
                     self.visit(member);
-                    self.emit_op(Op::JumpIfNoMatch);
-                    jump_idxes.push(self.len());
-                    self.emit_byte(0xff);
+                    jump_idxes.push(self.stub_jump(Op::JumpIfNoMatch));
                 }
-                self.emit_op(Op::Jump);
-                let jump_idx = self.len();
-                self.emit_byte(0xff);
+
+                let jump_idx = self.stub_jump(Op::Jump);
+
                 for idx in jump_idxes {
-                    self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
+                    self.patch_jump(idx, self.len() - idx - 2)
                 }
-                for _ in 0..members.len() {
-                    self.emit_op(Op::Pop);
-                }
-                self.chunk.bytecode[before_load_list_idx] =
-                    (self.len() - before_load_list_idx) as u8 - 1;
-                self.chunk.bytecode[jump_idx] = (self.len() - jump_idx) as u8 - 1;
+
+                self.emit_op(Op::PopN);
+                self.emit_byte(members.len());
+
+                self.patch_jump(before_load_tup_idx, self.len() - before_load_tup_idx - 3);
+
+                self.patch_jump(jump_idx, self.len() - jump_idx - 3);
+
+                self.match_depth = match_depth + members.len();
             }
             DictPattern(pairs) => {
                 self.emit_op(Op::MatchDict);
                 self.emit_byte(pairs.len());
-                self.emit_op(Op::JumpIfNoMatch);
-                let before_load_dict_idx = self.len();
-                self.emit_byte(0xff);
+                let before_load_dict_idx = self.stub_jump(Op::JumpIfNoMatch);
+
                 let mut jump_idxes = vec![];
                 let dict_stack_pos = self.stack_depth - 1;
                 for pair in pairs {
@@ -815,22 +827,19 @@ impl<'a> Compiler<'a> {
                     self.emit_op(Op::LoadDictValue);
                     self.emit_byte(dict_stack_pos);
                     self.visit(pattern);
-                    self.emit_op(Op::JumpIfNoMatch);
-                    jump_idxes.push(self.len());
-                    self.emit_byte(0xff);
+                    jump_idxes.push(self.stub_jump(Op::JumpIfNoMatch));
                 }
-                self.emit_op(Op::Jump);
-                let jump_idx = self.len();
-                self.emit_byte(0xff);
+                let jump_idx = self.stub_jump(Op::Jump);
+
                 for idx in jump_idxes {
-                    self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
+                    self.patch_jump(idx, self.len() - idx - 2);
                 }
-                for _ in 0..pairs.len() {
-                    self.emit_op(Op::Pop);
-                }
-                self.chunk.bytecode[before_load_dict_idx] =
-                    (self.len() - before_load_dict_idx) as u8 - 1;
-                self.chunk.bytecode[jump_idx] = (self.len() - jump_idx) as u8 - 1;
+
+                self.emit_op(Op::PopN);
+                self.emit_byte(pairs.len());
+
+                self.patch_jump(before_load_dict_idx, self.len() - before_load_dict_idx - 3);
+                self.patch_jump(jump_idx, self.len() - jump_idx - 3);
             }
             Splattern(..) => {
                 todo!()
@@ -864,9 +873,7 @@ impl<'a> Compiler<'a> {
                 self.emit_op(Op::MatchString);
                 self.emit_byte(pattern_idx);
 
-                self.emit_op(Op::JumpIfNoMatch);
-                let jnm_idx = self.len();
-                self.emit_byte(0xff);
+                let jnm_idx = self.stub_jump(Op::JumpIfNoMatch);
 
                 self.emit_op(Op::PushStringMatches);
                 self.emit_byte(pattern_idx);
@@ -882,7 +889,7 @@ impl<'a> Compiler<'a> {
                     self.stack_depth += 1;
                 }
 
-                self.chunk.bytecode[jnm_idx] = (self.len() - jnm_idx - 1) as u8;
+                self.patch_jump(jnm_idx, self.len() - jnm_idx - 3);
             }
             PairPattern(_, _) => unreachable!(),
             Tuple(members) => {
@@ -950,12 +957,10 @@ impl<'a> Compiler<'a> {
                             for arg in args {
                                 self.visit(arg);
                                 self.emit_op(Op::Stash);
-                                self.emit_op(Op::JumpIfTrue);
-                                jump_idxes.push(self.len());
-                                self.emit_byte(0xff);
+                                jump_idxes.push(self.stub_jump(Op::JumpIfTrue));
                             }
                             for idx in jump_idxes {
-                                self.chunk.bytecode[idx] = (self.len() - idx + 2) as u8;
+                                self.patch_jump(idx, self.len() - idx);
                             }
                             self.emit_op(Op::Load);
                         } else {
@@ -970,12 +975,10 @@ impl<'a> Compiler<'a> {
                             for arg in args {
                                 self.visit(arg);
                                 self.emit_op(Op::Stash);
-                                self.emit_op(Op::JumpIfFalse);
-                                jump_idxes.push(self.len());
-                                self.emit_byte(0xff);
+                                jump_idxes.push(self.stub_jump(Op::JumpIfFalse));
                             }
                             for idx in jump_idxes {
-                                self.chunk.bytecode[idx] = (self.len() - idx + 2) as u8;
+                                self.patch_jump(idx, self.len() - idx);
                             }
                             self.emit_op(Op::Load);
                         } else {
@@ -1045,20 +1048,16 @@ impl<'a> Compiler<'a> {
                 let mut clauses = clauses.iter();
                 while let Some((WhenClause(cond, body), _)) = clauses.next() {
                     self.visit(cond.as_ref());
-                    self.emit_op(Op::JumpIfFalse);
-                    let jif_jump_idx = self.len();
-                    self.emit_byte(0xff);
+                    let jif_jump_idx = self.stub_jump(Op::JumpIfFalse);
                     self.stack_depth -= 1;
                     self.visit(body);
                     self.stack_depth -= 1;
-                    self.emit_op(Op::Jump);
-                    jump_idxes.push(self.len());
-                    self.emit_byte(0xff);
-                    self.chunk.bytecode[jif_jump_idx] = (self.len() - jif_jump_idx) as u8 - 1;
+                    jump_idxes.push(self.stub_jump(Op::Jump));
+                    self.patch_jump(jif_jump_idx, self.len() - jif_jump_idx - 3);
                 }
                 self.emit_op(Op::PanicNoWhen);
                 for idx in jump_idxes {
-                    self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
+                    self.patch_jump(idx, self.len() - idx - 3);
                 }
                 self.stack_depth += 1;
             }
@@ -1073,16 +1072,12 @@ impl<'a> Compiler<'a> {
                     self.scope_depth += 1;
                     self.match_depth = 0;
                     self.visit(pattern);
-                    self.emit_op(Op::JumpIfNoMatch);
-                    no_match_jumps.push(self.len());
-                    self.emit_byte(0xff);
+                    no_match_jumps.push(self.stub_jump(Op::JumpIfNoMatch));
                     if guard.is_some() {
                         let guard_expr: &'static Spanned<Ast> =
                             Box::leak(Box::new(guard.clone().unwrap()));
                         self.visit(guard_expr);
-                        self.emit_op(Op::JumpIfFalse);
-                        no_match_jumps.push(self.len());
-                        self.emit_byte(0xff);
+                        no_match_jumps.push(self.stub_jump(Op::JumpIfFalse));
                         self.stack_depth -= 1;
                     }
                     self.visit(body);
@@ -1095,31 +1090,29 @@ impl<'a> Compiler<'a> {
                             break;
                         }
                     }
-                    while self.stack_depth > stack_depth {
-                        self.pop();
-                    }
-                    self.emit_op(Op::Jump);
-                    jump_idxes.push(self.len());
-                    self.emit_byte(0xff);
+                    self.pop_n(self.stack_depth - stack_depth);
+                    jump_idxes.push(self.stub_jump(Op::Jump));
                     for idx in no_match_jumps {
-                        self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
+                        self.patch_jump(idx, self.len() - idx - 3);
                     }
                 }
                 self.emit_op(Op::PanicNoMatch);
+
                 for idx in jump_idxes {
-                    self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
-                }
-                while self.stack_depth > stack_depth {
-                    self.pop();
+                    self.patch_jump(idx, self.len() - idx - 3);
+                    // self.chunk.bytecode[idx] = (self.len() - idx) as u8 - 1;
                 }
+                self.pop_n(self.stack_depth - stack_depth);
+                // while self.stack_depth > stack_depth {
+                //     self.pop();
+                // }
                 self.emit_op(Op::Load);
                 self.stack_depth += 1;
             }
             MatchClause(..) => unreachable!(),
             Fn(name, body, doc) => {
                 let name = if name.is_empty() { "anonymous" } else { name };
-                // first, declare the function
-                // TODO: or, check if the function has already been declared!
+
                 let FnBody(fn_body) = &body.as_ref().0 else {
                     unreachable!()
                 };
@@ -1161,35 +1154,26 @@ impl<'a> Compiler<'a> {
                         compiler.emit_op(Op::MatchDepth);
                         compiler.emit_byte(compiler.match_depth);
                         compiler.visit(member);
-                        compiler.emit_op(Op::JumpIfNoMatch);
-                        tup_jump_idxes.push(compiler.len());
-                        compiler.emit_byte(0xff);
+                        tup_jump_idxes.push(compiler.stub_jump(Op::JumpIfNoMatch));
                     }
                     if pattern.is_empty() {
                         compiler.emit_op(Op::Match);
                     }
-                    compiler.emit_op(Op::Jump);
-                    let jump_idx = compiler.len();
-                    compiler.emit_byte(0xff);
+                    let jump_idx = compiler.stub_jump(Op::Jump);
                     for idx in tup_jump_idxes {
-                        compiler.chunk.bytecode[idx] = compiler.len() as u8 - idx as u8 - 2;
+                        compiler.patch_jump(idx, compiler.len() - idx - 3);
                     }
                     for _ in 0..arity {
                         compiler.emit_op(Op::Pop);
                     }
-                    compiler.chunk.bytecode[jump_idx] = compiler.len() as u8 - jump_idx as u8 - 1;
+                    compiler.patch_jump(jump_idx, compiler.len() - jump_idx - 3);
                     let mut no_match_jumps = vec![];
-                    compiler.emit_op(Op::JumpIfNoMatch);
-                    // let jnm_idx = compiler.len();
-                    no_match_jumps.push(compiler.len());
-                    compiler.emit_byte(0xff);
+                    no_match_jumps.push(compiler.stub_jump(Op::JumpIfNoMatch));
                     if guard.is_some() {
                         let guard_expr: &'static Spanned<Ast> =
                             Box::leak(Box::new(guard.clone().unwrap()));
                         compiler.visit(guard_expr);
-                        compiler.emit_op(Op::JumpIfFalse);
-                        no_match_jumps.push(self.len());
-                        compiler.emit_byte(0xff);
+                        no_match_jumps.push(compiler.stub_jump(Op::JumpIfFalse));
                         compiler.stack_depth -= 1;
                     }
                     compiler.visit(clause_body);
@@ -1202,15 +1186,12 @@ impl<'a> Compiler<'a> {
                             break;
                         }
                     }
-                    while compiler.stack_depth > 0 {
-                        compiler.pop();
-                    }
+                    compiler.pop_n(compiler.stack_depth);
                     compiler.stack_depth = 0;
                     compiler.emit_op(Op::Return);
                     for idx in no_match_jumps {
-                        compiler.chunk.bytecode[idx] = compiler.len() as u8 - idx as u8 - 1;
+                        compiler.patch_jump(idx, compiler.len() - idx - 3);
                     }
-                    // compiler.chunk.bytecode[jnm_idx] = compiler.len() as u8 - jnm_idx as u8 - 1;
                     compiler.scope_depth -= 1;
 
                     std::mem::swap(&mut compiler.upvalues, &mut upvalues);
@@ -1251,42 +1232,27 @@ impl<'a> Compiler<'a> {
                 self.bind(name);
             }
             FnBody(_) => unreachable!(),
-            // TODO: add a check to make sure times >= 0
-            // TODO: fix this so that the algorithm is actually correct
-            // * truncate
-            // * check that it's 0 or more, panic if not
-            // * jump if 0 to end of loop
-            // * body
-            // * pop
-            // * decrement counter (which is now at the top of the stack)
-            // * jump back to jump if 0 instruction
             Repeat(times, body) => {
                 self.visit(times);
                 self.emit_op(Op::Truncate);
                 // skip the decrement the first time
                 self.emit_op(Op::Jump);
+                self.emit_byte(0);
                 self.emit_byte(1);
                 // begin repeat
                 self.emit_op(Op::Decrement);
                 let repeat_begin = self.len();
                 self.emit_op(Op::Duplicate);
-                // self.stack_depth += 1;
-                self.emit_op(Op::JumpIfZero);
-                self.emit_byte(0xff);
+                let jiz_idx = self.stub_jump(Op::JumpIfZero);
                 // compile the body
                 self.visit(body);
                 // pop whatever value the body returns
                 self.pop();
-                // self.emit_op(Op::Pop);
-                self.emit_op(Op::JumpBack);
+                let jump_back = self.stub_jump(Op::JumpBack);
                 // set jump points
-                let repeat_end = self.len();
-                self.emit_byte(repeat_end - repeat_begin);
-                self.chunk.bytecode[repeat_begin + 2] = (repeat_end - repeat_begin - 2) as u8;
-                // pop the counter
-                // self.pop();
-                // self.emit_op(Op::Pop);
-                // and emit nil
+                self.patch_jump(jump_back, self.len() - repeat_begin - 2);
+                self.patch_jump(jiz_idx, self.len() - repeat_begin - 4);
+                self.pop();
                 self.emit_constant(Value::Nil);
             }
             Loop(value, clauses) => {

src/main.rs

@@ -73,9 +73,11 @@ pub fn run(src: &'static str) {
 }
 
 pub fn main() {
-    // env::set_var("RUST_BACKTRACE", "1");
+    env::set_var("RUST_BACKTRACE", "1");
     let src = "
-let ((1)) = ((1))
+repeat 4 {
+    :foo
+}
    ";
     run(src);
 }

src/value.rs

@@ -45,6 +45,13 @@ impl LFn {
         }
     }
 
+    pub fn accepts(&self, arity: u8) -> bool {
+        match self {
+            LFn::Defined { chunks, .. } => chunks.iter().any(|(a, _)| *a == arity),
+            LFn::Declared { .. } => unreachable!(),
+        }
+    }
+
     pub fn name(&self) -> &'static str {
         match self {
             LFn::Declared { name } | LFn::Defined { name, .. } => name,
@@ -52,7 +59,6 @@ impl LFn {
     }
 
     pub fn chunk(&self, arity: u8) -> &Chunk {
-        // println!("Getting chunk of {arity} from {:?}", self);
         match self {
             LFn::Declared { .. } => unreachable!(),
             LFn::Defined { chunks, .. } => &chunks.iter().find(|(a, _)| *a == arity).unwrap().1,
@@ -268,4 +274,11 @@ impl Value {
             Partial(..) => "fn",
         }
     }
+
+    pub fn as_fn(&self) -> &LFn {
+        match self {
+            Value::Fn(inner) => inner.as_ref(),
+            _ => unreachable!(),
+        }
+    }
 }

src/vm.rs

@@ -790,6 +790,12 @@ impl Vm {
 
                     match val {
                         Value::Fn(_) => {
+                            if !val.as_fn().accepts(arity) {
+                                return self.panic_with(format!(
+                                    "wrong number of arguments to {} passing {arity} args",
+                                    val.show()
+                                ));
+                            }
                             let mut frame = CallFrame {
                                 function: val,
                                 arity,