Implement loop/return forms

examples/factorial.lysp

@@ -7,18 +7,20 @@
     )
   )
 
-(print 1  "\t" (factorial 1))
-(print 2  "\t" (factorial 2))
-(print 3  "\t" (factorial 3))
-(print 4  "\t" (factorial 4))
-(print 5  "\t" (factorial 5))
-(print 6  "\t" (factorial 6))
-(print 7  "\t" (factorial 7))
-(print 8  "\t" (factorial 8))
-(print 9  "\t" (factorial 9))
-(print 10 "\t" (factorial 10))
-(print 11 "\t" (factorial 11))
-(print 12 "\t" (factorial 12))
-(print 13 "\t" (factorial 13))
-(print 14 "\t" (factorial 14))
-(print 15 "\t" (factorial 15))
+(defun loop-factorial (x)
+  (let (i 0 acc 1)
+    (while (< i x)
+      (setq i (+ i 1))
+      (setq acc (* acc i))
+      )
+    acc
+    )
+  )
+
+;; TODO for loops?
+(let (i 0)
+  (while (<= i 15)
+    (print i "\t" (factorial i) "\t" (loop-factorial i))
+    (setq i (+ i 1))
+    )
+  )

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))

examples/scopes.lysp

@@ -0,0 +1,25 @@
+;; vi:ft=lisp:sw=2:ts=2
+
+(defun replace-argument (x) (let (x 1) x))
+
+(defun replace-let ()
+  (let (x 1)
+    (setq x 2) ;; mutate local x
+    x
+    )
+  )
+
+(defun reassignment-in-a-loop (iterations)
+  (let (x 0 y 1)
+    (while (< x iterations)
+      (setq y (* y 2))
+      (setq x (+ x 1))
+      )
+    y
+    )
+  )
+
+(assert (= (replace-argument 2) 1))
+(assert (= (replace-argument 3) 1))
+(assert (= (replace-let) 2))
+(assert (= (reassignment-in-a-loop 4) 16))

src/compile/block.rs

@@ -8,7 +8,8 @@ use crate::{
         module::CompilationModule,
         syntax::{
             CallExpression, CondExpression, DefmacroExpression, DefunExpression, Expression,
-            FunctionBody, IfExpression, LambdaExpression, LetExpression, SetqExpression,
+            FunctionBody, IfExpression, LambdaExpression, LetExpression, LoopExpression,
+            PrognExpression, SetqExpression, WhileExpression,
         },
         value::{BuiltinFunction, CompileConstant, CompileValue},
     },
@@ -136,13 +137,19 @@ impl FunctionBlock {
         self.labels.insert(label, self.instructions.len());
     }
 
-    pub fn resolve_labels(self) -> CompiledFunction {
+    pub fn resolve_labels(self, trace: bool) -> CompiledFunction {
         let mut instructions = vec![];
-        // eprintln!("RESOLVE LABELS");
+
+        if trace {
+            eprintln!("Instruction resolution:");
+        }
+
         for (function_offset, emitted) in self.instructions.into_iter().enumerate() {
             match emitted {
                 Emitted::Instruction(instruction) => {
-                    // eprintln!("{function_offset}: {instruction:?}");
+                    if trace {
+                        eprintln!("{function_offset}: {instruction:?}");
+                    }
                     instructions.push(instruction);
                 }
                 Emitted::Branch(label) => {
@@ -151,9 +158,11 @@ impl FunctionBlock {
                         todo!()
                     }
                     let diff = address.checked_signed_diff(function_offset).expect("TODO");
-                    // eprintln!(
-                    //     "{function_offset}: branch to label {label} (@ {address}) -> {diff:+}"
-                    // );
+                    if trace {
+                        eprintln!(
+                            "{function_offset}: Branch: label {label} (@ {address}) -> {diff:+}"
+                        );
+                    }
                     let offset = U::from_signed(diff as i64).expect("TODO");
                     instructions.push(Instruction::Branch(offset));
                 }
@@ -163,7 +172,11 @@ impl FunctionBlock {
                         todo!()
                     }
                     let diff = address.checked_signed_diff(function_offset).expect("TODO");
-                    // eprintln!("{function_offset}: jump to label {label} (@ {address}) -> {diff:+}");
+                    if trace {
+                        eprintln!(
+                            "{function_offset}: jump to label {label} (@ {address}) -> {diff:+}"
+                        );
+                    }
                     let offset = U::from_signed(diff as i64).expect("TODO");
                     instructions.push(Instruction::Jump(offset));
                 }
@@ -185,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(())
     }
@@ -311,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(())
     }
@@ -366,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 {
@@ -383,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
@@ -402,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)
@@ -413,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
         };
