Improvements and fixes for assigning to table keys

8 files changed, 148 insertions, 37 deletions

diff --git a/ast.c b/ast.c
index a1916089..421f806e 100644
--- a/ast.c
+++ b/ast.c
@@ -123,6 +123,7 @@ CORD ast_to_xml(ast_t *ast)
       ast_list_to_xml(data.items))
     T(Table, "<Table>%r%r%r%r</Table>",
       optional_tagged_type("key-type", data.key_type), optional_tagged_type("value-type", data.value_type),
+      optional_tagged("default-value", data.default_value),
       ast_list_to_xml(data.entries), optional_tagged("fallback", data.fallback))
     T(TableEntry, "<TableEntry>%r%r</TableEntry>", ast_to_xml(data.key), ast_to_xml(data.value))
     T(Channel, "<Channel>%r%r</Channel>", type_ast_to_xml(data.item_type), optional_tagged("max-size", data.max_size))
diff --git a/ast.h b/ast.h
index 0cd5a033..82d4de03 100644
--- a/ast.h
+++ b/ast.h
@@ -106,6 +106,7 @@ struct type_ast_s {
         } ArrayTypeAST, ChannelTypeAST;
         struct {
             type_ast_t *key, *value;
+            ast_t *default_value;
         } TableTypeAST;
         struct {
             type_ast_t *item;
@@ -211,6 +212,7 @@ struct ast_s {
         } Set;
         struct {
             type_ast_t *key_type, *value_type;
+            ast_t *default_value;
             ast_t *fallback;
             ast_list_t *entries;
         } Table;
diff --git a/compile.c b/compile.c
index 4da22daa..d570e338 100644
--- a/compile.c
+++ b/compile.c
@@ -357,6 +357,17 @@ static CORD compile_lvalue(env_t *env, ast_t *ast)
                                 ", ", heap_strf("%ld", ast->end - ast->file->text), ")");
             }
         } else if (container_t->tag == TableType) {
+            auto table_type = Match(container_t, TableType);
+            if (table_type->default_value) {
+                type_t *value_type = get_type(env, table_type->default_value);
+                return CORD_all("*Table$get_or_setdefault(",
+                                compile_lvalue(env, index->indexed), ", ",
+                                compile_type(table_type->key_type), ", ",
+                                compile_type(value_type), ", ",
+                                compile(env, index->index), ", ",
+                                compile(env, table_type->default_value), ", ",
+                                compile_type_info(env, container_t), ")");
+            }
             type_t *key_type = Match(container_t, TableType)->key_type;
             type_t *value_type = Match(container_t, TableType)->value_type;
             if (index->unchecked)
@@ -581,6 +592,9 @@ CORD compile_statement(env_t *env, ast_t *ast)
             if (!assign->targets->next && assign->targets->ast->tag == Var && is_idempotent(assign->targets->ast)) {
                 // Common case: assigning to one variable:
                 type_t *lhs_t = get_type(env, assign->targets->ast);
+                if (assign->targets->ast->tag == Index && lhs_t->tag == OptionalType
+                    && value_type(get_type(env, Match(assign->targets->ast, Index)->indexed))->tag == TableType)
+                    lhs_t = Match(lhs_t, OptionalType)->type;
                 env_t *val_scope = with_enum_scope(env, lhs_t);
                 CORD value = compile_to_type(val_scope, assign->values->ast, lhs_t);
                 return CORD_asprintf(
@@ -598,18 +612,24 @@ CORD compile_statement(env_t *env, ast_t *ast)
                 CORD code = "test(({ // Assignment\n";
 
                 int64_t i = 1;
+                type_t *first_type = NULL;
                 for (ast_list_t *target = assign->targets, *value = assign->values; target && value; target = target->next, value = value->next) {
-                    type_t *target_type = get_type(env, target->ast);
-                    env_t *val_scope = with_enum_scope(env, target_type);
-                    CORD val_code = compile_to_type(val_scope, value->ast, target_type);
-                    CORD_appendf(&code, "%r $%ld = %r;\n", compile_type(target_type), i++, val_code);
+                    type_t *lhs_t = get_type(env, target->ast);
+                    if (target->ast->tag == Index && lhs_t->tag == OptionalType
+                        && value_type(get_type(env, Match(target->ast, Index)->indexed))->tag == TableType)
+                        lhs_t = Match(lhs_t, OptionalType)->type;
+                    if (target == assign->targets)
+                        first_type = lhs_t;
+                    env_t *val_scope = with_enum_scope(env, lhs_t);
+                    CORD val_code = compile_to_type(val_scope, value->ast, lhs_t);
+                    CORD_appendf(&code, "%r $%ld = %r;\n", compile_type(lhs_t), i++, val_code);
                 }
                 i = 1;
                 for (ast_list_t *target = assign->targets; target; target = target->next)
                     code = CORD_all(code, compile_assignment(env, target->ast, CORD_asprintf("$%ld", i++)), ";\n");
 
                 CORD_appendf(&code, "$1; }), %r, %r, %ld, %ld);",
-                    compile_type_info(env, get_type(env, assign->targets->ast)),
+                    compile_type_info(env, first_type),
                     CORD_quoted(test->output),
                     (int64_t)(test->expr->start - test->expr->file->text),
                     (int64_t)(test->expr->end - test->expr->file->text));
@@ -621,7 +641,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
 
             if (update->lhs->tag == Index) {
                 type_t *indexed = value_type(get_type(env, Match(update->lhs, Index)->indexed));
-                if (indexed->tag == TableType)
+                if (indexed->tag == TableType && Match(indexed, TableType)->default_value == NULL)
                     code_err(update->lhs, "Update assignments are not currently supported for tables");
             }
 
@@ -674,8 +694,11 @@ CORD compile_statement(env_t *env, ast_t *ast)
     case Assign: {
         auto assign = Match(ast, Assign);
         // Single assignment, no temp vars needed:
-        if (assign->targets && !assign->targets->next && is_idempotent(assign->targets->ast)) {
+        if (assign->targets && !assign->targets->next) {
             type_t *lhs_t = get_type(env, assign->targets->ast);
+            if (assign->targets->ast->tag == Index && lhs_t->tag == OptionalType
+                && value_type(get_type(env, Match(assign->targets->ast, Index)->indexed))->tag == TableType)
+                lhs_t = Match(lhs_t, OptionalType)->type;
             env_t *val_env = with_enum_scope(env, lhs_t);
             CORD val = compile_to_type(val_env, assign->values->ast, lhs_t);
             return CORD_all(compile_assignment(env, assign->targets->ast, val), ";\n");
@@ -685,7 +708,8 @@ CORD compile_statement(env_t *env, ast_t *ast)
         int64_t i = 1;
         for (ast_list_t *value = assign->values, *target = assign->targets; value && target; value = value->next, target = target->next) {
             type_t *lhs_t = get_type(env, target->ast);
-            if (target->ast->tag == Index && get_type(env, Match(target->ast, Index)->indexed)->tag == TableType)
+            if (target->ast->tag == Index && lhs_t->tag == OptionalType
+                && value_type(get_type(env, Match(target->ast, Index)->indexed))->tag == TableType)
                 lhs_t = Match(lhs_t, OptionalType)->type;
             env_t *val_env = with_enum_scope(env, lhs_t);
             CORD val = compile_to_type(val_env, value->ast, lhs_t);
@@ -702,7 +726,7 @@ CORD compile_statement(env_t *env, ast_t *ast)
 
         if (update->lhs->tag == Index) {
             type_t *indexed = value_type(get_type(env, Match(update->lhs, Index)->indexed));
-            if (indexed->tag == TableType)
+            if (indexed->tag == TableType && Match(indexed, TableType)->default_value == NULL)
                 code_err(update->lhs, "Update assignments are not currently supported for tables");
         }
 
@@ -1826,7 +1850,35 @@ CORD compile_math_method(env_t *env, binop_e op, ast_t *lhs, ast_t *rhs, type_t
         }
         break;
     }
-    case BINOP_PLUS: case BINOP_MINUS: case BINOP_AND: case BINOP_OR: case BINOP_XOR: {
+    case BINOP_OR: case BINOP_CONCAT: {
+        type_t *lhs_value_t = value_type(lhs_t);
+        arg_t *arg_spec = new(arg_t, .type=lhs_value_t, .next=new(arg_t, .type=lhs_value_t));
+        if (lhs_value_t->tag == SetType) {
+            return CORD_all("Table$with(", compile_arguments(env, lhs, arg_spec, args),
+                            ", ", compile_type_info(env, lhs_value_t), ")");
+        }
+        goto fallthrough;
+    }
+    case BINOP_AND: {
+        type_t *lhs_value_t = value_type(lhs_t);
+        arg_t *arg_spec = new(arg_t, .type=lhs_value_t, .next=new(arg_t, .type=lhs_value_t));
+        if (lhs_value_t->tag == SetType) {
+            return CORD_all("Table$overlap(", compile_arguments(env, lhs, arg_spec, args),
+                            ", ", compile_type_info(env, lhs_value_t), ")");
+        }
+        goto fallthrough;
+    }
+    case BINOP_MINUS: {
+        type_t *lhs_value_t = value_type(lhs_t);
+        arg_t *arg_spec = new(arg_t, .type=lhs_value_t, .next=new(arg_t, .type=lhs_value_t));
+        if (lhs_value_t->tag == SetType) {
+            return CORD_all("Table$without(", compile_arguments(env, lhs, arg_spec, args),
+                            ", ", compile_type_info(env, lhs_value_t), ")");
+        }
+        goto fallthrough;
+    }
+    case BINOP_PLUS: case BINOP_XOR: {
+      fallthrough:
         if (type_eq(lhs_t, rhs_t)) {
             binding_t *b = get_namespace_binding(env, lhs, binop_method_names[op]);
             if (binding_works(b, lhs_t, rhs_t, lhs_t))
@@ -3600,15 +3652,30 @@ CORD compile(env_t *env, ast_t *ast)
                                 CORD_asprintf("%ld", (int64_t)(indexing->index->end - f->text)),
                                 ")");
         } else if (container_t->tag == TableType) {
-            type_t *key_type = Match(container_t, TableType)->key_type;
-            type_t *value_type = Match(container_t, TableType)->value_type;
+            auto table_type = Match(container_t, TableType);
             if (indexing->unchecked)
                 code_err(ast, "Table indexes cannot be unchecked");
-            return CORD_all("({ ", compile_declaration(Type(PointerType, value_type, .is_view=true), "value"),
-                            " = Table$get(", compile_to_pointer_depth(env, indexing->indexed, 0, false), ", ",
-                            compile_to_type(env, indexing->index, Type(PointerType, key_type, .is_view=true)), ", ",
-                            compile_type_info(env, container_t), "); \n"
-                            "value ? ", promote_to_optional(value_type, "*value"), " : ", compile_null(value_type), "; })");
+            if (table_type->default_value) {
+                type_t *value_type = get_type(env, table_type->default_value);
+                return CORD_all("Table$get_or_default(",
+                                compile(env, indexing->indexed), ", ",
+                                compile_type(table_type->key_type), ", ",
+                                compile_type(value_type), ", ",
+                                compile(env, indexing->index), ", ",
+                                compile(env, table_type->default_value), ", ",
+                                compile_type_info(env, container_t), ")");
+            } else if (table_type->value_type) {
+                return CORD_all("Table$get_optional(",
+                                compile(env, indexing->indexed), ", ",
+                                compile_type(table_type->key_type), ", ",
+                                compile_type(table_type->value_type), ", ",
+                                compile(env, indexing->index), ", "
+                                "_, ", promote_to_optional(table_type->value_type, "(*_)"), ", ",
+                                compile_null(table_type->value_type), ", ",
+                                compile_type_info(env, container_t), ")");
+            } else {
+                code_err(indexing->index, "This table doesn't have a value type or a default value");
+            }
         } else {
             code_err(ast, "Indexing is not supported for type: %T", container_t);
         }
diff --git a/parse.c b/parse.c
index 37956055..17ebcea0 100644
--- a/parse.c
+++ b/parse.c
@@ -567,11 +567,18 @@ type_ast_t *parse_table_type(parse_ctx_t *ctx, const char *pos) {
     if (!key_type) return NULL;
     pos = key_type->end;
     whitespace(&pos);
-    if (!match(&pos, ":")) return NULL;
-    type_ast_t *value_type = expect(ctx, start, &pos, parse_type, "I couldn't parse the rest of this table type");
+    type_ast_t *value_type = NULL;
+    ast_t *default_value = NULL;
+    if (match(&pos, ":")) {
+        value_type = expect(ctx, start, &pos, parse_type, "I couldn't parse the rest of this table type");
+    } else if (match(&pos, "=")) {
+        default_value = expect(ctx, start, &pos, parse_extended_expr, "I couldn't parse the rest of this table type");
+    } else {
+        return NULL;
+    }
     whitespace(&pos);
     expect_closing(ctx, &pos, "}", "I wasn't able to parse the rest of this table type");
-    return NewTypeAST(ctx->file, start, pos, TableTypeAST, .key=key_type, .value=value_type);
+    return NewTypeAST(ctx->file, start, pos, TableTypeAST, .key=key_type, .value=value_type, .default_value=default_value);
 }
 
 type_ast_t *parse_set_type(parse_ctx_t *ctx, const char *pos) {
@@ -805,13 +812,18 @@ PARSER(parse_table) {
 
     ast_list_t *entries = NULL;
     type_ast_t *key_type = NULL, *value_type = NULL;
+    ast_t *default_value = NULL;
     if (match(&pos, ":")) {
         whitespace(&pos);
         key_type = expect(ctx, pos-1, &pos, parse_type, "I couldn't parse a key type for this table");
         whitespace(&pos);
-        if (!match(&pos, ":"))
+        if (match(&pos, ":")) {
+            value_type = expect(ctx, pos-1, &pos, parse_type, "I couldn't parse the value type for this table");
+        } else if (match(&pos, "=")) {
+            default_value = expect(ctx, pos-1, &pos, parse_extended_expr, "I couldn't parse the default value for this table");
+        } else {
             return NULL;
-        value_type = expect(ctx, pos-1, &pos, parse_type, "I couldn't parse a value type for this table");
+        }
         whitespace(&pos);
         match(&pos, ",");
     }
@@ -864,7 +876,8 @@ PARSER(parse_table) {
     whitespace(&pos);
     expect_closing(ctx, &pos, "}", "I wasn't able to parse the rest of this table");
 
-    return NewAST(ctx->file, start, pos, Table, .key_type=key_type, .value_type=value_type, .entries=entries, .fallback=fallback);
+    return NewAST(ctx->file, start, pos, Table, .key_type=key_type, .value_type=value_type,
+                  .default_value=default_value, .entries=entries, .fallback=fallback);
 }
 
 PARSER(parse_set) {
diff --git a/stdlib/tables.h b/stdlib/tables.h
index cd19efdd..50ecc453 100644
--- a/stdlib/tables.h
+++ b/stdlib/tables.h
@@ -35,6 +35,14 @@ void *Table$get(Table_t t, const void *key, const TypeInfo_t *type);
     const Table_t t = table_expr; const key_t k = key_expr; \
     val_t *nonnull_var = Table$get(t, &k, info_expr); \
     nonnull_var ? nonnull_expr : null_expr; })
+#define Table$get_or_setdefault(table_expr, key_t, val_t, key_expr, default_expr, info_expr) ({ \
+    Table_t *t = &table_expr; const key_t k = key_expr; \
+    val_t *v = Table$get(*t, &k, info_expr); \
+    v ? v : (val_t*)Table$reserve(t, &k, (val_t[1]){default_expr}, info_expr); })
+#define Table$get_or_default(table_expr, key_t, val_t, key_expr, default_expr, info_expr) ({ \
+    const Table_t t = table_expr; const key_t k = key_expr; \
+    val_t *v = Table$get(t, &k, info_expr); \
+    v ? *v : default_expr; })
 #define Table$has_value(table_expr, key_expr, info_expr) ({ \
     const Table_t t = table_expr; __typeof(key_expr) k = key_expr; \
     (Table$get(t, &k, info_expr) != NULL); })
diff --git a/typecheck.c b/typecheck.c
index 0213dedb..64660f09 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -83,19 +83,25 @@ type_t *parse_type_ast(env_t *env, type_ast_t *ast)
         return Type(ChannelType, .item_type=item_t);
     }
     case TableTypeAST: {
-        type_ast_t *key_type_ast = Match(ast, TableTypeAST)->key;
+        auto table_type = Match(ast, TableTypeAST);
+        type_ast_t *key_type_ast = table_type->key;
         type_t *key_type = parse_type_ast(env, key_type_ast);
         if (!key_type) code_err(key_type_ast, "I can't figure out what type this is.");
         if (has_view_memory(key_type))
             code_err(key_type_ast, "Tables can't have stack references because the array may outlive the stack frame.");
-        type_ast_t *val_type_ast = Match(ast, TableTypeAST)->value;
-        type_t *val_type = parse_type_ast(env, val_type_ast);
-        if (!val_type) code_err(val_type_ast, "I can't figure out what type this is.");
-        if (has_view_memory(val_type))
-            code_err(val_type_ast, "Tables can't have stack references because the array may outlive the stack frame.");
-        else if (val_type->tag == OptionalType)
-            code_err(ast, "Tables with optional-typed values are not currently supported");
-        return Type(TableType, .key_type=key_type, .value_type=val_type);
+        if (table_type->value) {
+            type_t *val_type = parse_type_ast(env, table_type->value);
+            if (!val_type) code_err(table_type->value, "I can't figure out what type this is.");
+            if (has_view_memory(val_type))
+                code_err(table_type->value, "Tables can't have stack references because the array may outlive the stack frame.");
+            else if (val_type->tag == OptionalType)
+                code_err(ast, "Tables with optional-typed values are not currently supported");
+            return Type(TableType, .key_type=key_type, .value_type=val_type);
+        } else if (table_type->default_value) {
+            return Type(TableType, .key_type=key_type, .default_value=table_type->default_value);
+        } else {
+            code_err(ast, "No value type or default value!");
+        }
     }
     case FunctionTypeAST: {
         auto fn = Match(ast, FunctionTypeAST);
@@ -634,6 +640,9 @@ type_t *get_type(env_t *env, ast_t *ast)
         if (table->key_type && table->value_type) {
             key_type = parse_type_ast(env, table->key_type);
             value_type = parse_type_ast(env, table->value_type);
+        } else if (table->key_type && table->default_value) {
+            key_type = parse_type_ast(env, table->key_type);
+            value_type = get_type(env, table->default_value);
         } else {
             for (ast_list_t *entry = table->entries; entry; entry = entry->next) {
                 ast_t *entry_ast = entry->ast;
@@ -666,7 +675,7 @@ type_t *get_type(env_t *env, ast_t *ast)
         }
         if (has_view_memory(key_type) || has_view_memory(value_type))
             code_err(ast, "Tables cannot hold stack references because the table may outlive the reference's stack frame.");
-        return Type(TableType, .key_type=key_type, .value_type=value_type);
+        return Type(TableType, .key_type=key_type, .value_type=value_type, .default_value=table->default_value);
     }
     case TableEntry: {
         code_err(ast, "Table entries should not be typechecked directly");
@@ -720,7 +729,13 @@ type_t *get_type(env_t *env, ast_t *ast)
                 return Match(value_t, ArrayType)->item_type;
             code_err(indexing->index, "I only know how to index lists using integers, not %T", index_t);
         } else if (value_t->tag == TableType) {
-            return Type(OptionalType, Match(value_t, TableType)->value_type);
+            auto table_type = Match(value_t, TableType);
+            if (table_type->default_value)
+                return get_type(env, table_type->default_value);
+            else if (table_type->value_type)
+                return Type(OptionalType, table_type->value_type);
+            else
+                code_err(indexing->indexed, "This type doesn't have a value type or a default value");
         } else {
             code_err(ast, "I don't know how to index %T values", indexed_t);
         }
@@ -1028,10 +1043,10 @@ type_t *get_type(env_t *env, ast_t *ast)
             if (!type_eq(lhs_t, rhs_t))
                 code_err(ast, "The type on the left side of this concatenation doesn't match the right side: %T vs. %T",
                              lhs_t, rhs_t);
-            if (lhs_t->tag == ArrayType || lhs_t->tag == TextType)
+            if (lhs_t->tag == ArrayType || lhs_t->tag == TextType || lhs_t->tag == SetType)
                 return lhs_t;
 
-            code_err(ast, "Only array/text value types support concatenation, not %T", lhs_t);
+            code_err(ast, "Only array/set/text value types support concatenation, not %T", lhs_t);
         }
         case BINOP_EQ: case BINOP_NE: case BINOP_LT: case BINOP_LE: case BINOP_GT: case BINOP_GE: {
             if (!can_promote(lhs_t, rhs_t) && !can_promote(rhs_t, lhs_t))
diff --git a/types.c b/types.c
index 78d52f51..189cec89 100644
--- a/types.c
+++ b/types.c
@@ -43,7 +43,11 @@ CORD type_to_cord(type_t *t) {
         }
         case TableType: {
             auto table = Match(t, TableType);
-            return CORD_asprintf("{%r:%r}", type_to_cord(table->key_type), type_to_cord(table->value_type));
+            if (table->value_type)
+                return CORD_asprintf("{%r:%r}", type_to_cord(table->key_type), type_to_cord(table->value_type));
+            else
+                return CORD_asprintf("{%r=%.*s}", type_to_cord(table->key_type),
+                                     table->default_value->end - table->default_value->start, table->default_value->start);
         }
         case SetType: {
             auto set = Match(t, SetType);
diff --git a/types.h b/types.h
index 9cde2880..c496c391 100644
--- a/types.h
+++ b/types.h
@@ -91,6 +91,7 @@ struct type_s {
         } SetType;
         struct {
             type_t *key_type, *value_type;
+            ast_t *default_value;
         } TableType;
         struct {
             arg_t *args;