Add field accessor for typechecking enum fields

1 files changed, 40 insertions, 5 deletions

diff --git a/src/main.c b/src/main.c
index f2c4d4f..d493d8e 100644
--- a/src/main.c
+++ b/src/main.c
@@ -8,7 +8,7 @@
 #include "vm.c"
 
 // TODO: typecheck structs
-// TODO: add field accessor for struct/enum symbols
+// TODO: add field accessor for struct symbols
 
 typedef enum ExecMode {
     RUN_NORMAL,
@@ -162,6 +162,18 @@ find_fun(TypeScope *scope, Str type) {
     return NULL;
 }
 
+EnumMap *
+find_enum(TypeScope *scope, Str type) {
+    while (scope != NULL) {
+        EnumMap *val = enummap_lookup(&scope->enums, type);
+        if (val != NULL) {
+            return val;
+        }
+        scope = scope->parent;
+    }
+    return NULL;
+}
+
 void
 graph_scope(Scope *scope, Arena a) {
     if (!scope->symbols) {
@@ -507,6 +519,7 @@ type_inference(Analyzer *a, Node *node, TypeScope *scope) {
                     a->storage);
                 field->type = cstr("int");
             }
+            typemap_insert(&scope->types, symbol, symbol, a->storage);
             return node->type;
         } break;
         case NODE_IF: {
@@ -716,11 +729,33 @@ type_inference(Analyzer *a, Node *node, TypeScope *scope) {
             return node->type;
         } break;
         case NODE_SYMBOL: {
-            TypeMap *type = find_type(scope, node->value.str);
-            if (type) {
-                node->type = type->val;
-                return node->type;
+            Str symbol = node->value.str;
+            TypeMap *type = find_type(scope, symbol);
+            if (!type) {
+                eprintln("%s:%d:%d: error: couldn't resolve symbol '%s'",
+                         a->file_name, node->line, node->col, symbol);
+                return cstr("");
             }
+            EnumMap *e = find_enum(scope, type->val);
+            if (e && str_eq(symbol, type->val)) {
+                if (!node->next) {
+                    eprintln(
+                        "%s:%d:%d: error: unspecified enum field for symbol "
+                        "'%s'",
+                        a->file_name, node->line, node->col, symbol);
+                    return cstr("");
+                }
+                // Check if there is a next and it matches the enum field.
+                if (!fieldmap_lookup(&e->val.fields, node->next->value.str)) {
+                    eprintln(
+                        "%s:%d:%d: error: unknown enum field for symbol "
+                        "'%s'",
+                        a->file_name, node->line, node->col, symbol);
+                    return cstr("");
+                }
+                node->next->type = symbol;
+            }
+            node->type = type->val;
             return node->type;
         } break;
         case NODE_FUNCALL: {