@@ -436,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
@@ -454,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
         };
@@ -464,11 +497,61 @@ impl<'a> LocalBlock<'a> {
         Ok(CompileValue::Stack)
     }
 
-    fn compile_let(&mut self, binding: &LetExpression) -> 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, 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, Some(label_end))?;
+        }
+        self.compile_statement(&cloop.body.tail, Some(label_end))?;
+
+        // Jump back
+        self.function.emit(Emitted::Jump(label_start));
+
+        self.function.adjust_label(label_end);
+        Ok(CompileValue::Nil)
+    }
+
+    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);
         }
@@ -480,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();
@@ -493,10 +576,23 @@ 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)?;
-            self.compile_set_global(&pair.identifier, value)?;
+            let value = self.compile_expression(&pair.value, break_label)?;
+            if let Some(index) = self
+                .function
+                .local_scope_ref()
+                .and_then(|scope| scope.get(&pair.identifier))
+            {
+                self.compile_push(value)?;
+                self.function.emit(Instruction::SetLocal(index));
+            } else {
+                self.compile_set_global(&pair.identifier, value)?;
+            }
         }
         Ok(CompileValue::Nil)
     }
@@ -505,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),
@@ -517,13 +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, 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!(),
         }
@@ -599,8 +724,11 @@ mod tests {
             max_local_index: 0,
         };
         let mut local = LocalBlock::root(&mut function, &mut module);
-        let value = local.compile_expression(&Rc::new(expression)).unwrap();
-        (module, function.resolve_labels(), value)
+        let value = local
+            .compile_expression(&Rc::new(expression), None)
+            .unwrap();
+        let trace = module.options.trace_compile;
+        (module, function.resolve_labels(trace), value)
     }
 
     #[test]
