Begin large refactor to enhance type handling

dust-lang/src/chunk.rs

@@ -11,7 +11,7 @@ use std::{
 
 use serde::{Deserialize, Serialize};
 
-use crate::{Disassembler, Instruction, Operation, Span, Type, Value};
+use crate::{Disassembler, Instruction, Span, Type, Value};
 
 /// In-memory representation of a Dust program or function.
 ///
@@ -19,12 +19,12 @@ use crate::{Disassembler, Instruction, Operation, Span, Type, Value};
 #[derive(Clone, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct Chunk {
     name: Option<String>,
-    pub is_poisoned: bool,
 
     instructions: Vec<(Instruction, Span)>,
     constants: Vec<Value>,
     locals: Vec<Local>,
 
+    return_type: Type,
     current_scope: Scope,
     block_index: u8,
 }
@@ -33,10 +33,10 @@ impl Chunk {
     pub fn new(name: Option<String>) -> Self {
         Self {
             name,
-            is_poisoned: false,
             instructions: Vec::new(),
             constants: Vec::new(),
             locals: Vec::new(),
+            return_type: Type::None,
             current_scope: Scope::default(),
             block_index: 0,
         }
@@ -50,10 +50,10 @@ impl Chunk {
     ) -> Self {
         Self {
             name,
-            is_poisoned: false,
             instructions,
             constants,
             locals,
+            return_type: Type::None,
             current_scope: Scope::default(),
             block_index: 0,
         }
@@ -175,96 +175,26 @@ impl Chunk {
         }
     }
 
-    pub fn expect_not_poisoned(&self) -> Result<(), ChunkError> {
-        if self.is_poisoned {
-            Err(ChunkError::PoisonedChunk)
-        } else {
-            Ok(())
-        }
-    }
-
-    pub fn get_constant_type(&self, constant_index: u8) -> Option<Type> {
+    pub fn get_constant_type(&self, constant_index: u8) -> Result<Type, ChunkError> {
         self.constants
             .get(constant_index as usize)
             .map(|value| value.r#type())
-    }
-
-    pub fn get_local_type(&self, local_index: u8) -> Option<Type> {
-        self.locals.get(local_index as usize)?.r#type.clone()
-    }
-
-    pub fn get_register_type(&self, register_index: u8) -> Option<Type> {
-        let local_type_option = self
-            .locals
-            .iter()
-            .find(|local| local.register_index == register_index)
-            .map(|local| local.r#type.clone());
-
-        if let Some(local_type) = local_type_option {
-            return local_type;
-        }
-
-        self.instructions
-            .iter()
-            .enumerate()
-            .find_map(|(index, (instruction, _))| {
-                if let Operation::LoadList = instruction.operation() {
-                    if instruction.a() == register_index {
-                        let mut length = (instruction.c() - instruction.b() + 1) as usize;
-                        let mut item_type = Type::Any;
-                        let distance_to_end = self.len() - index;
-
-                        for (instruction, _) in self
-                            .instructions()
-                            .iter()
-                            .rev()
-                            .skip(distance_to_end)
-                            .take(length)
-                        {
-                            if let Operation::Close = instruction.operation() {
-                                length -= (instruction.c() - instruction.b()) as usize;
-                            } else if let Type::Any = item_type {
-                                item_type = instruction.yielded_type(self).unwrap_or(Type::Any);
-                            }
-                        }
-
-                        return Some(Type::List {
-                            item_type: Box::new(item_type),
-                            length,
-                        });
-                    }
-                }
-
-                if instruction.yields_value() && instruction.a() == register_index {
-                    instruction.yielded_type(self)
-                } else {
-                    None
-                }
+            .ok_or(ChunkError::ConstantIndexOutOfBounds {
+                index: constant_index as usize,
             })
     }
 
-    pub fn return_type(&self) -> Option<Type> {
-        let returns_value = self
-            .instructions()
-            .last()
-            .map(|(instruction, _)| {
-                debug_assert!(matches!(instruction.operation(), Operation::Return));
-
-                instruction.b_as_boolean()
+    pub fn get_local_type(&self, local_index: u8) -> Result<&Type, ChunkError> {
+        self.locals
+            .get(local_index as usize)
+            .map(|local| &local.r#type)
+            .ok_or(ChunkError::LocalIndexOutOfBounds {
+                index: local_index as usize,
             })
-            .unwrap_or(false);
+    }
 
-        if returns_value {
-            self.instructions.iter().rev().find_map(|(instruction, _)| {
-                if instruction.yields_value() {
-                    instruction.yielded_type(self)
-                } else {
-                    None
-                }
-            })
-        } else {
-            None
-        }
+    pub fn return_type(&self) -> &Type {
+        &self.return_type
     }
 
     pub fn disassembler(&self) -> Disassembler {
@@ -309,33 +239,23 @@ pub struct Local {
     pub identifier_index: u8,
 
     /// The expected type of the local's value.
-    pub r#type: Option<Type>,
+    pub r#type: Type,
 
     /// Whether the local is mutable.
     pub is_mutable: bool,
 
     /// Scope where the variable was declared.
     pub scope: Scope,
-
-    /// Expected location of a local's value.
-    pub register_index: u8,
 }
 
 impl Local {
     /// Creates a new Local instance.
-    pub fn new(
-        identifier_index: u8,
-        r#type: Option<Type>,
-        mutable: bool,
-        scope: Scope,
-        register_index: u8,
-    ) -> Self {
+    pub fn new(identifier_index: u8, r#type: Type, mutable: bool, scope: Scope) -> Self {
         Self {
             identifier_index,
             r#type,
             is_mutable: mutable,
             scope,
-            register_index,
         }
     }
 }
@@ -381,7 +301,6 @@ pub enum ChunkError {
     ConstantIndexOutOfBounds { index: usize },
     InstructionIndexOutOfBounds { index: usize },
     LocalIndexOutOfBounds { index: usize },
-    PoisonedChunk,
 }
 
 impl Display for ChunkError {
@@ -396,7 +315,6 @@ impl Display for ChunkError {
             ChunkError::LocalIndexOutOfBounds { index } => {
                 write!(f, "Local index {} out of bounds", index)
             }
-            ChunkError::PoisonedChunk => write!(f, "Chunk is poisoned"),
         }
     }
 }

dust-lang/src/compiler.rs

@@ -13,9 +13,9 @@ use std::{
 use colored::Colorize;
 
 use crate::{
-    optimize, AnnotatedError, Chunk, ChunkError, DustError, FunctionType, Instruction, LexError,
-    Lexer, Local, NativeFunction, Operation, Scope, Span, Token, TokenKind, TokenOwned, Type,
-    Value,
+    AnnotatedError, Chunk, ChunkError, DustError, FunctionType, Instruction, LexError, Lexer,
+    Local, NativeFunction, Operation, Optimizer, Scope, Span, Token, TokenKind, TokenOwned, Type,
+    TypeConflict, Value,
 };
 
 /// Compiles the input and returns a chunk.
@@ -48,9 +48,10 @@ pub fn compile(source: &str) -> Result<Chunk, DustError> {
 pub struct Compiler<'src> {
     chunk: Chunk,
     lexer: Lexer<'src>,
-    optimization_count: usize,
 
-    previous_is_expression: bool,
+    local_definitions: Vec<u8>,
+    optimization_count: usize,
+    previous_expression_type: Type,
     minimum_register: u8,
 
     current_token: Token<'src>,
@@ -74,8 +75,9 @@ impl<'src> Compiler<'src> {
         Ok(Compiler {
             chunk,
             lexer,
+            local_definitions: Vec::new(),
             optimization_count: 0,
-            previous_is_expression: false,
+            previous_expression_type: Type::None,
             minimum_register: 0,
             current_token,
             current_position,
@@ -165,10 +167,10 @@ impl<'src> Compiler<'src> {
             })
     }
 
-    pub fn declare_local(
+    fn declare_local(
         &mut self,
         identifier: &str,
-        r#type: Option<Type>,
+        r#type: Type,
         is_mutable: bool,
         scope: Scope,
         register_index: u8,
@@ -178,13 +180,10 @@ impl<'src> Compiler<'src> {
         let identifier = Value::string(identifier);
         let identifier_index = self.chunk.push_or_get_constant(identifier);
 
-        self.chunk.locals_mut().push(Local::new(
-            identifier_index,
-            r#type,
-            is_mutable,
-            scope,
-            register_index,
-        ));
+        self.chunk
+            .locals_mut()
+            .push(Local::new(identifier_index, r#type, is_mutable, scope));
+        self.local_definitions.push(register_index);
 
         (self.chunk.locals().len() as u8 - 1, identifier_index)
     }
@@ -231,35 +230,27 @@ impl<'src> Compiler<'src> {
             })
     }
 
-    fn get_last_value_operation(&self) -> Option<Operation> {
+    fn get_last_operations<const COUNT: usize>(&self) -> Option<[Operation; COUNT]> {
+        let mut n_operations = [Operation::Return; COUNT];
+
+        for (nth, operation) in n_operations.iter_mut().rev().zip(
             self.chunk
-            .instructions()
-            .iter()
-            .last()
-            .map(|(instruction, _)| instruction.operation())
-    }
-
-    fn get_last_instructions<const COUNT: usize>(&self) -> Option<[Operation; COUNT]> {
-        let mut operations = [Operation::Return; COUNT];
-
-        for (index, (instruction, _)) in self
-            .chunk
                 .instructions()
                 .iter()
                 .rev()
-            .take(COUNT)
-            .enumerate()
-        {
-            operations[index] = instruction.operation();
+                .map(|(instruction, _)| instruction.operation()),
+        ) {
+            *nth = operation;
         }
 
-        Some(operations)
+        Some(n_operations)
     }
 
     fn get_last_jumpable_mut(&mut self) -> Option<&mut Instruction> {
         self.chunk
             .instructions_mut()
             .iter_mut()
+            .rev()
             .find_map(|(instruction, _)| {
                 if let Operation::LoadBoolean | Operation::LoadConstant = instruction.operation() {
                     Some(instruction)
@@ -269,6 +260,101 @@ impl<'src> Compiler<'src> {
             })
     }
 
+    pub fn get_instruction_type(&self, instruction: &Instruction) -> Result<Type, CompileError> {
+        use Operation::*;
+
+        match instruction.operation() {
+            Add | Divide | Modulo | Multiply | Subtract => {
+                if instruction.b_is_constant() {
+                    self.chunk
+                        .get_constant_type(instruction.b())
+                        .map_err(|error| CompileError::Chunk {
+                            error,
+                            position: self.current_position,
+                        })
+                } else {
+                    self.get_register_type(instruction.b())
+                }
+            }
+            LoadBoolean | Not => Ok(Type::Boolean),
+            Negate => {
+                if instruction.b_is_constant() {
+                    self.chunk
+                        .get_constant_type(instruction.b())
+                        .map_err(|error| CompileError::Chunk {
+                            error,
+                            position: self.current_position,
+                        })
+                } else {
+                    self.get_register_type(instruction.b())
+                }
+            }
+            LoadConstant => self
+                .chunk
+                .get_constant_type(instruction.b())
+                .map_err(|error| CompileError::Chunk {
+                    error,
+                    position: self.current_position,
+                }),
+            LoadList => self.get_register_type(instruction.a()),
+            GetLocal => self
+                .chunk
+                .get_local_type(instruction.b())
+                .cloned()
+                .map_err(|error| CompileError::Chunk {
+                    error,
+                    position: self.current_position,
+                }),
+            CallNative => {
+                let native_function = NativeFunction::from(instruction.b());
+
+                Ok(*native_function.r#type().return_type)
+            }
+            _ => Ok(Type::None),
+        }
+    }
+
+    pub fn get_register_type(&self, register_index: u8) -> Result<Type, CompileError> {
+        for (index, (instruction, _)) in self.chunk.instructions().iter().enumerate() {
+            if let Operation::LoadList = instruction.operation() {
+                if instruction.a() == register_index {
+                    let mut length = (instruction.c() - instruction.b() + 1) as usize;
+                    let mut item_type = Type::Any;
+                    let distance_to_end = self.chunk.len() - index;
+
+                    for (instruction, _) in self
+                        .chunk
+                        .instructions()
+                        .iter()
+                        .rev()
+                        .skip(distance_to_end)
+                        .take(length)
+                    {
+                        if let Operation::Close = instruction.operation() {
+                            length -= (instruction.c() - instruction.b()) as usize;
+                        } else if let Type::Any = item_type {
+                            item_type = self.get_instruction_type(instruction)?;
+                        }
+                    }
+
+                    return Ok(Type::List {
+                        item_type: Box::new(item_type),
+                        length,
+                    });
+                }
+            }
+
+            if instruction.yields_value() && instruction.a() == register_index {
+                return self.get_instruction_type(instruction);
+            }
+        }
+
+        Err(CompileError::CannotResolveRegisterType {
+            register_index: register_index as usize,
+            position: self.current_position,
+        })
+    }
+
     fn emit_constant(&mut self, value: Value, position: Span) -> Result<(), CompileError> {
         let constant_index = self.chunk.push_or_get_constant(value);
         let register = self.next_register();
@@ -295,7 +381,7 @@ impl<'src> Compiler<'src> {
                 position,
             );
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Boolean;
 
             Ok(())
         } else {
@@ -319,7 +405,7 @@ impl<'src> Compiler<'src> {
 
             self.emit_constant(value, position)?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Byte;
 
             Ok(())
         } else {
@@ -341,7 +427,7 @@ impl<'src> Compiler<'src> {
 
             self.emit_constant(value, position)?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Character;
 
             Ok(())
         } else {
@@ -369,7 +455,7 @@ impl<'src> Compiler<'src> {
 
             self.emit_constant(value, position)?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Float;
 
             Ok(())
         } else {
@@ -397,7 +483,7 @@ impl<'src> Compiler<'src> {
 
             self.emit_constant(value, position)?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Integer;
 
             Ok(())
         } else {
@@ -419,7 +505,9 @@ impl<'src> Compiler<'src> {
 
             self.emit_constant(value, position)?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::String {
+                length: Some(text.len()),
+            };
 
             Ok(())
         } else {
@@ -436,8 +524,6 @@ impl<'src> Compiler<'src> {
         self.parse_expression()?;
         self.expect(Token::RightParenthesis)?;
 
-        self.previous_is_expression = true;
-
         Ok(())
     }
 
@@ -487,7 +573,9 @@ impl<'src> Compiler<'src> {
 
         self.emit_instruction(instruction, operator_position);
 
-        self.previous_is_expression = true;
+        if let TokenKind::Bang = operator.kind() {
+            self.previous_expression_type = Type::Boolean;
+        }
 
         Ok(())
     }
@@ -505,7 +593,22 @@ impl<'src> Compiler<'src> {
                 let local = self.get_local(local_index)?;
                 is_mutable_local = local.is_mutable;
 
-                local.register_index
+                *self
+                    .local_definitions
+                    .get(local_index as usize)
+                    .ok_or_else(|| {
+                        let identifier = self
+                            .chunk
+                            .constants()
+                            .get(local.identifier_index as usize)
+                            .unwrap()
+                            .to_string();
+
+                        CompileError::UndeclaredVariable {
+                            identifier,
+                            position: self.current_position,
+                        }
+                    })?
             }
             Operation::LoadConstant => {
                 is_constant = true;
@@ -626,17 +729,17 @@ impl<'src> Compiler<'src> {
         | Token::SlashEqual
         | Token::PercentEqual = operator
         {
-            self.previous_is_expression = false;
+            self.previous_expression_type = Type::None;
         } else {
-            self.previous_is_expression = true;
+            self.previous_expression_type = self.get_instruction_type(&left_instruction)?;
         }
 
         Ok(())
     }
 
     fn parse_comparison_binary(&mut self) -> Result<(), CompileError> {
-        if let Some(Operation::Equal | Operation::Less | Operation::LessEqual) =
-            self.get_last_value_operation()
+        if let Some([Operation::Equal | Operation::Less | Operation::LessEqual, _, _, _]) =
+            self.get_last_operations()
         {
             return Err(CompileError::CannotChainComparison {
                 position: self.current_position,
@@ -712,7 +815,6 @@ impl<'src> Compiler<'src> {
         let register = self.next_register();
 
         self.emit_instruction(instruction, operator_position);
-
         self.emit_instruction(Instruction::jump(1, true), operator_position);
         self.emit_instruction(
             Instruction::load_boolean(register, true, true),
@@ -723,7 +825,7 @@ impl<'src> Compiler<'src> {
             operator_position,
         );
 
-        self.previous_is_expression = true;
+        self.previous_expression_type = Type::Boolean;
 
         Ok(())
     }
@@ -756,7 +858,7 @@ impl<'src> Compiler<'src> {
         self.emit_instruction(Instruction::jump(jump_distance, true), operator_position);
         self.parse_sub_expression(&rule.precedence)?;
 
-        self.previous_is_expression = true;
+        self.previous_expression_type = Type::Boolean;
 
         Ok(())
     }
@@ -783,9 +885,9 @@ impl<'src> Compiler<'src> {
             let scope = self.chunk.current_scope();
 
             self.emit_instruction(Instruction::load_self(register), start_position);
-            self.declare_local(identifier, None, false, scope, register);
+            self.declare_local(identifier, Type::SelfChunk, false, scope, register);
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::SelfChunk;
 
             return Ok(());
         } else {
@@ -821,35 +923,25 @@ impl<'src> Compiler<'src> {
 
             self.parse_expression()?;
 
-            let (mut previous_instruction, previous_position) =
-                self.chunk.instructions_mut().pop().ok_or_else(|| {
-                    CompileError::ExpectedExpression {
-                        found: self.previous_token.to_owned(),
-                        position: self.previous_position,
-                    }
-                })?;
+            let register = self.next_register() - 1;
 
-            if previous_instruction.operation().is_math() {
-                let register_index = self.get_local(local_index)?.register_index;
-
-                log::trace!("Condensing SET_LOCAL to binary math expression");
-
-                previous_instruction.set_a(register_index);
-                self.emit_instruction(previous_instruction, self.current_position);
-
-                return Ok(());
-            }
-
-            let register = self.next_register();
-
-            self.emit_instruction(previous_instruction, previous_position);
             self.emit_instruction(
                 Instruction::set_local(register, local_index),
                 start_position,
             );
 
-            self.previous_is_expression = false;
-        } else {
+            self.previous_expression_type = Type::None;
+
+            let mut optimizer = Optimizer::new(self.chunk.instructions_mut());
+            let optimized = Optimizer::optimize_set_local(&mut optimizer);
+
+            if optimized {
+                self.optimization_count += 1;
+            }
+
+            return Ok(());
+        }
+
         let register = self.next_register();
 
         self.emit_instruction(
@@ -857,8 +949,9 @@ impl<'src> Compiler<'src> {
             self.previous_position,
         );
 
-            self.previous_is_expression = true;
-        }
+        let local = self.get_local(local_index)?;
+
+        self.previous_expression_type = local.r#type.clone();
 
         Ok(())
     }
@@ -886,7 +979,7 @@ impl<'src> Compiler<'src> {
         self.advance()?;
         self.chunk.begin_scope();
 
-        while !self.allow(Token::RightCurlyBrace)? && !self.is_eof() {
+        while !self.allow(Token::RightBrace)? && !self.is_eof() {
             self.parse(Precedence::None)?;
         }
 
@@ -901,12 +994,23 @@ impl<'src> Compiler<'src> {
         self.advance()?;
 
         let start_register = self.next_register();
+        let mut item_type = Type::Any;
 
-        while !self.allow(Token::RightSquareBrace)? && !self.is_eof() {
+        while !self.allow(Token::RightBracket)? && !self.is_eof() {
             let expected_register = self.next_register();
 
             self.parse_expression()?;
 
+            if expected_register > start_register {
+                if let Err(conflict) = item_type.check(&self.previous_expression_type) {
+                    return Err(CompileError::ListItemTypeConflict {
+                        conflict,
+                        position: self.previous_position,
+                    });
+                }
+            }
+
+            item_type = self.previous_expression_type.clone();
             let actual_register = self.next_register() - 1;
 
             if expected_register < actual_register {
@@ -927,7 +1031,10 @@ impl<'src> Compiler<'src> {
             Span(start, end),
         );
 
-        self.previous_is_expression = true;
+        self.previous_expression_type = Type::List {
+            item_type: Box::new(item_type),
+            length: (to_register - start_register) as usize,
+        };
 
         Ok(())
     }
@@ -937,12 +1044,12 @@ impl<'src> Compiler<'src> {
         self.parse_expression()?;
 
         if matches!(
-            self.get_last_instructions(),
+            self.get_last_operations(),
             Some([
-                Operation::LoadBoolean,
-                Operation::LoadBoolean,
+                Operation::Equal | Operation::Less | Operation::LessEqual,
                 Operation::Jump,
-                Operation::Equal | Operation::Less | Operation::LessEqual
+                Operation::LoadBoolean,
+                Operation::LoadBoolean,
             ])
         ) {
             self.chunk.instructions_mut().pop();
@@ -953,11 +1060,11 @@ impl<'src> Compiler<'src> {
         let if_block_start = self.chunk.len();
         let if_block_start_position = self.current_position;
 
-        if let Token::LeftCurlyBrace = self.current_token {
+        if let Token::LeftBrace = self.current_token {
             self.parse_block()?;
         } else {
             return Err(CompileError::ExpectedToken {
-                expected: TokenKind::LeftCurlyBrace,
+                expected: TokenKind::LeftBrace,
                 found: self.current_token.to_owned(),
                 position: self.current_position,
             });
@@ -965,47 +1072,44 @@ impl<'src> Compiler<'src> {
 
         let if_block_end = self.chunk.len();
         let mut if_block_distance = (if_block_end - if_block_start) as u8;
-
-        let if_block_is_expression = self
-            .chunk
-            .instructions()
-            .iter()
-            .rev()
-            .find_map(|(instruction, _)| {
-                if !matches!(instruction.operation(), Operation::Jump | Operation::Move) {
-                    Some(instruction.yields_value())
-                } else {
-                    None
-                }
-            })
-            .unwrap_or(false);
-
+        let if_block_type = self.previous_expression_type.clone();
         let if_last_register = self.next_register().saturating_sub(1);
 
-        if let Token::Else = self.current_token {
+        let has_else_statement = if let Token::Else = self.current_token {
             self.advance()?;
 
-            if let Token::LeftCurlyBrace = self.current_token {
+            if let Token::LeftBrace = self.current_token {
                 self.parse_block()?;
             } else {
                 return Err(CompileError::ExpectedTokenMultiple {
-                    expected: &[TokenKind::If, TokenKind::LeftCurlyBrace],
+                    expected: &[TokenKind::If, TokenKind::LeftBrace],
                     found: self.current_token.to_owned(),
                     position: self.current_position,
                 });
             }
-        } else {
-            self.previous_is_expression = false;
-        }
 
-        self.previous_is_expression = if_block_is_expression && self.previous_is_expression;
+            true
+        } else if self.previous_expression_type != Type::None {
+            return Err(CompileError::IfMissingElse {
+                position: Span(if_block_start_position.0, self.current_position.1),
+            });
+        } else {
+            false
+        };
 
         let else_block_end = self.chunk.len();
         let else_block_distance = (else_block_end - if_block_end) as u8;
 
+        if let Err(conflict) = if_block_type.check(&self.previous_expression_type) {
+            return Err(CompileError::IfElseBranchMismatch {
+                conflict,
+                position: Span(if_block_start_position.0, self.current_position.1),
+            });
+        }
+
         match else_block_distance {
             0 => {}
-            1 => {
+            1 if !has_else_statement => {
                 if let Some(skippable) = self.get_last_jumpable_mut() {
                     skippable.set_c_to_boolean(true);
                 } else {
@@ -1020,6 +1124,7 @@ impl<'src> Compiler<'src> {
                     );
                 }
             }
+            1 => {}
             2.. => {
                 if_block_distance += 1;
 
@@ -1042,13 +1147,12 @@ impl<'src> Compiler<'src> {
         );
 
         if self.chunk.len() >= 4 {
-            let possible_comparison_statement = {
-                let start = self.chunk.len() - 4;
+            let mut optimizer = Optimizer::new(self.chunk.instructions_mut());
+            let optimized = optimizer.optimize_comparison();
 
-                &mut self.chunk.instructions_mut()[start..]
-            };
-
-            self.optimization_count += optimize(possible_comparison_statement);
+            if optimized {
+                self.optimization_count += 1
+            }
         }
 
         let else_last_register = self.next_register().saturating_sub(1);
@@ -1071,12 +1175,12 @@ impl<'src> Compiler<'src> {
         self.parse_expression()?;
 
         if matches!(
-            self.get_last_instructions(),
+            self.get_last_operations(),
             Some([
-                Operation::LoadBoolean,
-                Operation::LoadBoolean,
+                Operation::Equal | Operation::Less | Operation::LessEqual,
                 Operation::Jump,
-                Operation::Equal | Operation::Less | Operation::LessEqual
+                Operation::LoadBoolean,
+                Operation::LoadBoolean,
             ],)
         ) {
             self.chunk.instructions_mut().pop();
@@ -1103,7 +1207,7 @@ impl<'src> Compiler<'src> {
 
         self.emit_instruction(jump_back, self.current_position);
 
-        self.previous_is_expression = false;
+        self.previous_expression_type = Type::None;
 
         Ok(())
     }
@@ -1134,7 +1238,8 @@ impl<'src> Compiler<'src> {
         let end = self.previous_position.1;
         let to_register = self.next_register();
         let argument_count = to_register - start_register;
-        self.previous_is_expression = function.r#type().return_type.is_some();
+
+        self.previous_expression_type = Type::Function(function.r#type());
 
         self.emit_instruction(
             Instruction::call_native(to_register, function, argument_count),
@@ -1147,7 +1252,7 @@ impl<'src> Compiler<'src> {
         loop {
             self.parse(Precedence::None)?;
 
-            if self.is_eof() || self.allow(Token::RightCurlyBrace)? {
+            if self.is_eof() || self.allow(Token::RightBrace)? {
                 self.parse_implicit_return()?;
 
                 break;
@@ -1160,7 +1265,7 @@ impl<'src> Compiler<'src> {
     fn parse_expression(&mut self) -> Result<(), CompileError> {
         self.parse(Precedence::None)?;
 
-        if !self.previous_is_expression || self.chunk.is_empty() {
+        if self.previous_expression_type == Type::None || self.chunk.is_empty() {
             return Err(CompileError::ExpectedExpression {
                 found: self.previous_token.to_owned(),
                 position: self.current_position,
@@ -1179,10 +1284,8 @@ impl<'src> Compiler<'src> {
 
         self.advance()?;
 
-        let has_return_value = if matches!(
-            self.current_token,
-            Token::Semicolon | Token::RightCurlyBrace
-        ) {
+        let has_return_value = if matches!(self.current_token, Token::Semicolon | Token::RightBrace)
+        {
             false
         } else {
             self.parse_expression()?;
@@ -1193,7 +1296,7 @@ impl<'src> Compiler<'src> {
 
         self.emit_instruction(Instruction::r#return(has_return_value), Span(start, end));
 
-        self.previous_is_expression = false;
+        self.previous_expression_type = Type::None;
 
         Ok(())
     }
@@ -1203,7 +1306,7 @@ impl<'src> Compiler<'src> {
             self.emit_instruction(Instruction::r#return(false), self.current_position);
         } else {
             self.emit_instruction(
-                Instruction::r#return(self.previous_is_expression),
+                Instruction::r#return(self.previous_expression_type != Type::None),
                 self.current_position,
             );
         }
@@ -1228,29 +1331,33 @@ impl<'src> Compiler<'src> {
                 position,
             });
         };
-        let r#type = if self.allow(Token::Colon)? {
-            let r#type = self.parse_type_from(self.current_token, self.current_position)?;
-
+        let explicit_type = if self.allow(Token::Colon)? {
             self.advance()?;
 
+            let r#type = self.parse_type_from(self.current_token, self.current_position)?;
+
             Some(r#type)
         } else {
             None
         };
-        let register = self.next_register();
 
         self.expect(Token::Equal)?;
         self.parse_expression()?;
 
+        let register = self.next_register() - 1;
+        let r#type = if let Some(r#type) = explicit_type {
+            r#type
+        } else {
+            self.get_register_type(register)?
+        };
         let (local_index, _) = self.declare_local(identifier, r#type, is_mutable, scope, register);
-        let register = self.next_register().saturating_sub(1);
 
         self.emit_instruction(
             Instruction::define_local(register, local_index, is_mutable),
             position,
         );
 
-        self.previous_is_expression = false;
+        self.previous_expression_type = Type::None;
 
         Ok(())
     }
@@ -1296,14 +1403,12 @@ impl<'src> Compiler<'src> {
 
             function_compiler.advance()?;
 
-            let register = value_parameters
-                .as_ref()
-                .map(|values| values.len() as u8)
-                .unwrap_or(0);
+            let register = function_compiler.next_register();
+
             let scope = function_compiler.chunk.current_scope();
             let (_, identifier_index) = function_compiler.declare_local(
                 parameter,
-                Some(r#type.clone()),
+                r#type.clone(),
                 is_mutable,
                 scope,
                 register,
@@ -1330,12 +1435,12 @@ impl<'src> Compiler<'src> {
 
             function_compiler.advance()?;
 
-            Some(Box::new(r#type))
+            Box::new(r#type)
         } else {
-            None
+            Box::new(Type::None)
         };
 
-        function_compiler.expect(Token::LeftCurlyBrace)?;
+        function_compiler.expect(Token::LeftBrace)?;
         function_compiler.parse_top_level()?;
 
         self.previous_token = function_compiler.previous_token;
@@ -1358,7 +1463,7 @@ impl<'src> Compiler<'src> {
             let scope = self.chunk.current_scope();
             let (local_index, _) = self.declare_local(
                 identifier,
-                Some(Type::Function(function_type)),
+                Type::Function(function_type),
                 false,
                 scope,
                 register,
@@ -1370,11 +1475,11 @@ impl<'src> Compiler<'src> {
                 identifier_position,
             );
 
-            self.previous_is_expression = false;
+            self.previous_expression_type = Type::None;
         } else {
             self.emit_constant(function, Span(function_start, function_end))?;
 
-            self.previous_is_expression = true;
+            self.previous_expression_type = Type::Function(function_type);
         }
 
         Ok(())
@@ -1398,6 +1503,15 @@ impl<'src> Compiler<'src> {
         }
 
         let function_register = last_instruction.a();
+        let function_return_type =
+            if let Type::Function(function_type) = self.get_register_type(function_register)? {
+                *function_type.return_type
+            } else {
+                return Err(CompileError::ExpectedFunction {
+                    found: self.previous_token.to_owned(),
+                    position: self.previous_position,
+                });
+            };
         let start = self.current_position.0;
 
         self.advance()?;
@@ -1428,7 +1542,7 @@ impl<'src> Compiler<'src> {
             Span(start, end),
         );
 
-        self.previous_is_expression = true;
+        self.previous_expression_type = function_return_type;
 
         Ok(())
     }
@@ -1436,7 +1550,7 @@ impl<'src> Compiler<'src> {
     fn parse_semicolon(&mut self) -> Result<(), CompileError> {
         self.advance()?;
 
-        self.previous_is_expression = false;
+        self.previous_expression_type = Type::None;
 
         Ok(())
     }
@@ -1673,7 +1787,7 @@ impl From<&Token<'_>> for ParseRule<'_> {
                 infix: None,
                 precedence: Precedence::None,
             },
-            Token::LeftCurlyBrace => ParseRule {
+            Token::LeftBrace => ParseRule {
                 prefix: Some(Compiler::parse_block),
                 infix: None,
                 precedence: Precedence::None,
@@ -1683,7 +1797,7 @@ impl From<&Token<'_>> for ParseRule<'_> {
                 infix: Some(Compiler::parse_call),
                 precedence: Precedence::Call,
             },
-            Token::LeftSquareBrace => ParseRule {
+            Token::LeftBracket => ParseRule {
                 prefix: Some(Compiler::parse_list),
                 infix: None,
                 precedence: Precedence::None,
@@ -1745,7 +1859,7 @@ impl From<&Token<'_>> for ParseRule<'_> {
                 infix: None,
                 precedence: Precedence::None,
             },
-            Token::RightCurlyBrace => ParseRule {
+            Token::RightBrace => ParseRule {
                 prefix: None,
                 infix: None,
                 precedence: Precedence::None,
@@ -1755,7 +1869,7 @@ impl From<&Token<'_>> for ParseRule<'_> {
                 infix: None,
                 precedence: Precedence::None,
             },
-            Token::RightSquareBrace => ParseRule {
+            Token::RightBracket => ParseRule {
                 prefix: None,
                 infix: None,
                 precedence: Precedence::None,
@@ -1828,6 +1942,10 @@ pub enum CompileError {
         found: TokenOwned,
         position: Span,
     },
+    ExpectedFunction {
+        found: TokenOwned,
+        position: Span,
+    },
     InvalidAssignmentTarget {
         found: TokenOwned,
         position: Span,
@@ -1856,6 +1974,27 @@ pub enum CompileError {
         position: Span,
     },
 
+    // Type errors
+    CannotResolveRegisterType {
+        register_index: usize,
+        position: Span,
+    },
+    CannotResolveVariableType {
+        identifier: String,
+        position: Span,
+    },
+    IfElseBranchMismatch {
+        conflict: TypeConflict,
+        position: Span,
+    },
+    IfMissingElse {
+        position: Span,
+    },
+    ListItemTypeConflict {
+        conflict: TypeConflict,
+        position: Span,
+    },
+
     // Wrappers around foreign errors
     Chunk {
         error: ChunkError,
@@ -1879,15 +2018,21 @@ impl AnnotatedError for CompileError {
 
     fn description(&self) -> &'static str {
         match self {
-            Self::CannotChainComparison { .. } => "Cannot chain comparison",
+            Self::CannotChainComparison { .. } => "Cannot chain comparison operations",
             Self::CannotMutateImmutableVariable { .. } => "Cannot mutate immutable variable",
+            Self::CannotResolveRegisterType { .. } => "Cannot resolve register type",
+            Self::CannotResolveVariableType { .. } => "Cannot resolve type",
             Self::Chunk { .. } => "Chunk error",
             Self::ExpectedExpression { .. } => "Expected an expression",
+            Self::ExpectedFunction { .. } => "Expected a function",
             Self::ExpectedMutableVariable { .. } => "Expected a mutable variable",
             Self::ExpectedToken { .. } => "Expected a specific token",
             Self::ExpectedTokenMultiple { .. } => "Expected one of multiple tokens",
+            Self::IfElseBranchMismatch { .. } => "Type mismatch in if/else branches",
+            Self::IfMissingElse { .. } => "If statement missing else branch",
             Self::InvalidAssignmentTarget { .. } => "Invalid assignment target",
             Self::Lex(error) => error.description(),
+            Self::ListItemTypeConflict { .. } => "List item type conflict",
             Self::ParseFloatError { .. } => "Failed to parse float",
             Self::ParseIntError { .. } => "Failed to parse integer",
             Self::UndeclaredVariable { .. } => "Undeclared variable",
@@ -1898,14 +2043,14 @@ impl AnnotatedError for CompileError {
 
     fn details(&self) -> Option<String> {
         match self {
-            Self::CannotChainComparison { .. } => {
-                Some("Cannot chain comparison operations".to_string())
-            }
             Self::CannotMutateImmutableVariable { identifier, .. } => {
-                Some(format!("Cannot mutate immutable variable {identifier}"))
+                Some(format!("{identifier} is immutable"))
             }
             Self::Chunk { error, .. } => Some(error.to_string()),
             Self::ExpectedExpression { found, .. } => Some(format!("Found {found}")),
+            Self::ExpectedFunction { found, .. } => {
+                Some(format!("Expected \"{found}\" to be a function"))
+            }
             Self::ExpectedToken {
                 expected, found, ..
             } => Some(format!("Expected {expected} but found {found}")),
@@ -1930,23 +2075,31 @@ impl AnnotatedError for CompileError {
 
                 Some(details)
             }
-            Self::ExpectedMutableVariable { found, .. } => {
-                Some(format!("Expected mutable variable, found {found}"))
-            }
+            Self::ExpectedMutableVariable { found, .. } => Some(format!("Found {found}")),
+            Self::IfElseBranchMismatch {
+                conflict: TypeConflict { expected, actual },
+                ..
+            } => Some(
+                format!("This if block evaluates to type \"{expected}\" but the else block evaluates to \"{actual}\"")
+            ),
+            Self::IfMissingElse { .. } => Some(
+                "This \"if\" expression evaluates to a value but is missing an else block"
+                    .to_string(),
+            ),
             Self::InvalidAssignmentTarget { found, .. } => {
-                Some(format!("Invalid assignment target, found {found}"))
+                Some(format!("Cannot assign to {found}"))
             }
             Self::Lex(error) => error.details(),
-
             Self::ParseFloatError { error, .. } => Some(error.to_string()),
             Self::ParseIntError { error, .. } => Some(error.to_string()),
             Self::UndeclaredVariable { identifier, .. } => {
-                Some(format!("Undeclared variable {identifier}"))
+                Some(format!("{identifier} has not been declared"))
             }
             Self::UnexpectedReturn { .. } => None,
             Self::VariableOutOfScope { identifier, .. } => {
-                Some(format!("Variable {identifier} is out of scope"))
+                Some(format!("{identifier} is out of scope"))
             }
+            _ => None,
         }
     }
 
@@ -1954,13 +2107,19 @@ impl AnnotatedError for CompileError {
         match self {
             Self::CannotChainComparison { position } => *position,
             Self::CannotMutateImmutableVariable { position, .. } => *position,
+            Self::CannotResolveRegisterType { position, .. } => *position,
+            Self::CannotResolveVariableType { position, .. } => *position,
             Self::Chunk { position, .. } => *position,
             Self::ExpectedExpression { position, .. } => *position,
+            Self::ExpectedFunction { position, .. } => *position,
             Self::ExpectedMutableVariable { position, .. } => *position,
             Self::ExpectedToken { position, .. } => *position,
             Self::ExpectedTokenMultiple { position, .. } => *position,
+            Self::IfElseBranchMismatch { position, .. } => *position,
+            Self::IfMissingElse { position } => *position,
             Self::InvalidAssignmentTarget { position, .. } => *position,
             Self::Lex(error) => error.position(),
+            Self::ListItemTypeConflict { position, .. } => *position,
             Self::ParseFloatError { position, .. } => *position,
             Self::ParseIntError { position, .. } => *position,
             Self::UndeclaredVariable { position, .. } => *position,

dust-lang/src/disassembler.rs

@@ -64,8 +64,8 @@ const CONSTANT_HEADER: [&str; 4] = [
 const LOCAL_HEADER: [&str; 4] = [
     "Locals",
     "------",
-    " i  IDENTIFIER       TYPE       MUTABLE  SCOPE  REGISTER",
-    "--- ---------- ---------------- ------- ------- --------",
+    " i  IDENTIFIER       TYPE       MUTABLE  SCOPE ",
+    "--- ---------- ---------------- ------- -------",
 ];
 
 /// Builder that constructs a human-readable representation of a chunk.
@@ -222,7 +222,16 @@ impl<'a> Disassembler<'a> {
             .map(|identifier| identifier.to_string())
             .unwrap_or_else(|| {
                 current_exe()
-                    .map(|path| path.to_string_lossy().to_string())
+                    .map(|path| {
+                        let path_string = path.to_string_lossy();
+                        let file_name = path_string
+                            .split('/')
+                            .last()
+                            .map(|slice| slice.to_string())
+                            .unwrap_or(path_string.to_string());
+
+                        file_name
+                    })
                     .unwrap_or("Chunk Disassembly".to_string())
             });
 
@@ -245,10 +254,7 @@ impl<'a> Disassembler<'a> {
             self.chunk.len(),
             self.chunk.constants().len(),
             self.chunk.locals().len(),
-            self.chunk
-                .return_type()
-                .map(|r#type| r#type.to_string())
-                .unwrap_or("none".to_string())
+            self.chunk.return_type()
         );
 
         self.push(&info_line, true, false, true, true);
@@ -262,23 +268,10 @@ impl<'a> Disassembler<'a> {
             let bytecode = format!("{:02X}", u32::from(instruction));
             let operation = instruction.operation().to_string();
             let info = instruction.disassembly_info(self.chunk);
-            let type_display = instruction
-                .yielded_type(self.chunk)
-                .map(|r#type| {
-                    let type_string = r#type.to_string();
-
-                    if type_string.len() > 16 {
-                        format!("{type_string:.13}...")
-                    } else {
-                        type_string
-                    }
-                })
-                .unwrap_or(String::with_capacity(0));
             let position = position.to_string();
 
-            let instruction_display = format!(
-                "{index:^3} {bytecode:>8} {operation:13} {info:^20} {type_display:^16} {position:10}"
-            );
+            let instruction_display =
+                format!("{index:^3} {bytecode:>8} {operation:13} {info:^20} {position:10}");
 
             self.push_details(&instruction_display);
         }
@@ -295,7 +288,6 @@ impl<'a> Disassembler<'a> {
                 identifier_index,
                 r#type,
                 scope,
-                register_index,
                 is_mutable: mutable,
             },
         ) in self.chunk.locals().iter().enumerate()
@@ -306,20 +298,9 @@ impl<'a> Disassembler<'a> {
                 .get(*identifier_index as usize)
                 .map(|value| value.to_string())
                 .unwrap_or_else(|| "unknown".to_string());
-            let type_display = r#type
-                .as_ref()
-                .map(|r#type| {
-                    let type_string = r#type.to_string();
-
-                    if type_string.len() > 16 {
-                        format!("{type_string:.13}...")
-                    } else {
-                        type_string
-                    }
-                })
-                .unwrap_or("unknown".to_string());
+            let type_display = r#type.to_string();
             let local_display = format!(
-                "{index:^3} {identifier_display:10} {type_display:16} {mutable:7} {scope:7} {register_index:8}"
+                "{index:^3} {identifier_display:10} {type_display:16} {mutable:7} {scope:7}"
             );
 
             self.push_details(&local_display);

dust-lang/src/formatter.rs

@@ -101,13 +101,13 @@ impl<'src> Formatter<'src> {
                 String(string) => {
                     self.push_colored(string.magenta());
                 }
-                LeftCurlyBrace => {
+                LeftBrace => {
                     self.next_line.push_str(self.current_token.as_str());
                     self.commit_line(LineKind::OpenBlock);
 
                     self.indent += 1;
                 }
-                RightCurlyBrace => {
+                RightBrace => {
                     self.commit_line(LineKind::CloseBlock);
                     self.next_line.push_str(self.current_token.as_str());
 

dust-lang/src/instruction.rs

@@ -362,6 +362,18 @@ impl Instruction {
         self
     }
 
+    pub fn returns_or_panics(&self) -> bool {
+        match self.operation() {
+            Operation::Return => true,
+            Operation::CallNative => {
+                let native_function = NativeFunction::from(self.b());
+
+                matches!(native_function, NativeFunction::Panic)
+            }
+            _ => false,
+        }
+    }
+
     pub fn yields_value(&self) -> bool {
         match self.operation() {
             Operation::Add
@@ -380,43 +392,12 @@ impl Instruction {
             Operation::CallNative => {
                 let native_function = NativeFunction::from(self.b());
 
-                native_function.r#type().return_type.is_some()
+                *native_function.r#type().return_type != Type::None
             }
             _ => false,
         }
     }
 
-    pub fn yielded_type(&self, chunk: &Chunk) -> Option<Type> {
-        use Operation::*;
-
-        match self.operation() {
-            Add | Divide | Modulo | Multiply | Subtract => {
-                if self.b_is_constant() {
-                    chunk.get_constant_type(self.b())
-                } else {
-                    chunk.get_register_type(self.b())
-                }
-            }
-            LoadBoolean | Not => Some(Type::Boolean),
-            Negate => {
-                if self.b_is_constant() {
-                    chunk.get_constant_type(self.b())
-                } else {
-                    chunk.get_register_type(self.b())
-                }
-            }
-            LoadConstant => chunk.get_constant_type(self.b()),
-            LoadList => chunk.get_register_type(self.a()),
-            GetLocal => chunk.get_local_type(self.b()),
-            CallNative => {
-                let native_function = NativeFunction::from(self.b());
-
-                native_function.r#type().return_type.map(|boxed| *boxed)
-            }
-            _ => None,
-        }
-    }
-
     pub fn disassembly_info(&self, chunk: &Chunk) -> String {
         let format_arguments = || {
             let first_argument = if self.b_is_constant() {
@@ -629,7 +610,7 @@ impl Instruction {
                 let mut output = String::new();
                 let native_function_name = native_function.as_str();
 
-                if native_function.r#type().return_type.is_some() {
+                if *native_function.r#type().return_type != Type::None {
                     output.push_str(&format!("R{} = {}(", to_register, native_function_name));
                 } else {
                     output.push_str(&format!("{}(", native_function_name));

dust-lang/src/lexer.rs

@@ -41,11 +41,10 @@ pub fn lex<'tokens, 'src: 'tokens>(
             error: CompileError::Lex(error),
             source,
         })?;
-        let is_eof = matches!(token, Token::Eof);
 
         tokens.push((token, span));
 
-        if is_eof {
+        if lexer.is_eof() {
             break;
         }
     }
@@ -60,6 +59,7 @@ pub fn lex<'tokens, 'src: 'tokens>(
 pub struct Lexer<'src> {
     source: &'src str,
     position: usize,
+    is_eof: bool,
 }
 
 impl<'src> Lexer<'src> {
@@ -68,6 +68,7 @@ impl<'src> Lexer<'src> {
         Lexer {
             source,
             position: 0,
+            is_eof: false,
         }
     }
 
@@ -76,7 +77,7 @@ impl<'src> Lexer<'src> {
     }
 
     pub fn is_eof(&self) -> bool {
-        self.position >= self.source.len()
+        self.is_eof
     }
 
     pub fn skip_to(&mut self, position: usize) {
@@ -87,272 +88,19 @@ impl<'src> Lexer<'src> {
     pub fn next_token(&mut self) -> Result<(Token<'src>, Span), LexError> {
         self.skip_whitespace();
 
-        let (token, span) = if let Some(c) = self.peek_char() {
-            match c {
-                '0'..='9' => self.lex_numeric()?,
-                '-' => {
-                    let second_char = self.peek_second_char();
+        let (token, span) = if let Some(character) = self.peek_char() {
+            let lexer = LexRule::from(&character).lexer;
 
-                    if let Some('=') = second_char {
-                        self.position += 2;
-
-                        (Token::MinusEqual, Span(self.position - 2, self.position))
-                    } else if let Some('>') = second_char {
-                        self.position += 2;
-
-                        (Token::ArrowThin, Span(self.position - 2, self.position))
-                    } else if let Some('0'..='9') = second_char {
-                        self.lex_numeric()?
-                    } else if "-Infinity" == self.peek_chars(9) {
-                        self.position += 9;
-
-                        (
-                            Token::Float("-Infinity"),
-                            Span(self.position - 9, self.position),
-                        )
+            lexer(self)?
         } else {
-                        self.position += 1;
+            self.is_eof = true;
 
-                        (Token::Minus, Span(self.position - 1, self.position))
-                    }
-                }
-                'a'..='z' | 'A'..='Z' => self.lex_alphanumeric()?,
-                '"' => self.lex_string()?,
-                '\'' => {
-                    self.position += 1;
-
-                    if let Some(c) = self.peek_char() {
-                        self.position += 1;
-
-                        let peek = self.peek_char();
-
-                        if let Some('\'') = peek {
-                            self.position += 1;
-
-                            (Token::Character(c), Span(self.position - 3, self.position))
-                        } else {
-                            return Err(LexError::ExpectedCharacter {
-                                expected: '\'',
-                                actual: peek.unwrap_or('\0'),
-                                position: self.position,
-                            });
-                        }
-                    } else {
-                        return Err(LexError::UnexpectedEndOfFile {
-                            position: self.position,
-                        });
-                    }
-                }
-                '+' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::PlusEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Plus, Span(self.position - 1, self.position))
-                    }
-                }
-                '*' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::StarEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Star, Span(self.position - 1, self.position))
-                    }
-                }
-                '(' => {
-                    self.position += 1;
-
-                    (
-                        Token::LeftParenthesis,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                ')' => {
-                    self.position += 1;
-
-                    (
-                        Token::RightParenthesis,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                '=' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::DoubleEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Equal, Span(self.position - 1, self.position))
-                    }
-                }
-                '[' => {
-                    self.position += 1;
-
-                    (
-                        Token::LeftSquareBrace,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                ']' => {
-                    self.position += 1;
-
-                    (
-                        Token::RightSquareBrace,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                ',' => {
-                    self.position += 1;
-
-                    (Token::Comma, Span(self.position - 1, self.position))
-                }
-                '.' => {
-                    if let Some('.') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::DoubleDot, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Dot, Span(self.position - 1, self.position))
-                    }
-                }
-                '>' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::GreaterEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Greater, Span(self.position - 1, self.position))
-                    }
-                }
-                '<' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::LessEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Less, Span(self.position - 1, self.position))
-                    }
-                }
-                '{' => {
-                    self.position += 1;
-
-                    (
-                        Token::LeftCurlyBrace,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                '}' => {
-                    self.position += 1;
-
-                    (
-                        Token::RightCurlyBrace,
-                        Span(self.position - 1, self.position),
-                    )
-                }
-                '/' => {
-                    if let Some('=') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::SlashEqual, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        (Token::Slash, Span(self.position - 1, self.position))
-                    }
-                }
-                '%' => {
-                    self.position += 1;
-
-                    (Token::Percent, Span(self.position - 1, self.position))
-                }
-                '&' => {
-                    if let Some('&') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (
-                            Token::DoubleAmpersand,
-                            Span(self.position - 2, self.position),
-                        )
-                    } else {
-                        self.position += 1;
-
-                        return Err(LexError::UnexpectedCharacter {
-                            actual: c,
-                            position: self.position,
-                        });
-                    }
-                }
-                ';' => {
-                    self.position += 1;
-
-                    (Token::Semicolon, Span(self.position - 1, self.position))
-                }
-                '|' => {
-                    if let Some('|') = self.peek_second_char() {
-                        self.position += 2;
-
-                        (Token::DoublePipe, Span(self.position - 2, self.position))
-                    } else {
-                        self.position += 1;
-
-                        return Err(LexError::UnexpectedCharacter {
-                            actual: c,
-                            position: self.position,
-                        });
-                    }
-                }
-                '!' => {
-                    self.position += 1;
-
-                    if let Some('=') = self.peek_char() {
-                        self.position += 1;
-
-                        (Token::BangEqual, Span(self.position - 2, self.position))
-                    } else {
-                        (Token::Bang, Span(self.position - 1, self.position))
-                    }
-                }
-                ':' => {
-                    self.position += 1;
-
-                    (Token::Colon, Span(self.position - 1, self.position))
-                }
-                _ => {
-                    return Err(LexError::UnexpectedCharacter {
-                        actual: c,
-                        position: self.position,
-                    });
-                }
-            }
-        } else {
             (Token::Eof, Span(self.position, self.position))
         };
 
         Ok((token, span))
     }
 
-    /// Peek at the next token without consuming the source.
-    pub fn peek_token(&mut self) -> Result<(Token<'src>, Span), LexError> {
-        let token = self.next_token()?;
-
-        self.position -= token.0.len();
-
-        Ok(token)
-    }
-
     /// Progress to the next character.
     fn next_char(&mut self) -> Option<char> {
         if let Some(c) = self.source[self.position..].chars().next() {
@@ -507,7 +255,7 @@ impl<'src> Lexer<'src> {
     }
 
     /// Lex an identifier token.
-    fn lex_alphanumeric(&mut self) -> Result<(Token<'src>, Span), LexError> {
+    fn lex_keyword_or_identifier(&mut self) -> Result<(Token<'src>, Span), LexError> {
         let start_pos = self.position;
 
         while let Some(c) = self.peek_char() {
@@ -564,6 +312,410 @@ impl<'src> Lexer<'src> {
 
         Ok((Token::String(text), Span(start_pos, self.position)))
     }
+
+    fn lex_char(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_position = self.position;
+
+        self.next_char();
+
+        let char = self.source[self.position..].chars().next().unwrap();
+
+        self.next_char();
+
+        if self.peek_char() == Some('\'') {
+            self.next_char();
+        } else {
+            return Err(LexError::ExpectedCharacter {
+                expected: '\'',
+                actual: self.peek_char().unwrap_or('\0'),
+                position: self.position,
+            });
+        }
+
+        let end_position = self.position;
+
+        Ok((Token::Character(char), Span(start_position, end_position)))
+    }
+
+    fn lex_plus(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::PlusEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Plus, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_minus(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_position = self.position;
+
+        if self
+            .peek_second_char()
+            .is_some_and(|char| char.is_ascii_digit())
+        {
+            return self.lex_numeric();
+        }
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            return Ok((Token::MinusEqual, Span(start_position, self.position)));
+        }
+
+        if let Some('>') = self.peek_char() {
+            self.next_char();
+
+            return Ok((Token::ArrowThin, Span(start_position, self.position)));
+        }
+
+        if self.peek_chars(8) == "Infinity" {
+            self.position += 8;
+
+            return Ok((
+                Token::Float("-Infinity"),
+                Span(start_position, self.position),
+            ));
+        }
+
+        Ok((Token::Minus, Span(start_position, self.position)))
+    }
+
+    fn lex_star(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::StarEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Star, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_slash(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::SlashEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Slash, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_percent(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::PercentEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Percent, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_unexpected(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        Err(LexError::UnexpectedCharacter {
+            actual: self.peek_char().unwrap_or('\0'),
+            position: self.position,
+        })
+    }
+
+    fn lex_exclamation_mark(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::BangEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Bang, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_equal(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::DoubleEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Equal, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_less_than(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::LessEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Less, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_greater_than(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('=') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::GreaterEqual, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Greater, Span(start_pos, self.position)))
+        }
+    }
+
+    fn lex_ampersand(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        let peek_char = self.peek_char();
+
+        if let Some('&') = peek_char {
+            self.next_char();
+
+            Ok((Token::DoubleAmpersand, Span(start_pos, self.position)))
+        } else if peek_char.is_none() {
+            Err(LexError::UnexpectedEndOfFile {
+                position: self.position,
+            })
+        } else {
+            Err(LexError::ExpectedCharacter {
+                expected: '&',
+                actual: self.peek_char().unwrap(),
+                position: self.position,
+            })
+        }
+    }
+
+    fn lex_pipe(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        let peek_char = self.peek_char();
+
+        if let Some('|') = peek_char {
+            self.next_char();
+
+            Ok((Token::DoublePipe, Span(start_pos, self.position)))
+        } else if peek_char.is_none() {
+            Err(LexError::UnexpectedEndOfFile {
+                position: self.position,
+            })
+        } else {
+            Err(LexError::ExpectedCharacter {
+                expected: '&',
+                actual: self.peek_char().unwrap(),
+                position: self.position,
+            })
+        }
+    }
+
+    fn lex_left_parenthesis(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::LeftParenthesis, Span(start_pos, self.position)))
+    }
+
+    fn lex_right_parenthesis(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::RightParenthesis, Span(start_pos, self.position)))
+    }
+
+    fn lex_left_bracket(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::LeftBracket, Span(start_pos, self.position)))
+    }
+
+    fn lex_right_bracket(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::RightBracket, Span(start_pos, self.position)))
+    }
+
+    fn lex_left_brace(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::LeftBrace, Span(start_pos, self.position)))
+    }
+
+    fn lex_right_brace(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::RightBrace, Span(start_pos, self.position)))
+    }
+
+    fn lex_semicolon(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::Semicolon, Span(start_pos, self.position)))
+    }
+
+    fn lex_colon(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::Colon, Span(start_pos, self.position)))
+    }
+
+    fn lex_comma(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        Ok((Token::Comma, Span(start_pos, self.position)))
+    }
+
+    fn lex_dot(&mut self) -> Result<(Token<'src>, Span), LexError> {
+        let start_pos = self.position;
+
+        self.next_char();
+
+        if let Some('.') = self.peek_char() {
+            self.next_char();
+
+            Ok((Token::DoubleDot, Span(start_pos, self.position)))
+        } else {
+            Ok((Token::Dot, Span(start_pos, self.position)))
+        }
+    }
+}
+
+type LexerFn<'src> = fn(&mut Lexer<'src>) -> Result<(Token<'src>, Span), LexError>;
+
+pub struct LexRule<'src> {
+    lexer: LexerFn<'src>,
+}
+
+impl<'src> From<&char> for LexRule<'src> {
+    fn from(char: &char) -> Self {
+        match char {
+            '0'..='9' => LexRule {
+                lexer: Lexer::lex_numeric,
+            },
+            char if char.is_alphabetic() => LexRule {
+                lexer: Lexer::lex_keyword_or_identifier,
+            },
+            '"' => LexRule {
+                lexer: Lexer::lex_string,
+            },
+            '\'' => LexRule {
+                lexer: Lexer::lex_char,
+            },
+            '+' => LexRule {
+                lexer: Lexer::lex_plus,
+            },
+            '-' => LexRule {
+                lexer: Lexer::lex_minus,
+            },
+            '*' => LexRule {
+                lexer: Lexer::lex_star,
+            },
+            '/' => LexRule {
+                lexer: Lexer::lex_slash,
+            },
+            '%' => LexRule {
+                lexer: Lexer::lex_percent,
+            },
+            '!' => LexRule {
+                lexer: Lexer::lex_exclamation_mark,
+            },
+            '=' => LexRule {
+                lexer: Lexer::lex_equal,
+            },
+            '<' => LexRule {
+                lexer: Lexer::lex_less_than,
+            },
+            '>' => LexRule {
+                lexer: Lexer::lex_greater_than,
+            },
+            '&' => LexRule {
+                lexer: Lexer::lex_ampersand,
+            },
+            '|' => LexRule {
+                lexer: Lexer::lex_pipe,
+            },
+            '(' => LexRule {
+                lexer: Lexer::lex_left_parenthesis,
+            },
+            ')' => LexRule {
+                lexer: Lexer::lex_right_parenthesis,
+            },
+            '[' => LexRule {
+                lexer: Lexer::lex_left_bracket,
+            },
+            ']' => LexRule {
+                lexer: Lexer::lex_right_bracket,
+            },
+            '{' => LexRule {
+                lexer: Lexer::lex_left_brace,
+            },
+            '}' => LexRule {
+                lexer: Lexer::lex_right_brace,
+            },
+            ';' => LexRule {
+                lexer: Lexer::lex_semicolon,
+            },
+            ':' => LexRule {
+                lexer: Lexer::lex_colon,
+            },
+            ',' => LexRule {
+                lexer: Lexer::lex_comma,
+            },
+            '.' => LexRule {
+                lexer: Lexer::lex_dot,
+            },
+            _ => LexRule {
+                lexer: Lexer::lex_unexpected,
+            },
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Clone)]
@@ -693,7 +845,7 @@ mod tests {
             lex(input),
             Ok(vec![
                 (Token::Map, Span(0, 3)),
-                (Token::LeftCurlyBrace, Span(4, 5)),
+                (Token::LeftBrace, Span(4, 5)),
                 (Token::Identifier("x"), Span(6, 7)),
                 (Token::Equal, Span(8, 9)),
                 (Token::String("1"), Span(10, 13)),
@@ -705,7 +857,7 @@ mod tests {
                 (Token::Identifier("z"), Span(22, 23)),
                 (Token::Equal, Span(24, 25)),
                 (Token::Float("3.0"), Span(26, 29)),
-                (Token::RightCurlyBrace, Span(30, 31)),
+                (Token::RightBrace, Span(30, 31)),
                 (Token::Eof, Span(31, 31)),
             ])
         );
@@ -769,7 +921,7 @@ mod tests {
             Ok(vec![
                 (Token::Struct, Span(0, 6)),
                 (Token::Identifier("FooBar"), Span(7, 13)),
-                (Token::LeftCurlyBrace, Span(14, 15)),
+                (Token::LeftBrace, Span(14, 15)),
                 (Token::Identifier("foo"), Span(16, 19)),
                 (Token::Colon, Span(19, 20)),
                 (Token::Int, Span(21, 24)),
@@ -777,7 +929,7 @@ mod tests {
                 (Token::Identifier("bar"), Span(26, 29)),
                 (Token::Colon, Span(29, 30)),
                 (Token::FloatKeyword, Span(31, 36)),
-                (Token::RightCurlyBrace, Span(37, 38)),
+                (Token::RightBrace, Span(37, 38)),
                 (Token::Eof, Span(38, 38))
             ])
         );
@@ -790,16 +942,16 @@ mod tests {
         assert_eq!(
             lex(input),
             Ok(vec![
-                (Token::LeftSquareBrace, Span(0, 1)),
+                (Token::LeftBracket, Span(0, 1)),
                 (Token::Integer("1"), Span(1, 2)),
                 (Token::Comma, Span(2, 3)),
                 (Token::Integer("2"), Span(4, 5)),
                 (Token::Comma, Span(5, 6)),
                 (Token::Integer("3"), Span(7, 8)),
-                (Token::RightSquareBrace, Span(8, 9)),
-                (Token::LeftSquareBrace, Span(9, 10)),
+                (Token::RightBracket, Span(8, 9)),
+                (Token::LeftBracket, Span(9, 10)),
                 (Token::Integer("1"), Span(10, 11)),
-                (Token::RightSquareBrace, Span(11, 12)),
+                (Token::RightBracket, Span(11, 12)),
                 (Token::Eof, Span(12, 12)),
             ])
         )
@@ -812,13 +964,13 @@ mod tests {
         assert_eq!(
             lex(input),
             Ok(vec![
-                (Token::LeftSquareBrace, Span(0, 1)),
+                (Token::LeftBracket, Span(0, 1)),
                 (Token::Integer("1"), Span(1, 2)),
                 (Token::Comma, Span(2, 3)),
                 (Token::Integer("2"), Span(4, 5)),
                 (Token::Comma, Span(5, 6)),
                 (Token::Integer("3"), Span(7, 8)),
-                (Token::RightSquareBrace, Span(8, 9)),
+                (Token::RightBracket, Span(8, 9)),
                 (Token::Eof, Span(9, 9)),
             ])
         )
@@ -831,7 +983,7 @@ mod tests {
         assert_eq!(
             lex(input),
             Ok(vec![
-                (Token::LeftCurlyBrace, Span(0, 1)),
+                (Token::LeftBrace, Span(0, 1)),
                 (Token::Identifier("a"), Span(1, 2)),
                 (Token::Equal, Span(3, 4)),
                 (Token::Integer("1"), Span(5, 6)),
@@ -843,7 +995,7 @@ mod tests {
                 (Token::Identifier("c"), Span(15, 16)),
                 (Token::Equal, Span(17, 18)),
                 (Token::Integer("3"), Span(19, 20)),
-                (Token::RightCurlyBrace, Span(20, 21)),
+                (Token::RightBrace, Span(20, 21)),
                 (Token::Dot, Span(21, 22)),
                 (Token::Identifier("c"), Span(22, 23)),
                 (Token::Eof, Span(23, 23)),
@@ -912,15 +1064,15 @@ mod tests {
                 (Token::Identifier("x"), Span(3, 4)),
                 (Token::Less, Span(5, 6)),
                 (Token::Integer("10"), Span(7, 9)),
-                (Token::LeftCurlyBrace, Span(10, 11)),
+                (Token::LeftBrace, Span(10, 11)),
                 (Token::Identifier("x"), Span(12, 13)),
                 (Token::Plus, Span(14, 15)),
                 (Token::Integer("1"), Span(16, 17)),
-                (Token::RightCurlyBrace, Span(18, 19)),
+                (Token::RightBrace, Span(18, 19)),
                 (Token::Else, Span(20, 24)),
-                (Token::LeftCurlyBrace, Span(25, 26)),
+                (Token::LeftBrace, Span(25, 26)),
                 (Token::Identifier("x"), Span(27, 28)),
-                (Token::RightCurlyBrace, Span(29, 30)),
+                (Token::RightBrace, Span(29, 30)),
                 (Token::Eof, Span(30, 30)),
             ])
         )
@@ -937,11 +1089,11 @@ mod tests {
                 (Token::Identifier("x"), Span(6, 7)),
                 (Token::Less, Span(8, 9)),
                 (Token::Integer("10"), Span(10, 12)),
-                (Token::LeftCurlyBrace, Span(13, 14)),
+                (Token::LeftBrace, Span(13, 14)),
                 (Token::Identifier("x"), Span(15, 16)),
                 (Token::PlusEqual, Span(17, 19)),
                 (Token::Integer("1"), Span(20, 21)),
-                (Token::RightCurlyBrace, Span(22, 23)),
+                (Token::RightBrace, Span(22, 23)),
                 (Token::Eof, Span(23, 23)),
             ])
         )
@@ -984,7 +1136,7 @@ mod tests {
         assert_eq!(
             lex(input),
             Ok(vec![
-                (Token::LeftCurlyBrace, Span(0, 1)),
+                (Token::LeftBrace, Span(0, 1)),
                 (Token::Identifier("x"), Span(2, 3)),
                 (Token::Equal, Span(4, 5)),
                 (Token::Integer("42"), Span(6, 8)),
@@ -992,7 +1144,7 @@ mod tests {
                 (Token::Identifier("y"), Span(10, 11)),
                 (Token::Equal, Span(12, 13)),
                 (Token::String("foobar"), Span(14, 22)),
-                (Token::RightCurlyBrace, Span(23, 24)),
+                (Token::RightBrace, Span(23, 24)),
                 (Token::Eof, Span(24, 24)),
             ])
         )
@@ -1333,8 +1485,8 @@ mod tests {
         assert_eq!(
             lex(input),
             Ok(vec![
-                (Token::LeftSquareBrace, Span(0, 1)),
-                (Token::RightSquareBrace, Span(1, 2)),
+                (Token::LeftBracket, Span(0, 1)),
+                (Token::RightBracket, Span(1, 2)),
                 (Token::Eof, Span(2, 2)),
             ])
         )

dust-lang/src/lib.rs

@@ -24,9 +24,9 @@ pub use crate::instruction::Instruction;
 pub use crate::lexer::{lex, LexError, Lexer};
 pub use crate::native_function::{NativeFunction, NativeFunctionError};
 pub use crate::operation::Operation;
-pub use crate::optimizer::{optimize, Optimizer};
+pub use crate::optimizer::Optimizer;
 pub use crate::r#type::{EnumType, FunctionType, RangeableType, StructType, Type, TypeConflict};
-pub use crate::token::{Token, TokenKind, TokenOwned};
+pub use crate::token::{output_token_list, Token, TokenKind, TokenOwned};
 pub use crate::value::{ConcreteValue, Function, Value, ValueError};
 pub use crate::vm::{run, Vm, VmError};
 

dust-lang/src/native_function.rs

@@ -56,7 +56,7 @@ macro_rules! define_native_function {
             pub fn returns_value(&self) -> bool {
                 match self {
                     $(
-                        NativeFunction::$name => $type.return_type.is_some(),
+                        NativeFunction::$name => *$type.return_type != Type::None,
                     )*
                 }
             }
@@ -100,7 +100,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: None,
-            return_type: None
+            return_type: Box::new(Type::None)
         }
     ),
     // (AssertEqual, 1_u8, "assert_equal", false),
@@ -112,7 +112,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: None,
-            return_type: Some(Box::new(Type::Any))
+            return_type: Box::new(Type::None)
         }
     ),
 
@@ -128,7 +128,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: Some(vec![(0, Type::Any)]),
-            return_type: Some(Box::new(Type::String { length: None }))
+            return_type: Box::new(Type::String { length: None })
         }
     ),
 
@@ -188,7 +188,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: None,
-            return_type: Some(Box::new(Type::String { length: None }))
+            return_type: Box::new(Type::String { length: None })
         }
     ),
     // (ReadTo, 51_u8, "read_to", false),
@@ -203,7 +203,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: Some(vec![(0, Type::String { length: None })]),
-            return_type: None
+            return_type: Box::new(Type::None)
         }
     ),
     // (WriteFile, 56_u8, "write_file", false),
@@ -214,7 +214,7 @@ define_native_function! {
         FunctionType {
             type_parameters: None,
             value_parameters: Some(vec![(0, Type::String { length: None })]),
-            return_type: None
+            return_type: Box::new(Type::None)
         }
     )
 
@@ -324,13 +324,14 @@ impl NativeFunction {
                 };
 
                 let first_index = to_register.saturating_sub(argument_count);
-                let arguments = vm.open_nonempty_registers(first_index..to_register, position)?;
 
-                for (index, argument) in arguments.into_iter().enumerate() {
+                for (index, register) in (first_index..to_register).enumerate() {
                     if index != 0 {
                         stdout.write(b" ").map_err(map_err)?;
                     }
 
+                    let argument = vm.open_register(register, position)?;
+
                     if let Value::Concrete(ConcreteValue::String(string)) = argument {
                         let bytes = string.as_bytes();
 

dust-lang/src/optimizer.rs

@@ -1,34 +1,41 @@
-//! Tools used by the compiler to optimize a chunk's bytecode.
-use std::{iter::Map, slice::Iter};
+//! Tool used by the compiler to optimize a chunk's bytecode.
 
 use crate::{Instruction, Operation, Span};
 
-type MapToOperation = fn(&(Instruction, Span)) -> Operation;
-
-type OperationIter<'iter> = Map<Iter<'iter, (Instruction, Span)>, MapToOperation>;
-
-/// Performs optimizations on a subset of instructions.
-pub fn optimize(instructions: &mut [(Instruction, Span)]) -> usize {
-    Optimizer::new(instructions).optimize()
-}
-
 /// An instruction optimizer that mutably borrows instructions from a chunk.
 #[derive(Debug, Eq, PartialEq, PartialOrd, Ord)]
-pub struct Optimizer<'chunk> {
-    instructions: &'chunk mut [(Instruction, Span)],
+pub struct Optimizer<'a> {
+    instructions: &'a mut Vec<(Instruction, Span)>,
 }
 
-impl<'chunk> Optimizer<'chunk> {
+impl<'a> Optimizer<'a> {
     /// Creates a new optimizer with a mutable reference to some of a chunk's instructions.
-    pub fn new(instructions: &'chunk mut [(Instruction, Span)]) -> Self {
+    pub fn new(instructions: &'a mut Vec<(Instruction, Span)>) -> Self {
         Self { instructions }
     }
 
-    /// Potentially mutates the instructions to optimize them.
-    pub fn optimize(&mut self) -> usize {
-        let mut optimizations = 0;
-
-        if matches!(
+    /// Optimizes a comparison operation.
+    ///
+    /// Comparison instructions (which are always followed by a JUMP) can be optimized when the
+    /// next instructions are two constant or boolean loaders. The first loader is set to skip an
+    /// instruction if it is run while the second loader is modified to use the first's register.
+    /// This makes the following two code snippets compile to the same bytecode:
+    ///
+    /// ```dust
+    /// 4 == 4
+    /// ```
+    ///
+    /// ```dust
+    /// if 4 == 4 { true } else { false }
+    /// ```
+    ///
+    /// The instructions must be in the following order:
+    ///     - `Operation::Equal | Operation::Less | Operation::LessEqual`
+    ///     - `Operation::Jump`
+    ///     - `Operation::LoadBoolean | Operation::LoadConstant`
+    ///     - `Operation::LoadBoolean | Operation::LoadConstant`
+    pub fn optimize_comparison(&mut self) -> bool {
+        if !matches!(
             self.get_operations(),
             Some([
                 Operation::Equal | Operation::Less | Operation::LessEqual,
@@ -37,22 +44,9 @@ impl<'chunk> Optimizer<'chunk> {
                 Operation::LoadBoolean | Operation::LoadConstant,
             ])
         ) {
-            self.optimize_comparison();
-
-            optimizations += 1;
+            return false;
         }
 
-        optimizations
-    }
-
-    /// Optimizes a comparison operation.
-    ///
-    /// The instructions must be in the following order:
-    ///     - `Operation::Equal | Operation::Less | Operation::LessEqual`
-    ///     - `Operation::Jump`
-    ///     - `Operation::LoadBoolean | Operation::LoadConstant`
-    ///     - `Operation::LoadBoolean | Operation::LoadConstant`
-    fn optimize_comparison(&mut self) {
         log::debug!("Optimizing comparison");
 
         let first_loader_register = {
@@ -72,12 +66,30 @@ impl<'chunk> Optimizer<'chunk> {
         second_loader_new.set_c_to_boolean(second_loader.c_is_constant());
 
         *second_loader = second_loader_new;
+
+        true
     }
 
-    fn operations_iter(&self) -> OperationIter {
-        self.instructions
-            .iter()
-            .map(|(instruction, _)| instruction.operation())
+    pub fn optimize_set_local(&mut self) -> bool {
+        if !matches!(
+            self.get_operations(),
+            Some([
+                Operation::Add
+                    | Operation::Subtract
+                    | Operation::Multiply
+                    | Operation::Divide
+                    | Operation::Modulo,
+                Operation::SetLocal,
+            ])
+        ) {
+            return false;
+        }
+
+        log::debug!("Optimizing set local");
+
+        self.instructions.pop();
+
+        true
     }
 
     fn get_operations<const COUNT: usize>(&self) -> Option<[Operation; COUNT]> {
@@ -87,7 +99,12 @@ impl<'chunk> Optimizer<'chunk> {
 
         let mut n_operations = [Operation::Return; COUNT];
 
-        for (nth, operation) in n_operations.iter_mut().zip(self.operations_iter()) {
+        for (nth, operation) in n_operations.iter_mut().rev().zip(
+            self.instructions
+                .iter()
+                .rev()
+                .map(|(instruction, _)| instruction.operation()),
+        ) {
             *nth = operation;
         }
 

dust-lang/src/token.rs

@@ -1,8 +1,30 @@
 //! Token, TokenOwned and TokenKind types.
-use std::fmt::{self, Display, Formatter};
+use std::{
+    fmt::{self, Display, Formatter},
+    io::Write,
+};
 
 use serde::{Deserialize, Serialize};
 
+use crate::Span;
+
+pub fn output_token_list<W: Write>(tokens: &[(Token, Span)], writer: &mut W) {
+    const HEADER: [&str; 2] = [
+        "TOKEN        KIND       POSITION  ",
+        "------------ ---------- ----------",
+    ];
+
+    writeln!(writer, "{}", HEADER[0]).unwrap();
+    writeln!(writer, "{}", HEADER[1]).unwrap();
+
+    for (token, position) in tokens {
+        let kind = token.kind().to_string();
+        let token = token.to_string();
+
+        writeln!(writer, "{token:<12} {kind:<10} {position}").unwrap();
+    }
+}
+
 macro_rules! define_tokens {
     ($($variant:ident $(($data_type:ty))?),+ $(,)?) => {
         /// Source token.
@@ -72,9 +94,9 @@ define_tokens! {
     Equal,
     Greater,
     GreaterEqual,
-    LeftCurlyBrace,
+    LeftBrace,
+    LeftBracket,
     LeftParenthesis,
-    LeftSquareBrace,
     Less,
     LessEqual,
     Minus,
@@ -83,9 +105,9 @@ define_tokens! {
     PercentEqual,
     Plus,
     PlusEqual,
-    RightCurlyBrace,
+    RightBrace,
+    RightBracket,
     RightParenthesis,
-    RightSquareBrace,
     Semicolon,
     Slash,
     SlashEqual,
@@ -133,9 +155,9 @@ impl<'src> Token<'src> {
             Token::Equal => 1,
             Token::Greater => 1,
             Token::GreaterEqual => 2,
-            Token::LeftCurlyBrace => 1,
+            Token::LeftBrace => 1,
             Token::LeftParenthesis => 1,
-            Token::LeftSquareBrace => 1,
+            Token::LeftBracket => 1,
             Token::Less => 1,
             Token::LessEqual => 2,
             Token::Minus => 1,
@@ -145,9 +167,9 @@ impl<'src> Token<'src> {
             Token::Plus => 1,
             Token::PlusEqual => 2,
             Token::Return => 6,
-            Token::RightCurlyBrace => 1,
+            Token::RightBrace => 1,
             Token::RightParenthesis => 1,
-            Token::RightSquareBrace => 1,
+            Token::RightBracket => 1,
             Token::Semicolon => 1,
             Token::Slash => 1,
             Token::SlashEqual => 2,
@@ -194,9 +216,9 @@ impl<'src> Token<'src> {
             Token::Equal => "=",
             Token::Greater => ">",
             Token::GreaterEqual => ">=",
-            Token::LeftCurlyBrace => "{",
+            Token::LeftBrace => "{",
             Token::LeftParenthesis => "(",
-            Token::LeftSquareBrace => "[",
+            Token::LeftBracket => "[",
             Token::Less => "<",
             Token::LessEqual => "<=",
             Token::Minus => "-",
@@ -206,9 +228,9 @@ impl<'src> Token<'src> {
             Token::Plus => "+",
             Token::PlusEqual => "+=",
             Token::Return => "return",
-            Token::RightCurlyBrace => "}",
+            Token::RightBrace => "}",
             Token::RightParenthesis => ")",
-            Token::RightSquareBrace => "]",
+            Token::RightBracket => "]",
             Token::Semicolon => ";",
             Token::Slash => "/",
             Token::SlashEqual => "/=",
@@ -247,9 +269,9 @@ impl<'src> Token<'src> {
             Token::If => TokenOwned::If,
             Token::Int => TokenOwned::Int,
             Token::Integer(integer) => TokenOwned::Integer(integer.to_string()),
-            Token::LeftCurlyBrace => TokenOwned::LeftCurlyBrace,
+            Token::LeftBrace => TokenOwned::LeftCurlyBrace,
             Token::LeftParenthesis => TokenOwned::LeftParenthesis,
-            Token::LeftSquareBrace => TokenOwned::LeftSquareBrace,
+            Token::LeftBracket => TokenOwned::LeftSquareBrace,
             Token::Let => TokenOwned::Let,
             Token::Less => TokenOwned::Less,
             Token::LessEqual => TokenOwned::LessOrEqual,
@@ -263,9 +285,9 @@ impl<'src> Token<'src> {
             Token::Plus => TokenOwned::Plus,
             Token::PlusEqual => TokenOwned::PlusEqual,
             Token::Return => TokenOwned::Return,
-            Token::RightCurlyBrace => TokenOwned::RightCurlyBrace,
+            Token::RightBrace => TokenOwned::RightCurlyBrace,
             Token::RightParenthesis => TokenOwned::RightParenthesis,
-            Token::RightSquareBrace => TokenOwned::RightSquareBrace,
+            Token::RightBracket => TokenOwned::RightSquareBrace,
             Token::Semicolon => TokenOwned::Semicolon,
             Token::Star => TokenOwned::Star,
             Token::StarEqual => TokenOwned::StarEqual,
@@ -308,9 +330,9 @@ impl<'src> Token<'src> {
             Token::If => TokenKind::If,
             Token::Int => TokenKind::Int,
             Token::Integer(_) => TokenKind::Integer,
-            Token::LeftCurlyBrace => TokenKind::LeftCurlyBrace,
+            Token::LeftBrace => TokenKind::LeftBrace,
             Token::LeftParenthesis => TokenKind::LeftParenthesis,
-            Token::LeftSquareBrace => TokenKind::LeftSquareBrace,
+            Token::LeftBracket => TokenKind::LeftBracket,
             Token::Let => TokenKind::Let,
             Token::Less => TokenKind::Less,
             Token::LessEqual => TokenKind::LessEqual,
@@ -324,9 +346,9 @@ impl<'src> Token<'src> {
             Token::Plus => TokenKind::Plus,
             Token::PlusEqual => TokenKind::PlusEqual,
             Token::Return => TokenKind::Return,
-            Token::RightCurlyBrace => TokenKind::RightCurlyBrace,
+            Token::RightBrace => TokenKind::RightBrace,
             Token::RightParenthesis => TokenKind::RightParenthesis,
-            Token::RightSquareBrace => TokenKind::RightSquareBrace,
+            Token::RightBracket => TokenKind::RightBracket,
             Token::Semicolon => TokenKind::Semicolon,
             Token::Star => TokenKind::Star,
             Token::StarEqual => TokenKind::StarEqual,
@@ -363,9 +385,9 @@ impl<'src> Token<'src> {
                 | Token::Equal
                 | Token::Greater
                 | Token::GreaterEqual
-                | Token::LeftCurlyBrace
+                | Token::LeftBrace
                 | Token::LeftParenthesis
-                | Token::LeftSquareBrace
+                | Token::LeftBracket
                 | Token::Less
                 | Token::LessEqual
                 | Token::Minus
@@ -413,9 +435,9 @@ impl<'src> Display for Token<'src> {
             Token::If => write!(f, "if"),
             Token::Int => write!(f, "int"),
             Token::Integer(value) => write!(f, "{value}"),
-            Token::LeftCurlyBrace => write!(f, "{{"),
+            Token::LeftBrace => write!(f, "{{"),
             Token::LeftParenthesis => write!(f, "("),
-            Token::LeftSquareBrace => write!(f, "["),
+            Token::LeftBracket => write!(f, "["),
             Token::Let => write!(f, "let"),
             Token::Less => write!(f, "<"),
             Token::LessEqual => write!(f, "<="),
@@ -429,9 +451,9 @@ impl<'src> Display for Token<'src> {
             Token::Plus => write!(f, "+"),
             Token::PlusEqual => write!(f, "+="),
             Token::Return => write!(f, "return"),
-            Token::RightCurlyBrace => write!(f, "}}"),
+            Token::RightBrace => write!(f, "}}"),
             Token::RightParenthesis => write!(f, ")"),
-            Token::RightSquareBrace => write!(f, "]"),
+            Token::RightBracket => write!(f, "]"),
             Token::Semicolon => write!(f, ";"),
             Token::Slash => write!(f, "/"),
             Token::SlashEqual => write!(f, "/="),
@@ -546,9 +568,9 @@ impl Display for TokenOwned {
             TokenOwned::If => Token::If.fmt(f),
             TokenOwned::Int => Token::Int.fmt(f),
             TokenOwned::Integer(integer) => Token::Integer(integer).fmt(f),
-            TokenOwned::LeftCurlyBrace => Token::LeftCurlyBrace.fmt(f),
+            TokenOwned::LeftCurlyBrace => Token::LeftBrace.fmt(f),
             TokenOwned::LeftParenthesis => Token::LeftParenthesis.fmt(f),
-            TokenOwned::LeftSquareBrace => Token::LeftSquareBrace.fmt(f),
+            TokenOwned::LeftSquareBrace => Token::LeftBracket.fmt(f),
             TokenOwned::Let => Token::Let.fmt(f),
             TokenOwned::Less => Token::Less.fmt(f),
             TokenOwned::LessOrEqual => Token::LessEqual.fmt(f),
@@ -562,9 +584,9 @@ impl Display for TokenOwned {
             TokenOwned::Plus => Token::Plus.fmt(f),
             TokenOwned::PlusEqual => Token::PlusEqual.fmt(f),
             TokenOwned::Return => Token::Return.fmt(f),
-            TokenOwned::RightCurlyBrace => Token::RightCurlyBrace.fmt(f),
+            TokenOwned::RightCurlyBrace => Token::RightBrace.fmt(f),
             TokenOwned::RightParenthesis => Token::RightParenthesis.fmt(f),
-            TokenOwned::RightSquareBrace => Token::RightSquareBrace.fmt(f),
+            TokenOwned::RightSquareBrace => Token::RightBracket.fmt(f),
             TokenOwned::Semicolon => Token::Semicolon.fmt(f),
             TokenOwned::Star => Token::Star.fmt(f),
             TokenOwned::StarEqual => Token::StarEqual.fmt(f),
@@ -586,10 +608,10 @@ impl Display for TokenKind {
             TokenKind::Bang => Token::Bang.fmt(f),
             TokenKind::BangEqual => Token::BangEqual.fmt(f),
             TokenKind::Bool => Token::Bool.fmt(f),
-            TokenKind::Boolean => write!(f, "boolean value"),
+            TokenKind::Boolean => write!(f, "boolean"),
             TokenKind::Break => Token::Break.fmt(f),
-            TokenKind::Byte => write!(f, "byte value"),
-            TokenKind::Character => write!(f, "character value"),
+            TokenKind::Byte => write!(f, "byte"),
+            TokenKind::Character => write!(f, "character"),
             TokenKind::Colon => Token::Colon.fmt(f),
             TokenKind::Comma => Token::Comma.fmt(f),
             TokenKind::Dot => Token::Dot.fmt(f),
@@ -600,7 +622,7 @@ impl Display for TokenKind {
             TokenKind::Else => Token::Else.fmt(f),
             TokenKind::Eof => Token::Eof.fmt(f),
             TokenKind::Equal => Token::Equal.fmt(f),
-            TokenKind::Float => write!(f, "float value"),
+            TokenKind::Float => write!(f, "float"),
             TokenKind::FloatKeyword => Token::FloatKeyword.fmt(f),
             TokenKind::Fn => Token::Fn.fmt(f),
             TokenKind::Greater => Token::Greater.fmt(f),
@@ -608,10 +630,10 @@ impl Display for TokenKind {
             TokenKind::Identifier => write!(f, "identifier"),
             TokenKind::If => Token::If.fmt(f),
             TokenKind::Int => Token::Int.fmt(f),
-            TokenKind::Integer => write!(f, "integer value"),
-            TokenKind::LeftCurlyBrace => Token::LeftCurlyBrace.fmt(f),
+            TokenKind::Integer => write!(f, "integer"),
+            TokenKind::LeftBrace => Token::LeftBrace.fmt(f),
             TokenKind::LeftParenthesis => Token::LeftParenthesis.fmt(f),
-            TokenKind::LeftSquareBrace => Token::LeftSquareBrace.fmt(f),
+            TokenKind::LeftBracket => Token::LeftBracket.fmt(f),
             TokenKind::Let => Token::Let.fmt(f),
             TokenKind::Less => Token::Less.fmt(f),
             TokenKind::LessEqual => Token::LessEqual.fmt(f),
@@ -625,16 +647,16 @@ impl Display for TokenKind {
             TokenKind::Plus => Token::Plus.fmt(f),
             TokenKind::PlusEqual => Token::PlusEqual.fmt(f),
             TokenKind::Return => Token::Return.fmt(f),
-            TokenKind::RightCurlyBrace => Token::RightCurlyBrace.fmt(f),
+            TokenKind::RightBrace => Token::RightBrace.fmt(f),
             TokenKind::RightParenthesis => Token::RightParenthesis.fmt(f),
-            TokenKind::RightSquareBrace => Token::RightSquareBrace.fmt(f),
+            TokenKind::RightBracket => Token::RightBracket.fmt(f),
             TokenKind::Semicolon => Token::Semicolon.fmt(f),
             TokenKind::Star => Token::Star.fmt(f),
             TokenKind::StarEqual => Token::StarEqual.fmt(f),
             TokenKind::Str => Token::Str.fmt(f),
             TokenKind::Slash => Token::Slash.fmt(f),
             TokenKind::SlashEqual => Token::SlashEqual.fmt(f),
-            TokenKind::String => write!(f, "string value"),
+            TokenKind::String => write!(f, "string"),
             TokenKind::Struct => Token::Struct.fmt(f),
             TokenKind::While => Token::While.fmt(f),
         }

dust-lang/src/type.rs

@@ -33,10 +33,12 @@ pub enum Type {
     Map {
         pairs: HashMap<u8, Type>,
     },
+    None,
     Number,
     Range {
         r#type: RangeableType,
     },
+    SelfChunk,
     String {
         length: Option<usize>,
     },
@@ -257,8 +259,10 @@ impl Display for Type {
 
                 write!(f, "}}")
             }
+            Type::None => write!(f, "none"),
             Type::Number => write!(f, "num"),
             Type::Range { r#type } => write!(f, "{type} range"),
+            Type::SelfChunk => write!(f, "self"),
             Type::String { .. } => write!(f, "str"),
             Type::Struct(struct_type) => write!(f, "{struct_type}"),
             Type::Tuple { fields } => {
@@ -343,12 +347,16 @@ impl Ord for Type {
                 left_pairs.iter().cmp(right_pairs.iter())
             }
             (Type::Map { .. }, _) => Ordering::Greater,
+            (Type::None, Type::None) => Ordering::Equal,
+            (Type::None, _) => Ordering::Greater,
             (Type::Number, Type::Number) => Ordering::Equal,
             (Type::Number, _) => Ordering::Greater,
             (Type::Range { r#type: left_type }, Type::Range { r#type: right_type }) => {
                 left_type.cmp(right_type)
             }
             (Type::Range { .. }, _) => Ordering::Greater,
+            (Type::SelfChunk, Type::SelfChunk) => Ordering::Equal,
+            (Type::SelfChunk, _) => Ordering::Greater,
             (Type::String { length: left }, Type::String { length: right }) => left.cmp(right),
             (Type::String { .. }, _) => Ordering::Greater,
             (Type::Struct(left_struct), Type::Struct(right_struct)) => {
@@ -366,7 +374,7 @@ impl Ord for Type {
 pub struct FunctionType {
     pub type_parameters: Option<Vec<u8>>,
     pub value_parameters: Option<Vec<(u8, Type)>>,
-    pub return_type: Option<Box<Type>>,
+    pub return_type: Box<Type>,
 }
 
 impl Display for FunctionType {
@@ -401,8 +409,8 @@ impl Display for FunctionType {
 
         write!(f, ")")?;
 
-        if let Some(return_type) = &self.return_type {
-            write!(f, " -> {return_type}")?;
+        if *self.return_type != Type::None {
+            write!(f, " -> {}", self.return_type)?;
         }
 
         Ok(())

dust-lang/src/vm.rs

@@ -1,19 +1,18 @@
 //! Virtual machine and errors
 use std::{
     cmp::Ordering,
+    collections::HashMap,
     fmt::{self, Display, Formatter},
-    ops::Range,
 };
 
 use crate::{
     compile, value::ConcreteValue, AnnotatedError, Chunk, ChunkError, DustError, FunctionType,
-    Instruction, Local, NativeFunction, NativeFunctionError, Operation, Span, Type, Value,
-    ValueError,
+    Instruction, NativeFunction, NativeFunctionError, Operation, Span, Type, Value, ValueError,
 };
 
 pub fn run(source: &str) -> Result<Option<Value>, DustError> {
-    let mut chunk = compile(source)?;
-    let mut vm = Vm::new(&mut chunk, None);
+    let chunk = compile(source)?;
+    let mut vm = Vm::new(&chunk, None);
 
     vm.run()
         .map(|option| option.cloned())
@@ -34,8 +33,9 @@ pub fn run_and_display_output(source: &str) {
 #[derive(Debug, Eq, PartialEq)]
 pub struct Vm<'chunk, 'parent> {
     ip: usize,
-    chunk: &'chunk mut Chunk,
+    chunk: &'chunk Chunk,
     stack: Vec<Register>,
+    local_definitions: HashMap<u8, u8>,
     last_assigned_register: Option<u8>,
     parent: Option<&'parent Vm<'chunk, 'parent>>,
 }
@@ -43,11 +43,12 @@ pub struct Vm<'chunk, 'parent> {
 impl<'chunk, 'parent> Vm<'chunk, 'parent> {
     const STACK_LIMIT: usize = u16::MAX as usize;
 
-    pub fn new(chunk: &'chunk mut Chunk, parent: Option<&'parent Vm<'chunk, 'parent>>) -> Self {
+    pub fn new(chunk: &'chunk Chunk, parent: Option<&'parent Vm<'chunk, 'parent>>) -> Self {
         Self {
             ip: 0,
             chunk,
             stack: Vec::new(),
+            local_definitions: HashMap::new(),
             last_assigned_register: None,
             parent,
         }
@@ -162,7 +163,7 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
                         FunctionType {
                             type_parameters: None,
                             value_parameters: None,
-                            return_type: None,
+                            return_type: Box::new(Type::None),
                         },
                     );
 
@@ -172,24 +173,39 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
                     let from_register = instruction.a();
                     let to_local = instruction.b();
 
-                    self.define_local(to_local, from_register, position)?;
+                    self.define_local(to_local, from_register)?;
                 }
                 Operation::GetLocal => {
                     let to_register = instruction.a();
                     let local_index = instruction.b();
-                    let local = self.get_local(local_index, position)?;
+                    let local_register = self.local_definitions.get(&local_index).copied().ok_or(
+                        VmError::UndefinedLocal {
+                            local_index,
+                            position,
+                        },
+                    )?;
 
                     self.set_register(
                         to_register,
-                        Register::StackPointer(local.register_index),
+                        Register::StackPointer(local_register),
                         position,
                     )?;
                 }
                 Operation::SetLocal => {
-                    let register = instruction.a();
-                    let local_index = instruction.b();
+                    let from_register = instruction.a();
+                    let to_local = instruction.b();
+                    let local_register = self.local_definitions.get(&to_local).copied().ok_or(
+                        VmError::UndefinedLocal {
+                            local_index: to_local,
+                            position,
+                        },
+                    )?;
 
-                    self.define_local(local_index, register, position)?;
+                    self.set_register(
+                        local_register,
+                        Register::StackPointer(from_register),
+                        position,
+                    )?;
                 }
                 Operation::Add => {
                     let to_register = instruction.a();
@@ -401,17 +417,17 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
                     let to_register = instruction.a();
                     let function_register = instruction.b();
                     let argument_count = instruction.c();
-                    let value = self.open_register(function_register, position)?.clone();
-                    let mut function =
-                        if let Value::Concrete(ConcreteValue::Function(function)) = value {
+                    let value = self.open_register(function_register, position)?;
+                    let function = if let Value::Concrete(ConcreteValue::Function(function)) = value
+                    {
                         function
                     } else {
                         return Err(VmError::ExpectedFunction {
-                                found: value,
+                            found: value.clone(),
                             position,
                         });
                     };
-                    let mut function_vm = Vm::new(function.chunk_mut(), Some(self));
+                    let mut function_vm = Vm::new(function.chunk(), Some(self));
                     let first_argument_index = function_register + 1;
 
                     for argument_index in
@@ -527,6 +543,8 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
                     position,
                 })?;
 
+        log::trace!("Open R{register_index} to {register}");
+
         match register {
             Register::Value(value) => Ok(value),
             Register::StackPointer(register_index) => self.open_register(*register_index, position),
@@ -556,67 +574,7 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
         }
     }
 
-    pub fn open_nonempty_registers(
-        &self,
-        register_index_range: Range<u8>,
-        position: Span,
-    ) -> Result<Vec<&Value>, VmError> {
-        let mut values = Vec::with_capacity(register_index_range.len());
-
-        for register_index in register_index_range.clone() {
-            let register_index = register_index as usize;
-            let register = self.stack.get(register_index).ok_or_else(|| {
-                VmError::RegisterIndexOutOfBounds {
-                    index: register_index,
-                    position,
-                }
-            })?;
-
-            let value = match register {
-                Register::Value(value) => value,
-                Register::StackPointer(register_index) => {
-                    self.open_register(*register_index, position)?
-                }
-                Register::ConstantPointer(constant_index) => {
-                    self.get_constant(*constant_index, position)?
-                }
-                Register::ParentStackPointer(register_index) => {
-                    let parent = self
-                        .parent
-                        .as_ref()
-                        .ok_or(VmError::ExpectedParent { position })?;
-
-                    parent.open_register(*register_index, position)?
-                }
-                Register::ParentConstantPointer(constant_index) => {
-                    let parent = self
-                        .parent
-                        .as_ref()
-                        .ok_or(VmError::ExpectedParent { position })?;
-
-                    parent.get_constant(*constant_index, position)?
-                }
-                Register::Empty => continue,
-            };
-
-            values.push(value);
-        }
-
-        if values.is_empty() {
-            Err(VmError::EmptyRegisters {
-                indexes: register_index_range,
-                position,
-            })
-        } else {
-            Ok(values)
-        }
-    }
-
     fn read(&mut self, position: Span) -> Result<&(Instruction, Span), VmError> {
-        self.chunk
-            .expect_not_poisoned()
-            .map_err(|error| VmError::Chunk { error, position })?;
-
         let max_ip = self.chunk.len() - 1;
 
         if self.ip > max_ip {
@@ -628,30 +586,14 @@ impl<'chunk, 'parent> Vm<'chunk, 'parent> {
         self.get_instruction(self.ip - 1, position)
     }
 
-    fn define_local(
-        &mut self,
-        local_index: u8,
-        register_index: u8,
-        position: Span,
-    ) -> Result<(), VmError> {
-        let local = self
-            .chunk
-            .get_local_mut(local_index)
-            .map_err(|error| VmError::Chunk { error, position })?;
-
+    fn define_local(&mut self, local_index: u8, register_index: u8) -> Result<(), VmError> {
         log::debug!("Define local L{}", local_index);
 
-        local.register_index = register_index;
+        self.local_definitions.insert(local_index, register_index);
 
         Ok(())
     }
 
-    fn get_local(&self, local_index: u8, position: Span) -> Result<&Local, VmError> {
-        self.chunk
-            .get_local(local_index)
-            .map_err(|error| VmError::Chunk { error, position })
-    }
-
     fn get_instruction(
         &self,
         index: usize,
@@ -694,9 +636,11 @@ pub enum VmError {
 
     // Register errors
     EmptyRegister { index: usize, position: Span },
-    EmptyRegisters { indexes: Range<u8>, position: Span },
     RegisterIndexOutOfBounds { index: usize, position: Span },
 
+    // Local errors
+    UndefinedLocal { local_index: u8, position: Span },
+
     // Execution errors
     ExpectedBoolean { found: Value, position: Span },
     ExpectedFunction { found: Value, position: Span },
@@ -717,7 +661,6 @@ impl AnnotatedError for VmError {
         match self {
             Self::Chunk { .. } => "Chunk error",
             Self::EmptyRegister { .. } => "Empty register",
-            Self::EmptyRegisters { .. } => "Empty registers",
             Self::ExpectedBoolean { .. } => "Expected boolean",
             Self::ExpectedFunction { .. } => "Expected function",
             Self::ExpectedParent { .. } => "Expected parent",
@@ -725,6 +668,7 @@ impl AnnotatedError for VmError {
             Self::RegisterIndexOutOfBounds { .. } => "Register index out of bounds",
             Self::StackOverflow { .. } => "Stack overflow",
             Self::StackUnderflow { .. } => "Stack underflow",
+            Self::UndefinedLocal { .. } => "Undefined local",
             Self::Value { .. } => "Value error",
         }
     }
@@ -733,10 +677,6 @@ impl AnnotatedError for VmError {
         match self {
             Self::Chunk { error, .. } => Some(error.to_string()),
             Self::EmptyRegister { index, .. } => Some(format!("Register R{index} is empty")),
-            Self::EmptyRegisters { indexes: range, .. } => Some(format!(
-                "Registers R{} to R{} are empty",
-                range.start, range.end
-            )),
             Self::ExpectedFunction { found, .. } => Some(format!("{found} is not a function")),
             Self::RegisterIndexOutOfBounds { index, .. } => {
                 Some(format!("Register {index} does not exist"))
@@ -751,7 +691,6 @@ impl AnnotatedError for VmError {
         match self {
             Self::Chunk { position, .. } => *position,
             Self::EmptyRegister { position, .. } => *position,
-            Self::EmptyRegisters { position, .. } => *position,
             Self::ExpectedBoolean { position, .. } => *position,
             Self::ExpectedFunction { position, .. } => *position,
             Self::ExpectedParent { position } => *position,
@@ -759,6 +698,7 @@ impl AnnotatedError for VmError {
             Self::RegisterIndexOutOfBounds { position, .. } => *position,
             Self::StackOverflow { position } => *position,
             Self::StackUnderflow { position } => *position,
+            Self::UndefinedLocal { position, .. } => *position,
             Self::Value { position, .. } => *position,
         }
     }

dust-lang/tests/control_flow.rs

@@ -23,16 +23,7 @@ fn equality_assignment_long() {
                 (Instruction::r#return(true), Span(44, 44)),
             ],
             vec![Value::integer(4), Value::string("a")],
-            vec![Local::new(
-                1,
-                None,
-                false,
-                Scope {
-                    depth: 0,
-                    block_index: 0
-                },
-                0
-            )]
+            vec![Local::new(1, Type::Boolean, false, Scope::default(),)]
         )),
     );
 
@@ -62,7 +53,7 @@ fn equality_assignment_short() {
                 (Instruction::r#return(true), Span(16, 16)),
             ],
             vec![Value::integer(4), Value::string("a")],
-            vec![Local::new(1, None, false, Scope::default(), 0)]
+            vec![Local::new(1, Type::Boolean, false, Scope::default())]
         )),
     );
 
@@ -120,7 +111,7 @@ fn if_else_assigment_false() {
                 Value::integer(42),
                 Value::string("a")
             ],
-            vec![Local::new(5, None, false, Scope::default(), 0)]
+            vec![Local::new(5, Type::Integer, false, Scope::default())]
         )),
     );
 
@@ -178,7 +169,7 @@ fn if_else_assigment_true() {
                 Value::integer(42),
                 Value::string("a")
             ],
-            vec![Local::new(5, None, false, Scope::default(), 0)]
+            vec![Local::new(5, Type::Integer, false, Scope::default())]
         )),
     );
 
@@ -303,7 +294,7 @@ fn if_else_complex() {
 
 #[test]
 fn if_else_false() {
-    let source = "if 1 == 2 { panic() } else { 42 }";
+    let source = "if 1 == 2 { panic(); 0 } else { 42 }";
 
     assert_eq!(
         compile(source),
@@ -335,7 +326,7 @@ fn if_else_false() {
 
 #[test]
 fn if_else_true() {
-    let source = "if 1 == 1 { 42 } else { panic() }";
+    let source = "if 1 == 1 { 42 } else { panic(); 0 }";
 
     assert_eq!(
         compile(source),
@@ -367,7 +358,7 @@ fn if_else_true() {
 
 #[test]
 fn if_false() {
-    let source = "if 1 == 2 { 2 }";
+    let source = "if 1 == 2 { panic() }";
 
     assert_eq!(
         compile(source),
@@ -381,8 +372,11 @@ fn if_false() {
                     Span(5, 7)
                 ),
                 (Instruction::jump(1, true), Span(10, 11)),
-                (Instruction::load_constant(0, 1, false), Span(12, 13)),
-                (Instruction::r#return(false), Span(15, 15))
+                (
+                    Instruction::call_native(0, NativeFunction::Panic, 0),
+                    Span(12, 19)
+                ),
+                (Instruction::r#return(false), Span(21, 21))
             ],
             vec![Value::integer(1), Value::integer(2)],
             vec![]
@@ -394,7 +388,7 @@ fn if_false() {
 
 #[test]
 fn if_true() {
-    let source = "if 1 == 1 { 2 }";
+    let source = "if 1 == 1 { panic() }";
 
     assert_eq!(
         compile(source),
@@ -408,13 +402,25 @@ fn if_true() {
                     Span(5, 7)
                 ),
                 (Instruction::jump(1, true), Span(10, 11)),
-                (Instruction::load_constant(0, 1, false), Span(12, 13)),
-                (Instruction::r#return(false), Span(15, 15))
+                (
+                    Instruction::call_native(0, NativeFunction::Panic, 0),
+                    Span(12, 19)
+                ),
+                (Instruction::r#return(false), Span(21, 21))
             ],
-            vec![Value::integer(1), Value::integer(2)],
+            vec![Value::integer(1)],
             vec![]
         )),
     );
 
-    assert_eq!(run(source), Ok(None));
+    assert_eq!(
+        run(source),
+        Err(DustError::Runtime {
+            error: VmError::NativeFunction(NativeFunctionError::Panic {
+                message: None,
+                position: Span(12, 19)
+            }),
+            source
+        })
+    );
 }

dust-lang/tests/functions.rs

@@ -15,14 +15,14 @@ fn function() {
                 ],
                 vec![Value::string("a"), Value::string("b"),],
                 vec![
-                    Local::new(0, Some(Type::Integer), false, Scope::default(), 0),
-                    Local::new(1, Some(Type::Integer), false, Scope::default(), 1)
+                    Local::new(0, Type::Integer, false, Scope::default()),
+                    Local::new(1, Type::Integer, false, Scope::default())
                 ]
             ),
             FunctionType {
                 type_parameters: None,
                 value_parameters: Some(vec![(0, Type::Integer), (1, Type::Integer)]),
-                return_type: Some(Box::new(Type::Integer)),
+                return_type: Box::new(Type::Integer),
             }
         )))
     );
@@ -53,14 +53,14 @@ fn function_call() {
                         ],
                         vec![Value::string("a"), Value::string("b"),],
                         vec![
-                            Local::new(0, Some(Type::Integer), false, Scope::default(), 0),
-                            Local::new(1, Some(Type::Integer), false, Scope::default(), 1)
+                            Local::new(0, Type::Integer, false, Scope::default()),
+                            Local::new(1, Type::Integer, false, Scope::default())
                         ]
                     ),
                     FunctionType {
                         type_parameters: None,
                         value_parameters: Some(vec![(0, Type::Integer), (1, Type::Integer)]),
-                        return_type: Some(Box::new(Type::Integer)),
+                        return_type: Box::new(Type::Integer),
                     }
                 ),
                 Value::integer(1),
@@ -97,28 +97,27 @@ fn function_declaration() {
                         ],
                         vec![Value::string("a"), Value::string("b")],
                         vec![
-                            Local::new(0, Some(Type::Integer), false, Scope::default(), 0),
-                            Local::new(1, Some(Type::Integer), false, Scope::default(), 1)
+                            Local::new(0, Type::Integer, false, Scope::default()),
+                            Local::new(1, Type::Integer, false, Scope::default())
                         ]
                     ),
                     FunctionType {
                         type_parameters: None,
                         value_parameters: Some(vec![(0, Type::Integer), (1, Type::Integer)]),
-                        return_type: Some(Box::new(Type::Integer)),
+                        return_type: Box::new(Type::Integer),
                     },
                 )
             ],
             vec![Local::new(
                 0,
-                Some(Type::Function(FunctionType {
+                Type::Function(FunctionType {
                     type_parameters: None,
                     value_parameters: Some(vec![(0, Type::Integer), (1, Type::Integer)]),
-                    return_type: Some(Box::new(Type::Integer)),
-                })),
+                    return_type: Box::new(Type::Integer),
+                }),
                 false,
                 Scope::default(),
-                0
-            ),],
+            )],
         )),
     );
 

dust-lang/tests/logic.rs

@@ -67,8 +67,8 @@ fn variable_and() {
             ],
             vec![Value::string("a"), Value::string("b"),],
             vec![
-                Local::new(0, None, false, Scope::default(), 0),
-                Local::new(1, None, false, Scope::default(), 1),
+                Local::new(0, Type::Boolean, false, Scope::default()),
+                Local::new(1, Type::Boolean, false, Scope::default()),
             ]
         ))
     );

dust-lang/tests/loops.rs

@@ -16,7 +16,7 @@ fn r#while() {
                     Span(23, 24)
                 ),
                 (Instruction::jump(2, true), Span(41, 42)),
-                (*Instruction::add(0, 0, 3).set_c_is_constant(), Span(39, 40)),
+                (*Instruction::add(0, 0, 3).set_c_is_constant(), Span(35, 36)),
                 (Instruction::jump(3, false), Span(41, 42)),
                 (Instruction::get_local(1, 0), Span(41, 42)),
                 (Instruction::r#return(true), Span(42, 42)),
@@ -27,7 +27,7 @@ fn r#while() {
                 Value::integer(5),
                 Value::integer(1),
             ],
-            vec![Local::new(1, None, true, Scope::default(), 0),]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         )),
     );
 

dust-lang/tests/math.rs

@@ -41,7 +41,7 @@ fn add_assign() {
                 (Instruction::r#return(true), Span(24, 24))
             ],
             vec![Value::integer(1), Value::string("a"), Value::integer(2)],
-            vec![Local::new(1, None, true, Scope::default(), 0)]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         ))
     );
 
@@ -124,7 +124,7 @@ fn divide_assign() {
                 (Instruction::r#return(true), Span(24, 24))
             ],
             vec![Value::integer(2), Value::string("a")],
-            vec![Local::new(1, None, true, Scope::default(), 0)]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         ))
     );
 
@@ -233,7 +233,7 @@ fn multiply_assign() {
                 (Instruction::r#return(true), Span(23, 23))
             ],
             vec![Value::integer(2), Value::string("a"), Value::integer(3)],
-            vec![Local::new(1, None, true, Scope::default(), 0),]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         ))
     );
 
@@ -300,7 +300,7 @@ fn subtract_assign() {
                 (Instruction::r#return(true), Span(25, 25)),
             ],
             vec![Value::integer(42), Value::string("x"), Value::integer(2)],
-            vec![Local::new(1, None, true, Scope::default(), 0)]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         )),
     );
 

dust-lang/tests/scopes.rs

@@ -55,11 +55,11 @@ fn block_scope() {
                 Value::string("e"),
             ],
             vec![
-                Local::new(1, None, false, Scope::new(0, 0), 0),
-                Local::new(3, None, false, Scope::new(1, 1), 1),
-                Local::new(5, None, false, Scope::new(2, 2), 2),
-                Local::new(7, None, false, Scope::new(1, 1), 3),
-                Local::new(8, None, false, Scope::new(0, 0), 4),
+                Local::new(1, Type::Integer, false, Scope::new(0, 0)),
+                Local::new(3, Type::Integer, false, Scope::new(1, 1)),
+                Local::new(5, Type::Integer, false, Scope::new(2, 2)),
+                Local::new(7, Type::Integer, false, Scope::new(1, 1)),
+                Local::new(8, Type::Integer, false, Scope::new(0, 0)),
             ]
         )),
     );
@@ -127,15 +127,15 @@ fn multiple_block_scopes() {
                 Value::string("e"),
             ],
             vec![
-                Local::new(1, None, false, Scope::new(0, 0), 0),
-                Local::new(3, None, false, Scope::new(1, 1), 1),
-                Local::new(5, None, false, Scope::new(2, 2), 2),
-                Local::new(7, None, false, Scope::new(1, 1), 3),
-                Local::new(8, None, false, Scope::new(0, 0), 4),
-                Local::new(3, None, false, Scope::new(1, 3), 5),
-                Local::new(5, None, false, Scope::new(2, 4), 6),
-                Local::new(7, None, false, Scope::new(1, 3), 7),
-                Local::new(9, None, false, Scope::new(0, 0), 8),
+                Local::new(1, Type::Integer, false, Scope::new(0, 0)),
+                Local::new(3, Type::Integer, false, Scope::new(1, 1)),
+                Local::new(5, Type::Integer, false, Scope::new(2, 2)),
+                Local::new(7, Type::Integer, false, Scope::new(1, 1)),
+                Local::new(8, Type::Integer, false, Scope::new(0, 0)),
+                Local::new(3, Type::Integer, false, Scope::new(1, 3)),
+                Local::new(5, Type::Integer, false, Scope::new(2, 4)),
+                Local::new(7, Type::Integer, false, Scope::new(1, 3)),
+                Local::new(9, Type::Integer, false, Scope::new(0, 0)),
             ]
         )),
     );

dust-lang/tests/variables.rs

@@ -14,7 +14,7 @@ fn define_local() {
                 (Instruction::r#return(false), Span(11, 11))
             ],
             vec![Value::integer(42), Value::string("x")],
-            vec![Local::new(1, None, false, Scope::default(), 0)]
+            vec![Local::new(1, Type::Integer, false, Scope::default())]
         )),
     );
 
@@ -55,7 +55,7 @@ fn set_local() {
                 (Instruction::r#return(true), Span(25, 25)),
             ],
             vec![Value::integer(41), Value::string("x"), Value::integer(42)],
-            vec![Local::new(1, None, true, Scope::default(), 0)]
+            vec![Local::new(1, Type::Integer, true, Scope::default())]
         )),
     );
 

dust-shell/src/main.rs

@@ -1,8 +1,11 @@
-use std::{fs::read_to_string, io::Write};
+use std::{
+    fs::read_to_string,
+    io::{stdout, Write},
+};
 
 use clap::Parser;
 use colored::Colorize;
-use dust_lang::{compile, format, run};
+use dust_lang::{compile, format, lex, output_token_list, run};
 use log::{Level, LevelFilter};
 
 #[derive(Parser)]
@@ -11,7 +14,7 @@ struct Cli {
     #[arg(short, long)]
     command: Option<String>,
 
-    /// Whether to output formatted source code
+    /// Whether to output formatted source code instead of running the program
     #[arg(short, long)]
     format: bool,
 
@@ -23,7 +26,7 @@ struct Cli {
     #[arg(long)]
     format_colored: Option<bool>,
 
-    /// Whether to output the disassembled chunk
+    /// Whether to output the disassembled chunk instead of running the program
     #[arg(short, long)]
     parse: bool,
 
@@ -31,6 +34,10 @@ struct Cli {
     #[arg(long)]
     style_disassembly: Option<bool>,
 
+    /// Whether to tokenize the source code instead of running the program
+    #[arg(short, long)]
+    tokenize: bool,
+
     /// Log level
     #[arg(short, long)]
     log: Option<LevelFilter>,
@@ -78,11 +85,11 @@ fn main() {
     };
 
     if args.format {
+        log::info!("Formatting source");
+
         let line_numbers = args.format_line_numbers.unwrap_or(true);
         let colored = args.format_colored.unwrap_or(true);
 
-        log::info!("Formatting source");
-
         match format(source, line_numbers, colored) {
             Ok(formatted) => println!("{}", formatted),
             Err(error) => {
@@ -91,11 +98,20 @@ fn main() {
         }
     }
 
-    if args.parse {
-        let styled = args.style_disassembly.unwrap_or(true);
+    if args.tokenize {
+        log::info!("Tokenizing source");
 
+        match lex(source) {
+            Ok(tokens) => output_token_list(&tokens, &mut stdout()),
+            Err(error) => eprintln!("{}", error.report()),
+        }
+    }
+
+    if args.parse {
         log::info!("Parsing source");
 
+        let styled = args.style_disassembly.unwrap_or(true);
+
         match compile(source) {
             Ok(chunk) => {
                 let disassembly = chunk
@@ -112,7 +128,7 @@ fn main() {
         }
     }
 
-    if args.format || args.parse {
+    if args.format || args.tokenize || args.parse {
         return;
     }