Add postfix `!` operator for optionals

6 files changed, 68 insertions, 26 deletions

diff --git a/ast.c b/ast.c
index 1e89d29b..dc2ff13b 100644
--- a/ast.c
+++ b/ast.c
@@ -151,6 +151,7 @@ CORD ast_to_xml(ast_t *ast)
     T(Index, "<Index>%r%r</Index>", optional_tagged("indexed", data.indexed), optional_tagged("index", data.index))
     T(FieldAccess, "<FieldAccess field=\"%s\">%r</FieldAccess>", data.field, ast_to_xml(data.fielded))
     T(Optional, "<Optional>%r</Optional>", ast_to_xml(data.value))
+    T(NonOptional, "<NonOptional>%r</NonOptional>", ast_to_xml(data.value))
     T(DocTest, "<DocTest>%r<output>%r</output></DocTest>", optional_tagged("expression", data.expr), xml_escape(data.output))
     T(Use, "<Use>%r</Use>", xml_escape(data.path))
     T(LinkerDirective, "<LinkerDirective>%r</LinkerDirective>", xml_escape(data.directive))
diff --git a/ast.h b/ast.h
index e3108cae..9da7ebe9 100644
--- a/ast.h
+++ b/ast.h
@@ -124,7 +124,7 @@ typedef enum {
     Return,
     Extern,
     StructDef, EnumDef, LangDef,
-    Index, FieldAccess, Optional,
+    Index, FieldAccess, Optional, NonOptional,
     DocTest,
     Use,
     LinkerDirective,
@@ -288,7 +288,7 @@ struct ast_s {
         } FieldAccess;
         struct {
             ast_t *value;
-        } Optional;
+        } Optional, NonOptional;
         struct {
             ast_t *expr;
             const char *output;
diff --git a/compile.c b/compile.c
index d684d1ab..cc334e76 100644
--- a/compile.c
+++ b/compile.c
@@ -317,19 +317,18 @@ static CORD compile_inline_block(env_t *env, ast_t *ast)
     return code;
 }
 
-static CORD optional_var_into_nonnull(binding_t *b)
+static CORD optional_into_nonnull(type_t *t, CORD value)
 {
-    type_t *t = b->type->tag == OptionalType ? Match(b->type, OptionalType)->type : b->type;
+    if (t->tag == OptionalType) t = Match(t, OptionalType)->type;
     switch (t->tag) {
-    case OptionalType:
     case IntType:
-        return CORD_all(b->code, ".i");
+        return CORD_all(value, ".i");
     case StructType:
         if (t == THREAD_TYPE)
-            return b->code;
-        return CORD_all(b->code, ".value");
+            return value;
+        return CORD_all(value, ".value");
     default:
-        return b->code;
+        return value;
     }
 }
 
@@ -1239,7 +1238,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
             if (for_->vars) {
                 naked_body = CORD_all(
                     compile_declaration(Match(fn->ret, OptionalType)->type, CORD_all("$", Match(for_->vars->ast, Var)->name)),
-                    " = ", optional_var_into_nonnull(new(binding_t, .type=fn->ret, .code="cur")), ";\n",
+                    " = ", optional_into_nonnull(fn->ret, "cur"), ";\n",
                     naked_body);
             }
 
@@ -1283,11 +1282,10 @@ CORD compile_statement(env_t *env, ast_t *ast)
                 truthy_scope = fresh_scope(env);
                 const char *varname = Match(condition, Var)->name;
                 binding_t *b = get_binding(env, varname);
-                binding_t *nonnull_b = new(binding_t);
-                *nonnull_b = *b;
-                nonnull_b->type = Match(cond_t, OptionalType)->type;
-                nonnull_b->code = optional_var_into_nonnull(b);
-                set_binding(truthy_scope, varname, nonnull_b);
+                assert(b);
+                set_binding(truthy_scope, varname,
+                            new(binding_t, .type=Match(cond_t, OptionalType)->type,
+                                .code=optional_into_nonnull(cond_t, b->code)));
             }
             condition_code = CORD_all("!", check_null(cond_t, compile(env, condition)));
         } else if (cond_t->tag == BoolType) {
@@ -1858,6 +1856,18 @@ CORD compile(env_t *env, ast_t *ast)
         CORD value_code = compile(env, value);
         return promote_to_optional(get_type(env, value), value_code);
     }