@@ -647,9 +775,9 @@ mod tests {
             &f.instructions,
             &[
                 Instruction::PushBool(false),             // 0
-                Instruction::Branch(U::truncate(4)),      // 1
+                Instruction::Branch(U::truncate(3)),      // 1
                 Instruction::PushInteger(U::truncate(1)), // 2
-                Instruction::Jump(U::truncate(5)),        // 3
+                Instruction::Jump(U::truncate(2)),        // 3
                 Instruction::PushInteger(U::truncate(2))  // 4
                                                           // 5
             ]
@@ -676,13 +804,13 @@ mod tests {
             &f.instructions,
             &[
                 Instruction::PushBool(false),             // 0
-                Instruction::Branch(U::truncate(4)),      // 1
+                Instruction::Branch(U::truncate(3)),      // 1
                 Instruction::PushInteger(U::truncate(1)), // 2
-                Instruction::Jump(U::truncate(9)),        // 3
+                Instruction::Jump(U::truncate(6)),        // 3
                 Instruction::PushBool(true),              // 4
-                Instruction::Branch(U::truncate(8)),      // 5
+                Instruction::Branch(U::truncate(3)),      // 5
                 Instruction::PushInteger(U::truncate(2)), // 6
-                Instruction::Jump(U::truncate(9)),        // 7
+                Instruction::Jump(U::truncate(2)),        // 7
                 Instruction::PushInteger(U::truncate(3)), // 8
                                                           // 9
             ]

src/compile/mod.rs

@@ -13,3 +13,8 @@ pub use syntax::{
     CallExpression, ExpectedWhat, ExpectedWhere, Expression, FunctionBody, LambdaExpression,
     ParseError, ParseErrorKind,
 };
+
+#[derive(Default, Clone)]
+pub struct CompileOptions {
+    pub trace_compile: bool,
+}

src/compile/module.rs

@@ -2,6 +2,7 @@ use std::{collections::HashMap, rc::Rc};
 
 use crate::{
     compile::{
+        CompileOptions,
         block::{CompiledFunction, FunctionBlock},
         error::CompileError,
         function::FunctionSignature,
@@ -19,12 +20,20 @@ use crate::{
 pub struct CompilationModule {
     pub(crate) constant_pool: Pool<CompileConstant, { ConstantId::BITS }>,
     pub(crate) local_functions: HashMap<u32, CompiledFunction>,
+    pub(crate) options: CompileOptions,
     macros: HashMap<Rc<str>, DefmacroExpression>,
     local_function_index: u32,
     root: Option<u32>,
 }
 
 impl CompilationModule {
+    pub fn new(options: CompileOptions) -> Self {
+        Self {
+            options,
+            ..Default::default()
+        }
+    }
+
     pub fn define_macro(&mut self, defmacro: DefmacroExpression) {
         self.macros.insert(defmacro.name.clone(), defmacro);
     }
@@ -49,7 +58,7 @@ impl CompilationModule {
         self.local_function_index += 1;
         let mut function = FunctionBlock::new(signature);
         function.compile_body(self, body)?;
-        let function = function.resolve_labels();
+        let function = function.resolve_labels(self.options.trace_compile);
         self.local_functions.insert(index, function);
         if root {
             self.root = Some(index);

src/compile/syntax/condition.rs

@@ -64,9 +64,9 @@ impl IfExpression {
             Some(Expression::parse_inner(if_false))
         };
         Ok(Self {
-            condition: condition.into(),
-            if_true: if_true.into(),
-            if_false: if_false.map(Into::into),
+            condition,
+            if_true,
+            if_false,
         })
     }
 }

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

@@ -0,0 +1,59 @@
+use std::rc::Rc;
+
+use crate::{
+    compile::{
+        ExpectedWhat, ExpectedWhere, Expression, FunctionBody, ParseError, ParseErrorKind,
+        syntax::CollectErrors,
+    },
+    vm::value::{ConsCell, Keyword, Value},
+};
+
+#[derive(Debug, PartialEq)]
+pub struct WhileExpression {
+    pub condition: Rc<Expression>,
+    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 {
+            return Err(ParseError {
+                input: input.clone(),
+                error: ParseErrorKind::Expected(
+                    ExpectedWhat::ProperList,
+                    ExpectedWhere::AfterKeyword(Keyword::While),
+                ),
+            });
+        };
+        let ConsCell(car, cdr) = cons.as_ref();
+        let condition = Expression::parse_inner(car);
+        let body = FunctionBody::parse(cdr, input, Keyword::While)?;
+        Ok(Self { condition, body })
+    }
+}
+
+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);
+        let b = self.body.collect_errors(errors);
+        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

@@ -8,6 +8,7 @@ mod condition;
 mod error;
 mod function;
 mod lambda;
+mod loops;
 mod macros;
 
 pub use binding::*;
@@ -16,6 +17,7 @@ pub use condition::*;
 pub use error::*;
 pub use function::*;
 pub use lambda::*;
+pub use loops::*;
 pub use macros::*;
 
 #[derive(Debug, PartialEq)]
@@ -35,6 +37,10 @@ pub enum Expression {
     Defmacro(DefmacroExpression),
     SyntaxError(ParseError),
     Quote(Rc<Value>),
+    While(WhileExpression),
+    Loop(LoopExpression),
+    Progn(PrognExpression),
+    Return,
 }
 
 impl Expression {
@@ -97,6 +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),
                 }
             }
@@ -123,7 +139,11 @@ impl CollectErrors<ParseError> for Expression {
             Self::Let(let_) => let_.collect_errors(errors),
             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/main.rs

@@ -3,11 +3,12 @@ use std::{
     io::{self, BufReader},
     path::{Path, PathBuf},
     process::ExitCode,
+    str::FromStr,
 };
 
 use clap::Parser;
 use lysp::{
-    compile::{CompileError, ParseError},
+    compile::{CompileError, CompileOptions, ParseError},
     error::{EvalError, MachineErrorKind},
     read::{InteractiveReader, ModuleReader, read},
     util::Either,
@@ -25,8 +26,32 @@ enum Error {
     Printed,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum Trace {
+    Compile,
+    Execute,
+}
+
+impl FromStr for Trace {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "compile" => Ok(Self::Compile),
+            "execute" => Ok(Self::Execute),
+            _ => Err(format!("Unknown trace flag: {s:?}")),
+        }
+    }
+}
+
 #[derive(Debug, Parser)]
 struct Args {
+    #[clap(
+        short = 'T',
+        long = "trace",
+        help = "Enable tracing of execution/compilation steps"
+    )]
+    trace: Vec<Trace>,
     module: Option<PathBuf>,
 }
 
@@ -88,8 +113,13 @@ fn handle_module_error(input: Either<EvalError, Vec<ParseError>>) -> Error {
     }
 }
 
