Move iterator to AST

core_ast.cpp

@@ -378,9 +378,10 @@ ast_var_unpack(Token *pos, Array<Ast_Decl *> vars, Ast_Expr *expr) {
     return result;
 }
 
-//-----------------------------------------------------------------------------
+//
 // Value
-//-----------------------------------------------------------------------------
+//
+
 CORE_Static Value
 value_bool(B32 v) {
     Value value;
@@ -438,3 +439,228 @@ is_atom(Ast *ast) {
     B32 result = is_flag_set(ast->flags, AST_ATOM);
     return result;
 }
+
+//
+// Iterator
+//
+
+const unsigned AST_NODE_END = 128;
+
+struct Ast_Iter {
+    Array<Ast *> stack;
+    Ast *ast;
+    Ast_Kind kind;
+    bool skip_children;
+    uint32_t visit_id;
+    uint32_t di;
+};
+
+uint32_t Ast_Iter_VisitIDGen;
+Ast_Iter iterate_depth_first(Allocator *a, Ast *ast) {
+    assert(ast);
+    Ast_Iter result = {};
+    result.stack = {a};
+    result.ast = ast;
+    result.kind = ast->kind;
+    result.visit_id = ++Ast_Iter_VisitIDGen;
+    return result;
+}
+
+void next(Ast_Iter *iter) {
+    Ast *ast = iter->ast;
+    ast->visit_id = iter->visit_id;
+
+    if (iter->skip_children) {
+        iter->skip_children = false;
+        goto end_of_switch;
+    }
+
+    switch (iter->kind) {
+        case AST_SCOPE: {
+            Ast_Scope *node = (Ast_Scope *)ast;
+            iter->stack.add(node);
+            For(node->stmts) iter->stack.add(it);
+            For(node->decls) iter->stack.add(it);
+        } break;
+
+        case AST_MODULE: break; // This happens when we import stuff
+        case AST_FILE: invalid_codepath; break;
+
+        case AST_IDENT:
+        case AST_VALUE: {
+            Ast_Atom *node = (Ast_Atom *)ast;
+        } break;
+
+        case AST_INDEX: {
+            Ast_Index *node = (Ast_Index *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->index);
+            iter->stack.add(node->expr);
+            assert(node->index);
+            assert(node->expr);
+        } break;
+
+        case AST_UNARY: {
+            Ast_Unary *node = (Ast_Unary *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->expr);
+            assert(node->expr);
+        } break;
+
+        case AST_BINARY: {
+            Ast_Binary *node = (Ast_Binary *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->right);
+            iter->stack.add(node->left);
+            assert(node->right);
+            assert(node->left);
+        } break;
+
+        case AST_CALL_ITEM: {
+            Ast_Call_Item *node = (Ast_Call_Item *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->item);
+            assert(node->item);
+            if (node->call_flags & CALL_INDEX) {
+                iter->stack.add(node->index);
+                assert(node->index);
+            }
+            else if (node->call_flags & CALL_NAME) {
+                iter->stack.add(node->name);
+                assert(node->name);
+            }
+        } break;
+
+        case AST_COMPOUND:
+        case AST_CALL: {
+            Ast_Call *node = (Ast_Call *)ast;
+            iter->stack.add(node);
+            For(node->exprs) iter->stack.add(it);
+            if (node->name) iter->stack.add(node->name);
+        } break;
+
+        case AST_TYPE_OF:
+        case AST_LENGTH_OF:
+        case AST_ALIGN_OF:
+        case AST_SIZE_OF:
+        case AST_RUNTIME_ASSERT:
+        case AST_CONSTANT_ASSERT: {
+            Ast_Builtin *node = (Ast_Builtin *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->expr);
+            assert(node->expr);
+        } break;
+
+        case AST_SWITCH: {
+            Ast_Switch *node = (Ast_Switch *)ast;
+            iter->stack.add(node);
+            if (node->default_scope) iter->stack.add(node->default_scope);
+            For(node->cases) iter->stack.add(it);
+            iter->stack.add(node->value);
+            assert(node->value);
+        } break;
+
+        case AST_SWITCH_CASE: {
+            Ast_Switch_Case *node = (Ast_Switch_Case *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->scope);
+            assert(node->scope);
+            For(node->labels) iter->stack.add(it);
+        } break;
+
+        case AST_VAR_UNPACK: {
+            Ast_Var_Unpack *node = (Ast_Var_Unpack *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->expr);
+            assert(node->expr);
+            For(node->vars) iter->stack.add(it);
+        } break;
+
+        case AST_PASS:
+        case AST_BREAK: {
+            Ast *node = (Ast *)ast;
+        } break;
+
+        case AST_NAMESPACE:
+        case AST_STRUCT:
+        case AST_UNION:
+        case AST_ENUM:
+        case AST_LAMBDA:
+        case AST_TYPE: // @cleanup: what is this used for?
+        case AST_CONST:
+        case AST_VAR: {
+            Ast_Decl *node = (Ast_Decl *)ast;
+
+            iter->stack.add(node);
+            if (node->scope) iter->stack.add(node->scope);
+            if (node->expr) iter->stack.add(node->expr);
+            if (node->typespec) iter->stack.add(node->typespec);
+
+            // omitting polymorphs
+            // omitting polymorph parameters
+        } break;
+
+        case AST_ARRAY: {
+            Ast_Array *node = (Ast_Array *)ast;
+            iter->stack.add(node);
+            if (node->expr) iter->stack.add(node->expr);
+        } break;
+
+        case AST_FOR: {
+            Ast_For *node = (Ast_For *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->scope);
+            assert(node->scope);
+            if (node->iter) iter->stack.add(node->iter);
+            if (node->cond) iter->stack.add(node->cond);
+            if (node->init) iter->stack.add(node->init);
+        } break;
+
+        case AST_IF: {
+            Ast_If *node = (Ast_If *)ast;
+            iter->stack.add(node);
+            For(node->ifs) iter->stack.add(it);
+        } break;
+
+        case AST_IF_NODE: {
+            Ast_If_Node *node = (Ast_If_Node *)ast;
+            iter->stack.add(node);
+            iter->stack.add(node->scope);
+            assert(node->scope);
+            if (node->expr) iter->stack.add(node->expr);
+            if (node->init) iter->stack.add(node->init);
+        } break;
+
+        case AST_RETURN: {
+            Ast_Return *node = (Ast_Return *)ast;
+            iter->stack.add(node);
+            For(node->expr) iter->stack.add(it);
+        } break;
+
+        case AST_LAMBDA_EXPR: {
+            Ast_Lambda *node = (Ast_Lambda *)ast;
+            iter->stack.add(node);
+            if (node->scope) iter->stack.add(node->scope);
+            For(node->ret) iter->stack.add(it);
+            For(node->args) iter->stack.add(it);
+        } break;
+
+        default: assert(!"Invalid default case");
+    }
+end_of_switch:
+
+    if (iter->stack.len <= 0) {
+        iter->ast = 0;
+        iter->kind = AST_NONE;
+        iter->stack.dealloc();
+        return;
+    }
+
+    iter->ast = iter->stack.pop();
+    assert(iter->ast != 0);
+    iter->di += 1;
+    iter->kind = iter->ast->kind;
+    if (iter->ast->visit_id == iter->visit_id) {
+        iter->kind = (Ast_Kind)((unsigned)iter->kind + AST_NODE_END);
+    }
+}