Add support for optional chaining via the ? operator

CHANGELOG.md

@@ -14,6 +14,9 @@ The Koto project adheres to
 
 - Type hints with runtime type checks have been added ([#298](https://github.com/koto-lang/koto/issues/298)).
   - Thanks to [@Tarbetu](https://github.com/Tarbetu) for the contributions.
+- Optional chaining via the `?` operator has been added to simplify writing
+  expressions that need to check for `.null` on intermediate values.
+  - E.g. `x = get_data()?.bar()`, where `get_data` could return `null`.
 - `export` can be used with multi-assignment expressions.
   - E.g. expressions like `export a, b, c = foo()` are now allowed.
 - Maps now support `[]` indexing, returning the Nth entry as a tuple.

crates/bytecode/src/compiler.rs

@@ -2433,12 +2433,14 @@ impl Compiler {
         // so we don't need to keep track of how many temporary registers we use.
         let stack_count = self.stack_count();
         let span_stack_count = self.span_stack.len();
-        // let chain_temp_registers_start = self.frame().next_temporary_register();
 
         let mut chain_node = root_node.clone();
 
+        let mut null_check_jump_placeholders = SmallVec::<[usize; 4]>::new();
+        let mut null_check_on_end_node = false;
+
         // Work through the chain, up until the last node, which will be handled separately
-        while next_node_index.is_some() {
+        while let Some(next) = next_node_index {
             match &chain_node {
                 ChainNode::Root(root_node) => {
                     if !chain_nodes.is_empty() {
@@ -2535,27 +2537,42 @@ impl Compiler {
                         ctx.with_fixed_register(node_register),
                     )?;
                 }
-            }
+                ChainNode::NullCheck => {
+                    // `?` null check
+                    let Some(check_register) = chain_nodes.previous() else {
+                        return self.error(ErrorKind::OutOfPositionChildNodeInChain);
+                    };
 
-            // Is the chain complete?
-            let Some(next) = next_node_index else { break };
+                    self.push_op(JumpIfNull, &[check_register]);
+                    null_check_jump_placeholders.push(self.push_offset_placeholder());
+                }
+            }
 
             let next_chain_node = ctx.node_with_span(next);
+            self.push_span(next_chain_node, ctx.ast);
 
-            match next_chain_node.node.clone() {
+            match &next_chain_node.node {
                 Node::Chain((node, next)) => {
-                    chain_node = node;
-                    next_node_index = next;
+                    chain_node = node.clone();
+                    next_node_index = *next;
+
+                    // If the last node in the chain is a null check then break out now,
+                    // allowing the final node (before the null check) to be held in `chain_node`
+                    // for further processing below, after this loop.
+                    if let Some(next) = *next {
+                        if ctx.node(next) == &Node::Chain((ChainNode::NullCheck, None)) {
+                            null_check_on_end_node = true;
+                            break;
+                        }
+                    }
                 }
                 unexpected => {
                     return self.error(ErrorKind::UnexpectedNode {
                         expected: "a chain node".into(),
-                        unexpected,
+                        unexpected: unexpected.clone(),
                     });
                 }
-            };
-
-            self.push_span(next_chain_node, ctx.ast);
+            }
         }
 
         // The chain is complete up until the last node, and now we need to handle:
@@ -2596,16 +2613,20 @@ impl Compiler {
         debug_assert!(rhs_op.is_none() || rhs_op.is_some() && rhs.is_some());
         let simple_assignment = rhs.is_some() && rhs_op.is_none();
         let compound_assignment = rhs.is_some() && rhs_op.is_some();
+        let access_end_node = !simple_assignment || null_check_on_end_node;
 
-        // Do we need to read the last value in the lookup chain?
-        // - No if it's a simple assignment and the last value is going to be overwritten.
-        // - Yes otherwise, either there's a compound assignment or the last value is being accessed.
+        // Do we need to access the last node in the lookup chain?
+        // - No if it's a simple assignment (without a null check) and the last node is going to be
+        //   overwritten.
+        // - Yes otherwise, either there's a compound assignment, or a null check,
+        //   or the last node is the result.
         match &end_node {
-            ChainNode::Id(id, ..) if !simple_assignment => {
+            ChainNode::Id(id, ..) if access_end_node => {
+                dbg!(result_register);
                 self.compile_access_id(result_register, container_register, *id);
                 chain_nodes.push(result_register, false);
             }
-            ChainNode::Str(_) if !simple_assignment => {
+            ChainNode::Str(_) if access_end_node => {
                 self.push_op(
                     AccessString,
                     &[
@@ -2616,7 +2637,7 @@ impl Compiler {
                 );
                 chain_nodes.push(result_register, false);
             }
-            ChainNode::Index(_) if !simple_assignment => {
+            ChainNode::Index(_) if access_end_node => {
                 self.push_op(
                     Index,
                     &[result_register, container_register, index.unwrap(self)?],
@@ -2626,7 +2647,12 @@ impl Compiler {
             ChainNode::Call { args, with_parens } => {
                 if simple_assignment {
                     return self.error(ErrorKind::AssigningToATemporaryValue);
-                } else if compound_assignment || piped_arg_register.is_none() || *with_parens {
+                } else if compound_assignment // e.g. `f() += 1 # (valid depending on return type)`
+                    // If there's a piped call on the chain result, defer it until below
+                    || piped_arg_register.is_none()
+                    // Parenthesized calls need to be made now, e.g. `42 -> foo.bar()`
+                    || *with_parens
+                {
                     let (function_register, instance_register) = chain_nodes.previous_two();
 
                     let Some(function_register) = function_register else {
@@ -2646,9 +2672,15 @@ impl Compiler {
             _ => {}
         }
 
+        // Is a null check needed on the last node?
+        if null_check_on_end_node {
+            self.push_op(JumpIfNull, &[result_register]);
+            null_check_jump_placeholders.push(self.push_offset_placeholder());
+        }
+
         // Do we need to modify the accessed value?
         if compound_assignment {
-            // compound_assignment can only be true when both rhs and rhs_op have values
+            // Compound_assignment can only be true when both rhs and rhs_op have values
             let rhs = rhs.unwrap();
             let rhs_op = rhs_op.unwrap();
 
@@ -2714,6 +2746,13 @@ impl Compiler {
             )?;
         }
 
+        // Now that all chain operations are complete, update the null check jump placeholders so
+        // that they point to the end of the chain.
+
+        for placeholder in null_check_jump_placeholders {
+            self.update_offset_placeholder(placeholder)?;
+        }
+
         // Clean up the span and register stacks
         self.span_stack.truncate(span_stack_count);
         self.truncate_register_stack(stack_count)?;

crates/bytecode/src/instruction.rs

@@ -218,6 +218,10 @@ pub enum Instruction {
         register: u8,
         offset: u16,
     },
+    JumpIfNull {
+        register: u8,
+        offset: u16,
+    },
     Call {
         result: u8,
         function: u8,
@@ -625,6 +629,9 @@ impl fmt::Debug for Instruction {
             JumpIfFalse { register, offset } => {
                 write!(f, "JumpIfFalse\tresult: {register}\toffset: {offset}")
             }
+            JumpIfNull { register, offset } => {
+                write!(f, "JumpIfNull\tresult: {register}\toffset: {offset}")
+            }
             Call {
                 result,
                 function,

crates/bytecode/src/instruction_reader.rs

@@ -316,6 +316,10 @@ impl Iterator for InstructionReader {
                 register: get_u8!(),
                 offset: get_u16!(),
             }),
+            Op::JumpIfNull => Some(JumpIfNull {
+                register: get_u8!(),
+                offset: get_u16!(),
+            }),
             Op::Call => Some(Call {
                 result: get_u8!(),
                 function: get_u8!(),

crates/bytecode/src/op.rs

@@ -303,15 +303,20 @@ pub enum Op {
     /// `[offset[2]]`
     JumpBack,
 
-    /// Causes the instruction pointer to jump forward, if a condition is true
+    /// Jumps the instruction pointer forward if a value is `false` or `null`
     ///
-    /// `[*condition, offset[2]]`
+    /// `[*value, offset[2]]`
+    JumpIfFalse,
+
+    /// Jumps the instruction pointer forward if a value is neither `false` or `null`
+    ///
+    /// `[*value, offset[2]]`
     JumpIfTrue,
 
-    /// Causes the instruction pointer to jump forward, if a condition is false
+    /// Jumps the instruction pointer forward if a value is `null`
     ///
-    /// `[*condition, offset[2]]`
-    JumpIfFalse,
+    /// `[*value, offset[2]]`
+    JumpIfNull,
 
     /// Calls a standalone function
     ///
@@ -522,7 +527,6 @@ pub enum Op {
     CheckType,
 
     // Unused opcodes, allowing for a direct transmutation from a byte to an Op.
-    Unused85,
     Unused86,
     Unused87,
     Unused88,

crates/cli/docs/language_guide.md

@@ -916,11 +916,49 @@ print! match ['a', 'b', 'c'].extend [1, 2, 3]
 check! Starts with 'a', followed by 'b', then 4 others
 ```
 
+### Optional Chaining
+
+The `?` operator can be used to short-circuit expression chains where `null`
+might be encountered as an intermediate value. The `?` checks the current value 
+in the expression chain and if `null` is found then the chain is short-circuited 
+with `null` given as the expression's result.
+
+This makes it easier to check for `null` when you want to avoid runtime errors.
+
+```koto
+info = {town: 'Hamburg', country: 'Germany'}
+
+# `info` contains a value for 'town', which is then passed to to_uppercase():
+print! info.get('town')?.to_uppercase()
+check! HAMBURG
+
+# `info` doesn't contain a value for 'state', 
+# so the `?` operator short-circuits the expression, resulting in `null`:
+print! info.get('state')?.to_uppercase()
+check! null
+
+# Without the `?` operator an intermediate step is necessary:
+country = info.get('country')
+print! if country then country.to_uppercase()
+check! GERMANY
+```
+
+Multiple `?` checks can be performed in an expression chain:
+
+```koto
+get_data = || {nested: {maybe_string: null}}
+print! get_data()?
+  .get('nested')?
+  .get('maybe_string')?
+  .to_uppercase()
+check! null
+```
+
 ## Loops
 
 Koto includes several ways of evaluating expressions repeatedly in a loop.
 
-### for
+### `for`
 
 `for` loops are repeated for each element in a sequence, 
 such as a list or tuple.
@@ -933,7 +971,7 @@ check! 20
 check! 30
 ```
 
-### while
+### `while`
 
 `while` loops continue to repeat _while_ a condition is true.
 
@@ -945,7 +983,7 @@ print! x
 check! 5
 ```
 
-### until
+### `until`
 
 `until` loops continue to repeat _until_ a condition is true.
 

crates/lexer/src/lexer.rs

@@ -36,6 +36,7 @@ pub enum Token {
     CurlyClose,
     Range,
     RangeInclusive,
+    QuestionMark,
 
     // operators
     Add,
@@ -736,6 +737,7 @@ impl<'a> TokenLexer<'a> {
         check_symbol!(":", Colon);
         check_symbol!(",", Comma);
         check_symbol!(".", Dot);
+        check_symbol!("?", QuestionMark);
         check_symbol!("(", RoundOpen);
         check_symbol!(")", RoundClose);
         check_symbol!("|", Function);

crates/parser/src/error.rs

@@ -193,6 +193,8 @@ pub enum SyntaxError {
     UnexpectedMetaKey,
     #[error("Unexpected 'else' in switch arm")]
     UnexpectedSwitchElse,
+    #[error("Unexpected '?'")]
+    UnexpectedNullCheck,
     #[error("Unexpected token")]
     UnexpectedToken,
     #[error("Unicode value out of range, the maximum is \\u{{10ffff}}")]

crates/parser/src/node.rs

@@ -427,8 +427,9 @@ pub struct AstCatch {
 /// In other words, some series of operations involving indexing, `.` accesses, and function calls.
 ///
 /// e.g.
-/// `foo.bar."baz"[0](42)`
-///  |  |   |     |  ^ Call {args: 42, with_parens: true}
+/// `foo.bar."baz"[0]?(42)`
+///  |  |   |     |  |^ Call {args: 42, with_parens: true}
+///  |  |   |     |  ^ NullCheck
 ///  |  |   |     ^ Index (0)
 ///  |  |   ^ Str (baz)
 ///  |  ^ Id (bar)
@@ -457,6 +458,8 @@ pub enum ChainNode {
         ///   `99 -> foo.bar(42)` is equivalent to `foo.bar(42)(99)`.
         with_parens: bool,
     },
+    /// A `?` short-circuiting null check
+    NullCheck,
 }
 
 /// An arm in a match expression

crates/parser/src/parser.rs

@@ -1320,7 +1320,10 @@ impl<'source> Parser<'source> {
     fn next_token_is_chain_start(&mut self, context: &ExpressionContext) -> bool {
         use Token::*;
 
-        if matches!(self.peek_token(), Some(Dot | SquareOpen | RoundOpen)) {
+        if matches!(
+            self.peek_token(),
+            Some(Dot | SquareOpen | RoundOpen | QuestionMark)
+        ) {
             true
         } else if context.allow_linebreaks {
             matches!(
@@ -1404,6 +1407,15 @@ impl<'source> Parser<'source> {
                         return self.consume_token_and_error(SyntaxError::ExpectedMapKey);
                     }
                 }
+                // Null check
+                Token::QuestionMark => {
+                    self.consume_token();
+                    if matches!(chain.last(), Some((ChainNode::NullCheck, _))) {
+                        return self.error(SyntaxError::UnexpectedNullCheck);
+                    }
+
+                    chain.push((ChainNode::NullCheck, node_start_span));
+                }
                 _ => {
                     let Some(peeked) = self.peek_token_with_context(&node_context) else {
                         break;