+    case NonOptional: {
+        ast_t *value = Match(ast, NonOptional)->value;
+        type_t *t = get_type(env, value);
+        CORD value_code = compile(env, value);
+        return CORD_all("({ ", compile_declaration(t, "opt"), " = ", value_code, "; ",
+                        "if (", check_null(t, "opt"), ")\n",
+                        CORD_asprintf("fail_source(%r, %ld, %ld, \"This value was expected to be non-null, but it's null!\");\n",
+                                      CORD_quoted(ast->file->filename),
+                                      (long)(value->start - value->file->text),
+                                      (long)(value->end - value->file->text)),
+                        optional_into_nonnull(t, "opt"), "; })");
+    }
     case BinaryOp: {
         auto binop = Match(ast, BinaryOp);
         CORD method_call = compile_math_method(env, binop->op, binop->lhs, binop->rhs, NULL);
@@ -2750,7 +2760,7 @@ CORD compile(env_t *env, ast_t *ast)
                 return CORD_all("({ ", compile_declaration(self_value_t, "opt"), " = ", self, "; ",
                                 check_null(self_value_t, "opt"), " ? ",
                                 compile_arguments(env, ast, arg_spec, call->args),
-                                " : ", optional_var_into_nonnull(new(binding_t, .type=self_value_t, .code="opt")), "; })");
+                                " : ", optional_into_nonnull(self_value_t, "opt"), "; })");
             } else if (streq(call->name, "or_fail")) {
                 CORD self = compile_to_pointer_depth(env, call->self, 0, false);
                 arg_t *arg_spec = new(arg_t, .name="message", .type=TEXT_TYPE,
@@ -2762,7 +2772,7 @@ CORD compile(env_t *env, ast_t *ast)
                                               (long)(call->self->start - call->self->file->text),
                                               (long)(call->self->end - call->self->file->text),
                                               compile_arguments(env, ast, arg_spec, call->args)),
-                                optional_var_into_nonnull(new(binding_t, .type=self_value_t, .code="opt")), "; })");
+                                optional_into_nonnull(self_value_t, "opt"), "; })");
             }
             code_err(ast, "There is no '%s' method for optional %T values", call->name, nonnull);
         }
@@ -3550,11 +3560,11 @@ CORD compile_cli_arg_call(env_t *env, CORD fn_name, type_t *fn_type)
 
     code = CORD_all(code, fn_name, "(");
     for (arg_t *arg = fn_info->args; arg; arg = arg->next) {
-        if (arg->type->tag == OptionalType)
-            code = CORD_all(code, "$", arg->name);
-        else
-            code = CORD_all(code, optional_var_into_nonnull(get_binding(main_env, arg->name)));
-            
+        CORD arg_code = CORD_all("$", arg->name);
+        if (arg->type->tag != OptionalType)
+            arg_code = optional_into_nonnull(arg->type, arg_code);
+
+        code = CORD_all(code, arg_code);
         if (arg->next) code = CORD_all(code, ", ");
     }
     code = CORD_all(code, ");\n");
diff --git a/parse.c b/parse.c
index e33b9ef2..8ce53944 100644
--- a/parse.c
+++ b/parse.c
@@ -83,6 +83,7 @@ static ast_t *parse_index_suffix(parse_ctx_t *ctx, ast_t *lhs);
 static ast_t *parse_comprehension_suffix(parse_ctx_t *ctx, ast_t *lhs);
 static ast_t *parse_optional_conditional_suffix(parse_ctx_t *ctx, ast_t *lhs);
 static ast_t *parse_optional_suffix(parse_ctx_t *ctx, ast_t *lhs);
+static ast_t *parse_non_optional_suffix(parse_ctx_t *ctx, ast_t *lhs);
 static arg_ast_t *parse_args(parse_ctx_t *ctx, const char **pos, bool allow_unnamed);
 static type_ast_t *parse_non_optional_type(parse_ctx_t *ctx, const char *pos);
 static type_ast_t *parse_type(parse_ctx_t *ctx, const char *pos);
