Implement loop/return forms

examples/loops.lysp

@@ -0,0 +1,32 @@
+;; vi:ft=lisp:sw=2:ts=2
+
+;; loop, broken by (return)
+(setq looped-loop (let (i 0)
+  (loop
+    (if (< i 10)
+      (setq i (+ i 1))
+      (return)
+      )
+    )
+  i
+  ))
+
+;; while, broken prematurely
+(setq looped-while (let (i 0)
+  (while (< i 20)
+    (setq i (+ i 1))
+    (if (= i 10) (return))
+    )
+  i
+  ))
+
+;; while, broken by condition
+(setq looped-while-full (let (i 0)
+  (while (< i 10)
+    (setq i (+ i 1))
+    )
+  i
+  ))
+
+;; All loops execute the same count of times
+(assert (= looped-loop looped-while looped-while-full))

src/compile/block.rs

@@ -8,8 +8,8 @@ use crate::{
         module::CompilationModule,
         syntax::{
             CallExpression, CondExpression, DefmacroExpression, DefunExpression, Expression,
-            FunctionBody, IfExpression, LambdaExpression, LetExpression, SetqExpression,
-            WhileExpression,
+            FunctionBody, IfExpression, LambdaExpression, LetExpression, LoopExpression,
+            PrognExpression, SetqExpression, WhileExpression,
         },
         value::{BuiltinFunction, CompileConstant, CompileValue},
     },
@@ -198,9 +198,9 @@ impl FunctionBlock {
     ) -> Result<(), CompileError> {
         let mut local = LocalBlock::root(self, module);
         for statement in &body.head {
-            local.compile_statement(statement)?;
+            local.compile_statement(statement, None)?;
         }
-        let value = local.compile_expression(&body.tail)?;
+        let value = local.compile_expression(&body.tail, None)?;
         local.compile_return(value)?;
         Ok(())
     }
@@ -324,8 +324,12 @@ impl<'a> LocalBlock<'a> {
         Ok(())
     }
 