-fn eval(vm: &mut Machine, env: &mut Environment, value: Value) -> Option<Value> {
-    let result = vm.eval_value(env, value.clone());
+fn eval(
+    options: &CompileOptions,
+    vm: &mut Machine,
+    env: &mut Environment,
+    value: Value,
+) -> Option<Value> {
+    let result = vm.eval_value(options.clone(), env, value.clone());
     match result {
         Ok(r) => Some(r),
         Err(error) => {
@@ -99,7 +129,11 @@ fn eval(vm: &mut Machine, env: &mut Environment, value: Value) -> Option<Value>
     }
 }
 
-fn run_interactive(vm: &mut Machine, env: &mut Environment) -> Result<(), Error> {
+fn run_interactive(
+    compile_options: &CompileOptions,
+    vm: &mut Machine,
+    env: &mut Environment,
+) -> Result<(), Error> {
     let mut reader = InteractiveReader::new("> ", ">> ");
 
     loop {
@@ -111,7 +145,7 @@ fn run_interactive(vm: &mut Machine, env: &mut Environment) -> Result<(), Error>
                 continue;
             }
         };
-        if let Some(value) = eval(vm, env, value) {
+        if let Some(value) = eval(compile_options, vm, env, value) {
             println!("== {value}");
         } else {
             reader.reset();
@@ -122,6 +156,7 @@ fn run_interactive(vm: &mut Machine, env: &mut Environment) -> Result<(), Error>
 }
 
 fn run_module<P: AsRef<Path>>(
+    compile_options: &CompileOptions,
     vm: &mut Machine,
     env: &mut Environment,
     path: P,
@@ -129,7 +164,7 @@ fn run_module<P: AsRef<Path>>(
     let path = path.as_ref();
     let reader = BufReader::new(File::open(path)?);
     let module_reader = ModuleReader::new(reader);
-    let module = match module_reader.compile(env) {
+    let module = match module_reader.compile(compile_options, env) {
         Ok(module) => module,
         Err(error) => return Err(handle_module_error(error)),
     };
@@ -143,11 +178,15 @@ fn run_module<P: AsRef<Path>>(
 fn main() -> ExitCode {
     let args = Args::parse();
     let mut vm = Machine::default();
+    let compile_options = CompileOptions {
+        trace_compile: args.trace.contains(&Trace::Compile),
+    };
+    vm.trace_instructions = args.trace.contains(&Trace::Execute);
     let mut env = Environment::default();
     prelude::load(&mut env);
     let result = match args.module.as_ref() {
-        Some(module) => run_module(&mut vm, &mut env, module),
-        None => run_interactive(&mut vm, &mut env),
+        Some(module) => run_module(&compile_options, &mut vm, &mut env, module),
+        None => run_interactive(&compile_options, &mut vm, &mut env),
     };
     match result {
         Ok(()) => ExitCode::SUCCESS,

src/read.rs

@@ -5,7 +5,9 @@ use std::{
 };
 
 use crate::{
-    compile::{CompilationModule, Expression, FunctionBody, FunctionSignature, ParseError},
+    compile::{
+        CompilationModule, CompileOptions, Expression, FunctionBody, FunctionSignature, ParseError,
+    },
     error::EvalError,
     parse::{self, parse_value},
     util::Either,
@@ -106,6 +108,7 @@ impl<R: BufRead> ModuleReader<R> {
 
     pub fn read_expression(
         &mut self,
+        options: &CompileOptions,
         env: &mut Environment,
     ) -> Result<Option<Rc<Expression>>, Either<EvalError, Vec<ParseError>>> {
         loop {
@@ -117,7 +120,7 @@ impl<R: BufRead> ModuleReader<R> {
             let expression = Expression::parse(&value).map_err(Either::Right)?;
             if let Expression::Defmacro(_) = expression.as_ref() {
                 self.macro_machine
-                    .eval_value(env, value)
+                    .eval_value(options.clone(), env, value)
                     .map_err(Either::Left)?;
                 continue;
             }
@@ -127,9 +130,10 @@ impl<R: BufRead> ModuleReader<R> {
 
     pub fn compile(
         mut self,
+        options: &CompileOptions,
         env: &mut Environment,
     ) -> Result<ModuleRef, Either<EvalError, Vec<ParseError>>> {
-        let mut module = CompilationModule::default();
+        let mut module = CompilationModule::new(options.clone());
         let mut body = FunctionBody {
             head: vec![],
             tail: Rc::new(Expression::Nil),
@@ -138,7 +142,7 @@ impl<R: BufRead> ModuleReader<R> {
         let mut syntax_errors = vec![];
 
         loop {
-            let expression = match self.read_expression(env) {
+            let expression = match self.read_expression(options, env) {
                 Ok(Some(expression)) => expression,
                 Ok(None) => break,
                 Err(Either::Left(error)) => return Err(Either::Left(error)),

src/vm/instruction.rs

@@ -91,7 +91,7 @@ pub enum Instruction {
 }
 
 pub type ConstantId = U<12>;
-pub type FunctionOffset = U<12>;
+pub type FunctionOffset = U<10>;
 pub type LocalId = U<8>;
 pub type ArgumentId = U<6>;
 pub type ArgumentCount = U<6>;
@@ -222,7 +222,7 @@ impl TryFrom<u32> for Instruction {
 
 #[cfg(test)]
 mod tests {
-    use crate::vm::instruction::U;
+    use crate::vm::instruction::{FunctionOffset, U};
 
     #[test]
     fn test_u_convert() {
@@ -238,4 +238,10 @@ mod tests {
         let t = T::from_signed(-0x1000001);
         assert!(t.is_none());
     }
+
+    #[test]
+    fn test_branch_target_extend() {
+        let target = FunctionOffset::truncate(1022);
+        assert_eq!(target.sign_extend_i64(), -2);
+    }
 }

src/vm/machine.rs

@@ -1,6 +1,7 @@
 use std::{collections::HashMap, fmt};
 
 use crate::{
+    compile::CompileOptions,
     error::{EvalError, MachineError, MachineErrorKind},
     vm::{
         env::Environment,
@@ -516,11 +517,12 @@ impl Machine {
 
     pub fn eval_value(
         &mut self,
+        compile_options: CompileOptions,
         environment: &mut Environment,
         value: Value,
     ) -> Result<Value, EvalError> {
         let value = value.macro_expand(self, environment, false)?;
-        let module = Module::compile_value(&value)?;
+        let module = Module::compile_value(compile_options, &value)?;
         let module = ModuleRef::from(module);
         self.eval_module(environment, module)
     }

src/vm/module.rs

@@ -7,7 +7,10 @@ use std::{
 };
 
 use crate::{
-    compile::{CompilationModule, CompileError, Expression, FunctionBody, FunctionSignature},
+    compile::{
+        CompilationModule, CompileError, CompileOptions, Expression, FunctionBody,
+        FunctionSignature,
+    },
     vm::{
         instruction::{ConstantId, Instruction},
         pool::Pool,
@@ -122,9 +125,9 @@ impl Module {
         self.entry
     }
 
-    pub fn compile_value(value: &Value) -> Result<Self, CompileError> {
+    pub fn compile_value(options: CompileOptions, value: &Value) -> Result<Self, CompileError> {
         let expression = Expression::parse(value).map_err(CompileError::Parse)?;
-        let mut module = CompilationModule::default();
+        let mut module = CompilationModule::new(options);
         module.compile_function(
             FunctionSignature::EMPTY,
             &FunctionBody {
@@ -215,10 +218,13 @@ impl fmt::Display for ModuleConstant {
 
 #[cfg(test)]
 mod tests {
-    use crate::vm::{
-        instruction::{Instruction, MathInstruction, U},
-        module::Module,
-        value::Value,
+    use crate::{
+        compile::CompileOptions,
+        vm::{
+            instruction::{Instruction, MathInstruction, U},
+            module::Module,
+            value::Value,
+        },
     };
 
     #[test]
@@ -228,7 +234,7 @@ mod tests {
             Value::Integer(1),
             Value::Integer(2),
         ]);
-        let m = Module::compile_value(&v).unwrap();
+        let m = Module::compile_value(CompileOptions::default(), &v).unwrap();
         assert!(m.constants.is_empty());
         let is = [
             Instruction::PushInteger(U::truncate(2)),

src/vm/prelude/mod.rs

@@ -102,7 +102,7 @@ pub fn load(env: &mut Environment) {
             [_, _] => todo!(),
             _ => todo!(),
         };
-        let value = match vm.eval_value(env, value.clone()) {
+        let value = match vm.eval_value(Default::default(), env, value.clone()) {
             Ok(result) => result,
             _ => todo!(),
         };

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",
     }
 }

tests/integration.rs

@@ -30,7 +30,7 @@ fn eval_str_in(code: &str, env: &mut Environment) -> Result<Value, EvalError> {
             Err(error) => panic!("{error}"),
         };
 
-        last_value = Some(machine.eval_value(env, value));
+        last_value = Some(machine.eval_value(Default::default(), env, value));
     }
 
     last_value.expect("no expressions evaluated")