aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--docs/tables.md38
-rw-r--r--src/compile.c9
-rw-r--r--src/typecheck.c1
3 files changed, 48 insertions, 0 deletions
diff --git a/docs/tables.md b/docs/tables.md
index ee2603f5..d83a678b 100644
--- a/docs/tables.md
+++ b/docs/tables.md
@@ -154,6 +154,7 @@ iterating over any of the new values.
- [`func clear(t:&{K=V})`](#clear)
- [`func get(t:{K=V}, key: K -> V?)`](#get)
+- [`func get_or_set(t:&{K=V}, key: K, default: V -> V)`](#get_or_set)
- [`func has(t:{K=V}, key: K -> Bool)`](#has)
- [`func remove(t:{K=V}, key: K -> Void)`](#remove)
- [`func set(t:{K=V}, key: K, value: V -> Void)`](#set)
@@ -211,6 +212,43 @@ The value associated with the key or `none` if the key is not found.
---
+### `get_or_set`
+If the given key is in the table, return the associated value. Otherwise,
+insert the given default value into the table and return it.
+
+**Note:** If no default value is provided explicitly, but the table has a
+default value associated with it, the table's default value will be used.
+
+**Note:** The default value is only evaluated if the key is missing.
+
+```tomo
+func get_or_set(t: &{K=V}, key: K, default: V -> V?)
+```
+
+- `t`: The table.
+- `key`: The key whose associated value is to be retrieved.
+- `default`: The default value to insert and return if the key is not present in the table.
+
+**Returns:**
+Either the value associated with the key (if present) or the default value. The
+table will be mutated if the key is not already present.
+
+**Example:**
+```tomo
+>> t := &{"A"=@[1, 2, 3]; default=@[]}
+>> t.get_or_set("A").insert(4)
+>> t.get_or_set("B").insert(99)
+>> t
+= &{"A"=@[1, 2, 3, 4], "B"=@[99]}
+
+>> t.get_or_set("C", @[0, 0, 0])
+= @[0, 0, 0]
+>> t
+= &{"A"=@[1, 2, 3, 4], "B"=@[99], "C"=@[0, 0, 0]}
+```
+
+---
+
### `has`
Checks if the table contains a specified key.
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);