1 files changed, 24 insertions, 13 deletions

diff --git a/environment.c b/environment.c
index 1fe698d2..a5d935f3 100644
--- a/environment.c
+++ b/environment.c
@@ -334,22 +334,33 @@ env_t *for_scope(env_t *env, ast_t *ast)
     }
     case FunctionType: case ClosureType: {
         auto fn = iter_t->tag == ClosureType ? Match(Match(iter_t, ClosureType)->fn, FunctionType) : Match(iter_t, FunctionType);
-        arg_t *next_arg = fn->args;
+        if (fn->ret->tag != EnumType)
+            code_err(for_->iter, "Iterator functions must return an enum with a Done and Next field");
+        auto iter_enum = Match(fn->ret, EnumType);
+        type_t *next_type = NULL;
+        for (tag_t *tag = iter_enum->tags; tag; tag = tag->next) {
+            if (streq(tag->name, "Done")) {
+                if (Match(tag->type, StructType)->fields)
+                    code_err(for_->iter, "This iterator function returns an enum with a Done field that has values, when none are allowed");
+            } else if (streq(tag->name, "Next")) {
+                next_type = tag->type;
+            } else {
+                code_err(for_->iter, "This iterator function returns an enum with a value that isn't Done or Next: %s", tag->name);
+            }
+        }
+
+        if (!next_type)
+            code_err(for_->iter, "This iterator function returns an enum that doesn't have a Next field");
+
+        arg_t *iter_field = Match(next_type, StructType)->fields;
         for (ast_list_t *var = for_->vars; var; var = var->next) {
-            if (next_arg == NULL)
-                code_err(var->ast, "This is too many variables for this iterator function");
+            if (!iter_field)
+                code_err(var->ast, "This is one variable too many for this iterator, which returns a %T", fn->ret);
             const char *name = Match(var->ast, Var)->name;
-            type_t *t = get_arg_type(env, next_arg);
-            if (t->tag != PointerType)
-                code_err(for_->iter, "This iterator has type %T, but I need all its arguments to be mutable stack pointers", iter_t);
-            auto ptr = Match(t, PointerType);
-            if (!ptr->is_stack || ptr->is_readonly)
-                code_err(for_->iter, "This iterator has type %T, but I need all its arguments to be mutable stack pointers", iter_t);
-            set_binding(scope, name, new(binding_t, .type=ptr->pointed, .code=CORD_cat("$", name)));
-            next_arg = next_arg->next;
+            type_t *t = get_arg_type(env, iter_field);
+            set_binding(scope, name, new(binding_t, .type=t, .code=CORD_cat("cur.Next.", iter_field->name)));
+            iter_field = iter_field->next;
         }
-        if (next_arg)
-            code_err(ast, "There are not enough variables given for this loop with an iterator that has type %T", iter_t);
         return scope;
     }
     default: code_err(for_->iter, "Iteration is not implemented for type: %T", iter_t);