Fix bugs in the VM and compiler

dust-cli/src/main.rs

@@ -331,10 +331,6 @@ fn main() {
         let chunk = compiler.finish(name.or(file_name));
         let compile_end = start_time.elapsed();
 
-        if time {
-            print_time(compile_end);
-        }
-
         let vm = Vm::new(chunk);
         let return_value = vm.run();
         let run_end = start_time.elapsed();
@@ -347,30 +343,33 @@ fn main() {
 
         if time {
             let run_time = run_end - compile_end;
+            let total_time = compile_end + run_time;
 
-            print_time(run_time);
+            print_time("Compile Time", compile_end);
+            print_time("Run Time", run_time);
+            print_time("Total Time", total_time);
         }
     }
 }
 
-fn print_time(instant: Duration) {
+fn print_time(phase: &str, instant: Duration) {
     let seconds = instant.as_secs_f64();
 
     match seconds {
         ..=0.001 => {
             println!(
-                "Compile time: {microseconds} microseconds",
+                "{phase:12}: {microseconds}µs",
-                microseconds = seconds * 1_000_000.0
+                microseconds = (seconds * 1_000_000.0).round()
             );
         }
-        ..=0.1 => {
+        ..=0.199 => {
             println!(
-                "Compile time: {milliseconds} milliseconds",
+                "{phase:12}: {milliseconds}ms",
-                milliseconds = seconds * 1000.0
+                milliseconds = (seconds * 1000.0).round()
             );
         }
         _ => {
-            println!("Compile time: {seconds} seconds");
+            println!("{phase:12}: {seconds}s");
         }
     }
 }

dust-lang/src/chunk/mod.rs

@@ -119,20 +119,10 @@ impl Chunk {
             self.record_index,
         );
 
-        if records.is_empty() {
+        records.push(record);
-            records.push(record);
 
-            for chunk in self.prototypes {
+        for chunk in self.prototypes {
-                chunk.into_records(records);
+            chunk.into_records(records);
-            }
-        } else {
-            for chunk in self.prototypes {
-                chunk.into_records(records);
-            }
-
-            debug_assert!(record.index() as usize == records.len());
-
-            records.push(record);
         }
     }
 

dust-lang/src/compiler/mod.rs

@@ -53,6 +53,7 @@ use crate::{
 pub fn compile(source: &str) -> Result<Chunk, DustError> {
     let lexer = Lexer::new(source);
     let mut compiler = Compiler::new(lexer).map_err(|error| DustError::compile(error, source))?;
+    compiler.is_main_chunk = true;
 
     compiler
         .compile()
@@ -102,11 +103,6 @@ pub struct Compiler<'src> {
     /// [`Compiler::finish`] is called.
     stack_size: usize,
 
-    current_token: Token<'src>,
-    current_position: Span,
-    previous_token: Token<'src>,
-    previous_position: Span,
-
     /// The first register index that the compiler should use. This is used to avoid reusing the
     /// registers that are used for the function's arguments, thus it is zero in the program's main
     /// chunk.
@@ -128,6 +124,14 @@ pub struct Compiler<'src> {
     /// maintain the depth-first index. After the function is compiled, its `next_record_index`
     /// is assigned back to this chunk.
     next_record_index: u8,
+
+    /// Whether the compiler is compiling the main chunk.
+    is_main_chunk: bool,
+
+    current_token: Token<'src>,
+    current_position: Span,
+    previous_token: Token<'src>,
+    previous_position: Span,
 }
 
 impl<'src> Compiler<'src> {
@@ -148,15 +152,16 @@ impl<'src> Compiler<'src> {
             prototypes: Vec::new(),
             stack_size: 0,
             lexer,
-            current_token,
-            current_position,
-            previous_token: Token::Eof,
-            previous_position: Span(0, 0),
             minimum_register: 0,
             block_index: 0,
             current_scope: Scope::default(),
             record_index: 0,
             next_record_index: 1,
+            is_main_chunk: true,
+            current_token,
+            current_position,
+            previous_token: Token::Eof,
+            previous_position: Span(0, 0),
         })
     }
 
@@ -437,6 +442,15 @@ impl<'src> Compiler<'src> {
     ///
     /// If [`Self::type`] is already set, it will check if the given [Type] is compatible.
     fn update_return_type(&mut self, new_return_type: Type) -> Result<(), CompileError> {
+        if !self.is_main_chunk {
+            Type::function(self.r#type.clone())
+                .check(&new_return_type)
+                .map_err(|conflict| CompileError::ReturnTypeConflict {
+                    conflict,
+                    position: self.current_position,
+                })?;
+        }
+
         if self.r#type.return_type != Type::None {
             self.r#type
                 .return_type
@@ -1427,8 +1441,10 @@ impl<'src> Compiler<'src> {
                 true
             };
         let end = self.current_position.1;
+        let return_register = self.next_register() - 1;
         let r#return = Instruction::from(Return {
             should_return_value,
+            return_register,
         });
 
         self.emit_instruction(r#return, Type::None, Span(start, end));
@@ -1454,22 +1470,27 @@ impl<'src> Compiler<'src> {
 
     fn parse_implicit_return(&mut self) -> Result<(), CompileError> {
         if self.allow(Token::Semicolon)? {
-            let r#return = Instruction::r#return(false);
+            let r#return = Instruction::r#return(false, 0);
 
             self.emit_instruction(r#return, Type::None, self.current_position);
         } else {
-            let previous_expression_type =
+            let (previous_expression_type, previous_register) = self
-                self.instructions
+                .instructions
-                    .last()
+                .last()
-                    .map_or(Type::None, |(instruction, r#type, _)| {
+                .map(|(instruction, r#type, _)| {
-                        if instruction.yields_value() {
+                    if instruction.yields_value() {
-                            r#type.clone()
+                        (r#type.clone(), instruction.a_field())
-                        } else {
+                    } else {
-                            Type::None
+                        (Type::None, 0)
-                        }
+                    }
-                    });
+                })
+                .ok_or_else(|| CompileError::ExpectedExpression {
+                    found: self.previous_token.to_owned(),
+                    position: self.previous_position,
+                })?;
+
             let should_return_value = previous_expression_type != Type::None;
-            let r#return = Instruction::r#return(should_return_value);
+            let r#return = Instruction::r#return(should_return_value, previous_register);
 
             self.update_return_type(previous_expression_type.clone())?;
             self.emit_instruction(r#return, Type::None, self.current_position);

dust-lang/src/instruction/mod.rs

@@ -395,9 +395,10 @@ impl Instruction {
         })
     }
 
-    pub fn r#return(should_return_value: bool) -> Instruction {
+    pub fn r#return(should_return_value: bool, return_register: u8) -> Instruction {
         Instruction::from(Return {
             should_return_value,
+            return_register,
         })
     }
 
@@ -732,12 +733,13 @@ impl Instruction {
             Operation::RETURN => {
                 let Return {
                     should_return_value,
+                    return_register,
                 } = Return::from(self);
 
                 if should_return_value {
-                    "RETURN".to_string()
+                    format!("RETURN R{return_register}")
                 } else {
-                    "".to_string()
+                    String::new()
                 }
             }
             _ => {

dust-lang/src/instruction/return.rs

@@ -1,15 +1,33 @@
 use crate::{Instruction, Operation};
 
+use super::InstructionData;
+
 pub struct Return {
     pub should_return_value: bool,
+    pub return_register: u8,
+}
+
+impl From<InstructionData> for Return {
+    fn from(data: InstructionData) -> Self {
+        let InstructionData {
+            b_field, c_field, ..
+        } = data;
+
+        Return {
+            should_return_value: b_field != 0,
+            return_register: c_field,
+        }
+    }
 }
 
 impl From<&Instruction> for Return {
     fn from(instruction: &Instruction) -> Self {
         let should_return_value = instruction.b_field() != 0;
+        let return_register = instruction.c_field();
 
         Return {
             should_return_value,
+            return_register,
         }
     }
 }
@@ -18,7 +36,8 @@ impl From<Return> for Instruction {
     fn from(r#return: Return) -> Self {
         let operation = Operation::RETURN;
         let b = r#return.should_return_value as u8;
+        let c = r#return.return_register;
 
-        Instruction::new(operation, 0, b, 0, false, false, false)
+        Instruction::new(operation, 0, b, c, false, false, false)
     }
 }

dust-lang/src/native_function/io.rs

@@ -1,4 +1,4 @@
-use std::io::{stdin, stdout, Write};
+use std::io::{self, stdin, stdout, Write};
 use std::ops::Range;
 
 use crate::vm::{Register, ThreadSignal};
@@ -66,31 +66,22 @@ pub fn write_line(
     _destination: Option<u8>,
     argument_range: Range<u8>,
 ) -> Result<ThreadSignal, NativeFunctionError> {
+    let map_err = |io_error: io::Error| NativeFunctionError::Io {
+        error: io_error.kind(),
+        position: record.current_position(),
+    };
     let mut stdout = stdout().lock();
 
     for register_index in argument_range {
         if let Some(value) = record.open_register_allow_empty(register_index) {
             let string = value.display(record);
 
-            stdout
+            stdout.write(string.as_bytes()).map_err(map_err)?;
-                .write(string.as_bytes())
+            stdout.write(b"\n").map_err(map_err)?;
-                .map_err(|io_error| NativeFunctionError::Io {
-                    error: io_error.kind(),
-                    position: record.current_position(),
-                })?;
-            stdout
-                .write(b"\n")
-                .map_err(|io_error| NativeFunctionError::Io {
-                    error: io_error.kind(),
-                    position: record.current_position(),
-                })?;
         }
     }
 
-    stdout.flush().map_err(|io_error| NativeFunctionError::Io {
+    stdout.flush().map_err(map_err)?;
-        error: io_error.kind(),
-        position: record.current_position(),
-    })?;
 
     Ok(ThreadSignal::Continue)
 }

dust-lang/src/vm/record.rs

@@ -121,8 +121,6 @@ impl Record {
     }
 
     pub fn open_register(&self, register_index: u8) -> &Value {
-        trace!("Open register R{register_index}");
-
         let register_index = register_index as usize;
 
         assert!(
@@ -133,8 +131,16 @@ impl Record {
         let register = &self.stack[register_index];
 
         match register {
-            Register::Value(value) => value,
+            Register::Value(value) => {
-            Register::Pointer(pointer) => self.follow_pointer(*pointer),
+                trace!("Register R{register_index} openned to value {value}");
+
+                value
+            }
+            Register::Pointer(pointer) => {
+                trace!("Open register R{register_index} openned to pointer {pointer}");
+
+                self.follow_pointer(*pointer)
+            }
             Register::Empty => panic!("VM Error: Register {register_index} is empty"),
         }
     }
@@ -152,8 +158,16 @@ impl Record {
         let register = &self.stack[register_index];
 
         match register {
-            Register::Value(value) => Some(value),
+            Register::Value(value) => {
-            Register::Pointer(pointer) => Some(self.follow_pointer(*pointer)),
+                trace!("Register R{register_index} openned to value {value}");
+
+                Some(value)
+            }
+            Register::Pointer(pointer) => {
+                trace!("Open register R{register_index} openned to pointer {pointer}");
+
+                Some(self.follow_pointer(*pointer))
+            }
             Register::Empty => None,
         }
     }

dust-lang/src/vm/run_action.rs

@@ -3,7 +3,7 @@ use tracing::trace;
 use crate::{
     instruction::{
         Call, CallNative, Close, GetLocal, LoadBoolean, LoadConstant, LoadFunction, LoadList,
-        LoadSelf, Point,
+        LoadSelf, Point, Return,
     },
     vm::VmError,
     AbstractList, ConcreteValue, Instruction, InstructionData, Type, Value,
@@ -551,9 +551,15 @@ pub fn call_native(instruction_data: InstructionData, record: &mut Record) -> Th
 }
 
 pub fn r#return(instruction_data: InstructionData, _: &mut Record) -> ThreadSignal {
-    let should_return_value = instruction_data.b_field != 0;
+    let Return {
+        should_return_value,
+        return_register,
+    } = instruction_data.into();
 
-    ThreadSignal::Return(should_return_value)
+    ThreadSignal::Return {
+        should_return_value,
+        return_register,
+    }
 }
 
 #[cfg(test)]

dust-lang/src/vm/thread.rs

@@ -1,11 +1,11 @@
-use tracing::{info, span, trace, Level};
+use tracing::{info, trace};
 
 use crate::{
     vm::{FunctionCall, Register},
     Chunk, DustString, Value,
 };
 
-use super::{record::Record, CallStack, VmError};
+use super::{record::Record, CallStack};
 
 pub struct Thread {
     call_stack: CallStack,
@@ -26,13 +26,13 @@ impl Thread {
     }
 
     pub fn run(&mut self) -> Option<Value> {
-        let mut current_call = FunctionCall {
-            name: None,
-            record_index: 0,
-            return_register: 0,
-            ip: 0,
-        };
         let mut active = &mut self.records[0];
+        let mut current_call = FunctionCall {
+            name: active.name().cloned(),
+            record_index: active.index(),
+            return_register: active.stack_size() as u8 - 1,
+            ip: active.ip,
+        };
 
         info!(
             "Starting thread with {}",
@@ -43,15 +43,6 @@ impl Thread {
         );
 
         loop {
-            assert!(
-                active.ip < active.actions.len(),
-                "{}",
-                VmError::InstructionIndexOutOfBounds {
-                    call_stack: self.call_stack.clone(),
-                    ip: active.ip,
-                }
-            );
-
             trace!(
                 "Run \"{}\" | Record = {} | IP = {}",
                 active
@@ -62,6 +53,10 @@ impl Thread {
                 active.ip
             );
 
+            if active.ip >= active.actions.len() {
+                return None;
+            }
+
             let action = active.actions[active.ip];
             let signal = (action.logic)(action.data, active);
 
@@ -96,16 +91,24 @@ impl Thread {
                         active.ip = 0;
                     }
 
-                    self.call_stack.push(current_call);
+                    let next_call = FunctionCall {
-                    active.reserve_registers(arguments.len());
-
-                    current_call = FunctionCall {
                         name: active.name().cloned(),
                         record_index: active.index(),
                         return_register,
                         ip: active.ip,
                     };
 
+                    if self
+                        .call_stack
+                        .last()
+                        .is_some_and(|call| call != &next_call)
+                        || self.call_stack.is_empty()
+                    {
+                        self.call_stack.push(current_call);
+                    }
+
+                    current_call = next_call;
+
                     active = &mut self.records[record_index];
 
                     for (index, argument) in arguments.into_iter().enumerate() {
@@ -128,34 +131,50 @@ impl Thread {
 
                     active.set_register(to_register_index, register);
                 }
-                ThreadSignal::Return(should_return_value) => {
+                ThreadSignal::Return {
-                    trace!("{:#?}", self.call_stack);
+                    should_return_value,
+                    return_register,
+                } => {
+                    trace!("\n{:#?}{}", self.call_stack, current_call);
 
-                    if self.call_stack.is_empty() {
+                    let outer_call = if let Some(call) = self.call_stack.pop() {
-                        return None;
+                        call
-                    }
+                    } else if should_return_value {
-
-                    let outer_call = self.call_stack.pop().unwrap();
-                    let record_index = outer_call.record_index as usize;
-
-                    if should_return_value {
-                        let return_register = active
-                            .last_assigned_register()
-                            .unwrap_or_else(|| panic!("Expected return value"));
                         let return_value = active
                             .empty_register_or_clone_constant(return_register, Register::Empty);
+                        let next_call = self.call_stack.pop();
 
-                        active = &mut self.records[record_index];
+                        if next_call.is_none() {
+                            return Some(return_value);
+                        }
 
-                        active.reserve_registers((current_call.return_register + 1) as usize);
+                        let next_index = active.index() as usize - 1;
-                        active.set_register(
+                        let next_record = &mut self.records[next_index];
+
+                        next_record.set_register(
                             current_call.return_register,
                             Register::Value(return_value),
                         );
+
+                        current_call = next_call.unwrap();
+                        active = next_record;
+
+                        continue;
                     } else {
-                        active = &mut self.records[record_index];
+                        return None;
+                    };
+                    let record_index = outer_call.record_index as usize;
+                    let destination = current_call.return_register;
+
+                    if should_return_value {
+                        let return_value = active
+                            .empty_register_or_clone_constant(return_register, Register::Empty);
+                        let outer_record = &mut self.records[record_index];
+
+                        outer_record.set_register(destination, Register::Value(return_value));
                     }
 
+                    active = &mut self.records[record_index];
                     current_call = outer_call;
                 }
             }
@@ -171,7 +190,10 @@ pub enum ThreadSignal {
         return_register: u8,
         argument_count: u8,
     },
-    Return(bool),
+    Return {
+        should_return_value: bool,
+        return_register: u8,
+    },
     LoadFunction {
         from_record_index: u8,
         to_register_index: u8,

examples/fibonacci.ds

@@ -5,4 +5,4 @@ fn fib (n: int) -> int {
     fib(n - 1) + fib(n - 2)
 }
 
-write_line(fib(2))
+write_line(fib(10))