Add table.get_or_set()

2 files changed, 10 insertions, 0 deletions

diff --git a/src/compile.c b/src/compile.c
index 59c25f6a..3f6cb713 100644
--- a/src/compile.c
+++ b/src/compile.c
@@ -3312,6 +3312,15 @@ CORD compile(env_t *env, ast_t *ast)
                     compile_type(table->value_type), ", ", compile_arguments(env, ast, arg_spec, call->args), ", ",
                     "_, ", optional_into_nonnone(table->value_type, "(*_)"), ", ", compile_none(table->value_type), ", ",
                     compile_type_info(self_value_t), ")");
+            } else if (streq(call->name, "get_or_set")) {
+                self = compile_to_pointer_depth(env, call->self, 1, false);
+                arg_t *arg_spec = new(arg_t, .name="key", .type=table->key_type,
+                                      .next=new(arg_t, .name="default", .type=table->value_type, .default_val=table->default_value));
+                return CORD_all("*Table$get_or_setdefault(",
+                                self, ", ", compile_type(table->key_type), ", ",
+                                compile_type(table->value_type), ", ",
+                                compile_arguments(env, ast, arg_spec, call->args), ", ",
+                                compile_type_info(self_value_t), ")");
             } else if (streq(call->name, "has")) {
                 self = compile_to_pointer_depth(env, call->self, 0, false);
                 arg_t *arg_spec = new(arg_t, .name="key", .type=table->key_type);
diff --git a/src/typecheck.c b/src/typecheck.c
index e42f6ca8..58b74739 100644
--- a/src/typecheck.c
+++ b/src/typecheck.c
@@ -938,6 +938,7 @@ type_t *get_type(env_t *env, ast_t *ast)
             auto table = Match(self_value_t, TableType);
             if (streq(call->name, "clear")) return Type(VoidType);
             else if (streq(call->name, "get")) return Type(OptionalType, .type=table->value_type);
+            else if (streq(call->name, "get_or_set")) return table->value_type;
             else if (streq(call->name, "has")) return Type(BoolType);
             else if (streq(call->name, "remove")) return Type(VoidType);
             else if (streq(call->name, "set")) return Type(VoidType);