@@ -883,6 +884,15 @@ ast_t *parse_optional_suffix(parse_ctx_t *ctx, ast_t *lhs) {
         return NULL;
 }
 
+ast_t *parse_non_optional_suffix(parse_ctx_t *ctx, ast_t *lhs) {
+    if (!lhs) return NULL;
+    const char *pos = lhs->end;
+    if (match(&pos, "!"))
+        return NewAST(ctx->file, lhs->start, pos, NonOptional, .value=lhs);
+    else
+        return NULL;
+}
+
 PARSER(parse_reduction) {
     const char *start = pos;
     if (!match(&pos, "(")) return NULL;
@@ -905,6 +915,7 @@ PARSER(parse_reduction) {
                 || (new_term=parse_method_call_suffix(ctx, key))
                 || (new_term=parse_fncall_suffix(ctx, key))
                 || (new_term=parse_optional_suffix(ctx, key))
+                || (new_term=parse_non_optional_suffix(ctx, key))
                 );
             if (progress) key = new_term;
         }
@@ -1162,8 +1173,12 @@ PARSER(parse_heap_alloc) {
     pos = val->end;
 
     ast_t *ast = NewAST(ctx->file, start, pos, HeapAllocate, .value=val);
-    ast_t *optional = parse_optional_suffix(ctx, ast);
-    if (optional) ast = optional;
+    for (;;) {
+        ast_t *next = parse_optional_suffix(ctx, ast);
+        if (!next) next = parse_non_optional_suffix(ctx, ast);
+        if (!next) break;
+        ast = next;
+    }
     return ast;
 }
 
@@ -1184,8 +1199,12 @@ PARSER(parse_stack_reference) {
     pos = val->end;
 
     ast_t *ast = NewAST(ctx->file, start, pos, StackReference, .value=val);
-    ast_t *optional = parse_optional_suffix(ctx, ast);
-    if (optional) ast = optional;
+    for (;;) {
+        ast_t *next = parse_optional_suffix(ctx, ast);
+        if (!next) next = parse_non_optional_suffix(ctx, ast);
+        if (!next) break;
+        ast = next;
+    }
     return ast;
 }
 
@@ -1529,6 +1548,7 @@ PARSER(parse_term) {
             || (new_term=parse_method_call_suffix(ctx, term))
             || (new_term=parse_fncall_suffix(ctx, term))
             || (new_term=parse_optional_suffix(ctx, term))
+            || (new_term=parse_non_optional_suffix(ctx, term))
             );
         if (progress) term = new_term;
     }
@@ -1672,6 +1692,7 @@ static ast_t *parse_infix_expr(parse_ctx_t *ctx, const char *pos, int min_tightn
                     || (new_term=parse_method_call_suffix(ctx, key))
                     || (new_term=parse_fncall_suffix(ctx, key))
                     || (new_term=parse_optional_suffix(ctx, key))
+                    || (new_term=parse_non_optional_suffix(ctx, key))
                     );
                 if (progress) key = new_term;
             }
diff --git a/test/optionals.tm b/test/optionals.tm
index 26e825ae..c0c600a7 100644
--- a/test/optionals.tm
+++ b/test/optionals.tm
@@ -271,3 +271,6 @@ func main():
         >> yep
         = 123 : Int
     else: fail("Unreachable")
+
+    >> maybe_int(yes)!
+    = 123 : Int
diff --git a/typecheck.c b/typecheck.c
index 9ad151a3..b70db784 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -567,6 +567,13 @@ type_t *get_type(env_t *env, ast_t *ast)
             code_err(ast, "This value is already optional, it can't be converted to optional");
         return Type(OptionalType, .type=t);
     }
+    case NonOptional: {
+        ast_t *value = Match(ast, NonOptional)->value;
+        type_t *t = get_type(env, value);
+        if (t->tag != OptionalType)
+            code_err(value, "This value is not optional. Only optional values can use the '!' operator.");
+        return Match(t, OptionalType)->type;
+    }
     case TextLiteral: return TEXT_TYPE;
     case TextJoin: {
         const char *lang = Match(ast, TextJoin)->lang;