-    pub fn compile_statement(&mut self, expression: &Rc<Expression>) -> Result<(), CompileError> {
-        let value = self.compile_expression(expression)?;
+    pub fn compile_statement(
+        &mut self,
+        expression: &Rc<Expression>,
+        break_label: Option<u32>,
+    ) -> Result<(), CompileError> {
+        let value = self.compile_expression(expression, break_label)?;
         self.compile_drop(value)?;
         Ok(())
     }
@@ -379,10 +383,11 @@ impl<'a> LocalBlock<'a> {
         &mut self,
         math: MathInstruction,
         args: &[Rc<Expression>],
+        break_label: Option<u32>,
     ) -> Result<CompileValue, CompileError> {
         // TODO optimize
         for arg in args.iter().rev() {
-            let arg = self.compile_expression(arg)?;
+            let arg = self.compile_expression(arg, break_label)?;
             self.compile_push(arg)?;
         }
         let Some(count) = U::new(args.len() as u32) else {
@@ -396,18 +401,25 @@ impl<'a> LocalBlock<'a> {
         &mut self,
         builtin: BuiltinFunction,
         args: &[Rc<Expression>],
+        break_label: Option<u32>,
     ) -> Result<CompileValue, CompileError> {
         match builtin {
-            BuiltinFunction::Math(math) => self.compile_builtin_math_generic(math, args),
+            BuiltinFunction::Math(math) => {
+                self.compile_builtin_math_generic(math, args, break_label)
+            }
         }
     }
 
-    fn compile_call(&mut self, call: &CallExpression) -> Result<CompileValue, CompileError> {
+    fn compile_call(
+        &mut self,
+        call: &CallExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         match call.callee.as_ref() {
             Expression::Identifier(identifier)
                 if let Some(builtin) = BuiltinFunction::from_identifier(identifier.as_ref()) =>
             {
-                self.compile_call_builtin(builtin, &call.arguments)
+                self.compile_call_builtin(builtin, &call.arguments, break_label)
             }
             _ => {
                 // Push arguments in reverse order
@@ -415,10 +427,10 @@ impl<'a> LocalBlock<'a> {
                     todo!();
                 };
                 for arg in call.arguments.iter().rev() {
-                    let arg = self.compile_expression(arg)?;
+                    let arg = self.compile_expression(arg, break_label)?;
                     self.compile_push(arg)?;
                 }
-                let callee = self.compile_expression(&call.callee)?;
+                let callee = self.compile_expression(&call.callee, break_label)?;
                 self.compile_push(callee)?;
                 self.function.emit(Instruction::Call(count));
                 Ok(CompileValue::Stack)
@@ -426,21 +438,25 @@ impl<'a> LocalBlock<'a> {
         }
     }
 
-    fn compile_if(&mut self, condition: &IfExpression) -> Result<CompileValue, CompileError> {
+    fn compile_if(
+        &mut self,
+        condition: &IfExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         let label_false = self.function.new_label();
         let label_end = self.function.new_label();
         // Emit condition
-        let condition_value = self.compile_expression(&condition.condition)?;
+        let condition_value = self.compile_expression(&condition.condition, break_label)?;
         self.compile_push(condition_value)?;
         self.function.emit(Emitted::Branch(label_false));
         // True branch
-        let true_value = self.compile_expression(&condition.if_true)?;
+        let true_value = self.compile_expression(&condition.if_true, break_label)?;
         self.compile_push(true_value)?;
         self.function.emit(Emitted::Jump(label_end));
         // False branch
         self.function.adjust_label(label_false);
         let false_value = if let Some(if_false) = condition.if_false.as_ref() {
-            self.compile_expression(if_false)?
+            self.compile_expression(if_false, break_label)?
         } else {
             CompileValue::Nil
         };
@@ -449,16 +465,20 @@ impl<'a> LocalBlock<'a> {
         Ok(CompileValue::Stack)
     }
 
-    fn compile_cond(&mut self, condition: &CondExpression) -> Result<CompileValue, CompileError> {
+    fn compile_cond(
+        &mut self,
+        condition: &CondExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         let label_end = self.function.new_label();
 
         for arm in condition.arms.iter() {
             let label_condition_end = self.function.new_label();
-            let condition_value = self.compile_expression(&arm.condition)?;
+            let condition_value = self.compile_expression(&arm.condition, break_label)?;
             self.compile_push(condition_value)?;
             self.function.emit(Emitted::Branch(label_condition_end));
             // Condition true
-            let then_value = self.compile_expression(&arm.then)?;
+            let then_value = self.compile_expression(&arm.then, break_label)?;
             self.compile_push(then_value)?;
             self.function.emit(Emitted::Jump(label_end));
             // Condition false
@@ -467,7 +487,7 @@ impl<'a> LocalBlock<'a> {
 
         // Otherwise branch
         let otherwise_value = if let Some(othwerise_arm) = condition.otherwise_arm.as_ref() {
-            self.compile_expression(othwerise_arm)?
+            self.compile_expression(othwerise_arm, break_label)?
         } else {
             CompileValue::Nil
         };
@@ -477,22 +497,26 @@ impl<'a> LocalBlock<'a> {
         Ok(CompileValue::Stack)
     }
 
-    fn compile_while(&mut self, cloop: &WhileExpression) -> Result<CompileValue, CompileError> {
+    fn compile_while(
+        &mut self,
+        cloop: &WhileExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         let label_start = self.function.new_label();
         let label_end = self.function.new_label();
 
         self.function.adjust_label(label_start);
 
         // Condition
-        let condition_value = self.compile_expression(&cloop.condition)?;
+        let condition_value = self.compile_expression(&cloop.condition, break_label)?;
         self.compile_push(condition_value)?;
         self.function.emit(Emitted::Branch(label_end));
 
         // Body
         for expression in &cloop.body.head {
-            self.compile_statement(expression)?;
+            self.compile_statement(expression, Some(label_end))?;
         }
-        self.compile_statement(&cloop.body.tail)?;
+        self.compile_statement(&cloop.body.tail, Some(label_end))?;
 
         // Jump back
         self.function.emit(Emitted::Jump(label_start));
@@ -501,11 +525,33 @@ impl<'a> LocalBlock<'a> {
         Ok(CompileValue::Nil)
     }
 
-    fn compile_let(&mut self, binding: &LetExpression) -> Result<CompileValue, CompileError> {
+    fn compile_loop(&mut self, cloop: &LoopExpression) -> Result<CompileValue, CompileError> {
+        let label_start = self.function.new_label();
+        let label_end = self.function.new_label();
+
+        self.function.adjust_label(label_start);
+
+        for expression in &cloop.body.head {
+            self.compile_statement(expression, Some(label_end))?;
+        }
+        self.compile_statement(&cloop.body.tail, Some(label_end))?;
+
+        self.function.emit(Emitted::Jump(label_start));
+
+        self.function.adjust_label(label_end);
+
+        Ok(CompileValue::Nil)
+    }
+
+    fn compile_let(
+        &mut self,
+        binding: &LetExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         self.function.push_local_scope();
         let mut indices = vec![];
         for pair in &binding.bindings {
-            let value = self.compile_expression(&pair.value)?;
+            let value = self.compile_expression(&pair.value, break_label)?;
             let index = self.compile_set_local(&pair.identifier, value, binding.sequential)?;
             indices.push(index);
         }
@@ -517,9 +563,9 @@ impl<'a> LocalBlock<'a> {
             }
         }
         for expr in &binding.body.head {
-            self.compile_statement(expr)?;
+            self.compile_statement(expr, break_label)?;
         }
-        let value = self.compile_expression(&binding.body.tail)?;
+        let value = self.compile_expression(&binding.body.tail, break_label)?;
         if let CompileValue::Local(_) = value {
             self.compile_push(value)?;
             self.function.pop_local_scope();
@@ -530,9 +576,13 @@ impl<'a> LocalBlock<'a> {
         }
     }
 
-    fn compile_setq(&mut self, setq: &SetqExpression) -> Result<CompileValue, CompileError> {
+    fn compile_setq(
+        &mut self,
+        setq: &SetqExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
         for pair in setq.pairs.iter() {
-            let value = self.compile_expression(&pair.value)?;
+            let value = self.compile_expression(&pair.value, break_label)?;
             if let Some(index) = self
                 .function
                 .local_scope_ref()
@@ -551,9 +601,34 @@ impl<'a> LocalBlock<'a> {
         Ok(CompileValue::Quote(value.clone()))
     }
 
+    fn compile_progn(
+        &mut self,
+        value: &PrognExpression,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
+        for expression in &value.body.head {
+            self.compile_statement(expression, break_label)?;
+        }
+        self.compile_statement(&value.body.tail, break_label)?;
+        Ok(CompileValue::Nil)
+    }
+
+    fn compile_call_return(
+        &mut self,
+        break_label: Option<u32>,
+    ) -> Result<CompileValue, CompileError> {
+        if let Some(break_label) = break_label {
+            self.function.emit(Emitted::Jump(break_label));
+        } else {
+            todo!("return outside a loop")
+        }
+        Ok(CompileValue::Nil)
+    }
+
     pub fn compile_expression(
         &mut self,
         expression: &Rc<Expression>,
+        break_label: Option<u32>,
     ) -> Result<CompileValue, CompileError> {
         match expression.as_ref() {
             Expression::Nil => Ok(CompileValue::Nil),
@@ -563,14 +638,17 @@ impl<'a> LocalBlock<'a> {
             Expression::Identifier(identifier) => self.compile_identifier(identifier),
             Expression::Lambda(lambda) => self.compile_lambda(lambda),
             Expression::Defun(defun) => self.compile_defun(defun),
-            Expression::Call(call) => self.compile_call(call),
-            Expression::If(condition) => self.compile_if(condition),
-            Expression::Cond(condition) => self.compile_cond(condition),
-            Expression::Let(binding) => self.compile_let(binding),
-            Expression::Setq(setq) => self.compile_setq(setq),
+            Expression::Call(call) => self.compile_call(call, break_label),
+            Expression::If(condition) => self.compile_if(condition, break_label),
+            Expression::Cond(condition) => self.compile_cond(condition, break_label),
+            Expression::Let(binding) => self.compile_let(binding, break_label),
+            Expression::Setq(setq) => self.compile_setq(setq, break_label),
             Expression::Defmacro(defmacro) => self.compile_defmacro(defmacro),
             Expression::Quote(quote) => self.compile_quote(quote),
-            Expression::While(cloop) => self.compile_while(cloop),
+            Expression::While(cloop) => self.compile_while(cloop, break_label),
+            Expression::Loop(cloop) => self.compile_loop(cloop),
+            Expression::Progn(progn) => self.compile_progn(progn, break_label),
+            Expression::Return => self.compile_call_return(break_label),
 
             Expression::SyntaxError(_) => unreachable!(),
         }
@@ -646,7 +724,9 @@ mod tests {
             max_local_index: 0,
         };
         let mut local = LocalBlock::root(&mut function, &mut module);
-        let value = local.compile_expression(&Rc::new(expression)).unwrap();
+        let value = local
+            .compile_expression(&Rc::new(expression), None)
+            .unwrap();
         let trace = module.options.trace_compile;
         (module, function.resolve_labels(trace), value)
     }

src/compile/syntax/function.rs

@@ -21,6 +21,11 @@ pub struct DefunExpression {
     pub body: FunctionBody,
 }
 
+#[derive(Debug, PartialEq)]
+pub struct PrognExpression {
+    pub body: FunctionBody,
+}
+
 impl FunctionSignature {
     pub(super) fn parse(mut value: &Value, input: &Value) -> Result<Self, ParseError> {
         enum Mode {
@@ -171,6 +176,13 @@ impl FunctionBody {
     }
 }
 
+impl PrognExpression {
+    pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
+        let body = FunctionBody::parse(value, input, Keyword::Progn)?;
+        Ok(Self { body })
+    }
+}
+
 impl DefunExpression {
     pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
         let Value::Cons(value) = value else {
@@ -232,6 +244,12 @@ impl CollectErrors<ParseError> for DefunExpression {
     }
 }
 
+impl CollectErrors<ParseError> for PrognExpression {
+    fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
+        self.body.collect_errors(errors)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::rc::Rc;

src/compile/syntax/loops.rs

@@ -14,6 +14,11 @@ pub struct WhileExpression {
     pub body: FunctionBody,
 }
 
+#[derive(Debug, PartialEq)]
+pub struct LoopExpression {
+    pub body: FunctionBody,
+}
+
 impl WhileExpression {
     pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
         let Value::Cons(cons) = value else {
@@ -32,6 +37,13 @@ impl WhileExpression {
     }
 }
 
+impl LoopExpression {
+    pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
+        let body = FunctionBody::parse(value, input, Keyword::Loop)?;
+        Ok(Self { body })
+    }
+}
+
 impl CollectErrors<ParseError> for WhileExpression {
     fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
         let a = self.condition.collect_errors(errors);
@@ -39,3 +51,9 @@ impl CollectErrors<ParseError> for WhileExpression {
         a | b
     }
 }
+
+impl CollectErrors<ParseError> for LoopExpression {
+    fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
+        self.body.collect_errors(errors)
+    }
+}

src/compile/syntax/mod.rs

@@ -38,6 +38,9 @@ pub enum Expression {
     SyntaxError(ParseError),
     Quote(Rc<Value>),
     While(WhileExpression),
+    Loop(LoopExpression),
+    Progn(PrognExpression),
+    Return,
 }
 
 impl Expression {
@@ -100,9 +103,16 @@ impl Expression {
                         };
                         Rc::new(Self::Quote(value.into()))
                     }
+                    Value::Keyword(Keyword::Progn) => {
+                        Self::map_or(PrognExpression::parse(cdr, value), Expression::Progn)
+                    }
+                    Value::Keyword(Keyword::Loop) => {
+                        Self::map_or(LoopExpression::parse(cdr, value), Expression::Loop)
+                    }
                     Value::Keyword(Keyword::While) => {
                         Self::map_or(WhileExpression::parse(cdr, value), Expression::While)
                     }
+                    Value::Keyword(Keyword::Return) => Rc::new(Self::Return),
                     _ => Self::map_or(CallExpression::parse(cons, value), Expression::Call),
                 }
             }
@@ -130,7 +140,10 @@ impl CollectErrors<ParseError> for Expression {
             Self::Setq(setq) => setq.collect_errors(errors),
             Self::Defmacro(defmacro) => defmacro.collect_errors(errors),
             Self::While(cloop) => cloop.collect_errors(errors),
+            Self::Loop(cloop) => cloop.collect_errors(errors),
+            Self::Progn(progn) => progn.collect_errors(errors),
             Self::Nil
+            | Self::Return
             | Self::IntegerLiteral(_)
             | Self::Identifier(_)
             | Self::BooleanLiteral(_)

src/vm/value/keyword.rs

@@ -47,6 +47,9 @@ impl_keyword! {
         LetStar => "let*",
         Defmacro => "defmacro",
         Quote => "quote",
+        Progn => "progn",
+        Loop => "loop",
+        Return => "return",
         // Cons => "cons",
     }
 }