pass scope in walk_for_each

lang/src/ast.asm

@@ -1358,25 +1358,38 @@ ast_build_symtable_for_each:
 ;; rsi: start_index
 ;; rdx: ctx
 ;; rcx: for_each
-;; define-fn: fn ast_walk_for_each(ast: *mut Ast, start_index: u64, ctx: *mut (), for_each: unsafe extern "C" fn(ctx: *mut (), *mut Ast, node_index: u64))
+;; define-fn: fn ast_walk_for_each(ast: *mut Ast, start_index: u64, ctx: *mut (), for_each: unsafe extern "C" fn(ctx: *mut (), *mut Ast, node_index: u64, scope: u64))
 ast_walk_for_each:
     push rbp
     push r15
     push r14
     push rbx
 
-    ; current_index [24..32]
+    ; current_index [40..48]
+    ; *current_scope [32..40]
+    ; current_node_ptr [24..32]
     ; for_each [16..24]
     ; ctx [8..16]
     ; ast [0..8]
-    sub rsp, 32
+    sub rsp, 48
     mov [rsp], rdi              ; Ast
     mov [rsp + 8], rdx          ; ctx
     mov [rsp + 16], rcx         ; for_each
-    mov [rsp + 24], rsi         ; current_node_ptr
+    mov qword [rsp + 24], 0     ; current_node_ptr
+    mov [rsp + 40], rsi         ; current_index
 
     mov rbp, rsp
+    lea rdi, [rbp + 32]         ; current_scope
+    push rdi
     push rsi
+    mov [rbp + 32], rsp         ; initialise current_scope pointer
+    push rsi
+
+    ; `current_scope` points to the index of the current scope on the stack.
+    ; When we enter a new scope, we push `current_scope` onto the stack, then
+    ; update it to point to the new scope index.
+    ; When `rsp` is equal to `current_scope`, we need to additionally pop into
+    ; `current_scope` after popping the current index.
 
 .loop:
     cmp rsp, rbp
@@ -1385,6 +1398,8 @@ ast_walk_for_each:
     mov rdi, [rbp + 8]          ; ctx
     mov rsi, [rbp]              ; Ast
     mov rdx, [rsp]              ; current_index
+    mov rcx, [rbp + 32]         ; current_scope
+    mov rcx, [rcx]              ; current_scope value
     mov rax, [rbp + 16]         ; for_each
 
     ; align stack to 16 bytes before call
@@ -1398,6 +1413,8 @@ ast_walk_for_each:
     ; get current_node_ptr
     mov rdi, [rbp]              ; Ast
     pop rsi                     ; current_index
+    mov [rbp + 40], rsi         ; update current_index
+
     call vec_get
     mov [rbp + 24], rax         ; current_node_ptr
     mov bl, byte [rax]          ; AstNode.kind
@@ -1419,9 +1436,14 @@ ast_walk_for_each:
     je .address_of
     cmp bl, AST_RETURN_STATEMENT
     je .return_statement
-    jmp .loop
+    jmp .check_scope
 
 .func:
+    ; push scope
+    push qword [rbp + 32]       ; scope-ptr
+    push qword [rbp + 40]       ; current_index
+    mov [rbp + 32], rsp         ; update current_scope
+    
     ; push child indices to stack
     mov rbx, [rax + 8]          ; AstNode.data
 
@@ -1439,7 +1461,7 @@ ast_walk_for_each:
     inc r14
     jmp .arg_loop
 .arg_loop_done:
-    jmp .loop
+    jmp .check_scope
 
 .block:
     mov rbx, [rax + 8]          ; AstNode.data
@@ -1453,7 +1475,7 @@ ast_walk_for_each:
     push rdx                     ; push statement index
     jmp .stmt_loop
 .stmt_loop_done:
-    jmp .loop
+    jmp .check_scope
 
 .binary_op:
     mov rbx, [rax + 8]          ; AstNode.data
@@ -1461,14 +1483,14 @@ ast_walk_for_each:
     push rdx                     ; push right index
     mov rdx, [rbx + 0]          ; left index
     push rdx                     ; push left index
-    jmp .loop
+    jmp .check_scope
 
 .assignment:
     mov rbx, [rax + 8]          ; AstNode.data = dest
     mov rdx, [rax + 16]         ; AstNode.extra = source
     push rdx                     ; push source index
     push rbx                     ; push dest index
-    jmp .loop
+    jmp .check_scope
 
 .value_to_place:
 .place_to_value:
@@ -1477,11 +1499,35 @@ ast_walk_for_each:
 .return_statement:
     mov rbx, [rax + 8]          ; AstNode.data
     push rbx                     ; push inner expr index
+    jmp .check_scope
+
+.check_scope:
+    cmp rsp, [rbp + 32]         ; current_scope
+    je .pop_scope
     jmp .loop
 
-    ; weird alloca thing
+.pop_scope:
+    ; pop current_scope
+    ; the stack may look something like this:
+    ; current_scope---+
+    ; == stack ==  ^  | points here
+    ; scope-ptr[0]-+  |
+    ; scope: func0 <-----+
+    ; stmt0           |  |
+    ; stmt1           |  |
+    ; scope-ptr[1] ------+
+    ; scope: block0 <-+
+    ; stmt2
+    ; stmt3
+    ; ...
+
+    pop rax                     ; scope
+    pop rax                     ; scope_ptr
+    mov [rbp + 32], rax         ; update current_scope pointer
+    jmp .check_scope
+
 .done:
-    add rsp, 32
+    add rsp, 48
     pop rbx
     pop r14
     pop r15

lang/tests/ast.rs

@@ -29,22 +29,27 @@ fn main() {
             eprintln!("Parsed expression ID: {}", expr_id);
             println!("{:#}", &ast);
 
-            // unsafe extern "C" fn visit_node(_this: *mut (), ast: *mut Ast, node_id: u64) {
+            unsafe extern "C" fn visit_node(
-            //     let ast = unsafe { &*ast };
+                _this: *mut (),
-            //     let node = ast.nodes.get(node_id as usize).unwrap();
+                ast: *mut Ast,
-            //     eprintln!("Visiting node {node_id}: {node}");
+                node_id: u64,
-            // }
+                scope: u64,
+            ) {
+                let ast = unsafe { &*ast };
+                let node = ast.nodes.get(node_id as usize).unwrap();
+                eprintln!("{scope}: Visiting node {node_id}: {node}");
+            }
 
-            // util::defs::ast_walk_for_each(&mut ast, expr_id, core::ptr::null_mut(), visit_node);
+            util::defs::ast_walk_for_each(&mut ast, expr_id, core::ptr::null_mut(), visit_node);
 
-            let mut symtable = core::mem::MaybeUninit::<util::defs::SymbolTable>::uninit();
+            // let mut symtable = core::mem::MaybeUninit::<util::defs::SymbolTable>::uninit();
-            util::defs::ast_build_symtable(&mut ast, expr_id, &mut symtable);
+            // util::defs::ast_build_symtable(&mut ast, expr_id, &mut symtable);
-            let symtable = symtable.assume_init();
+            // let symtable = symtable.assume_init();
-            use util::DisplayedSliceExt;
+            // use util::DisplayedSliceExt;
-            println!(
+            // println!(
-                "Symbol Table: {:#?}",
+            //     "Symbol Table: {:#?}",
-                symtable.symtable.as_slice().displayed()
+            //     symtable.symtable.as_slice().displayed()
-            );
+            // );
         };
     }
 

lang/tests/shared/defs.rs

@@ -15,7 +15,7 @@ unsafe extern "C" {
     pub unsafe fn ast_parse_let(ast: *mut Ast) -> (u64, bool);
     pub unsafe fn symkey_cmp(a: *const SymKey, b: *const SymKey) -> i32;
     pub unsafe fn ast_build_symtable(ast: *mut Ast, root_index: u64, symtable: *mut core::mem::MaybeUninit<SymbolTable>);
-    pub unsafe fn ast_walk_for_each(ast: *mut Ast, start_index: u64, ctx: *mut (), for_each: unsafe extern "C" fn(ctx: *mut (), *mut Ast, node_index: u64));
+    pub unsafe fn ast_walk_for_each(ast: *mut Ast, start_index: u64, ctx: *mut (), for_each: unsafe extern "C" fn(ctx: *mut (), *mut Ast, node_index: u64, scope: u64));
     pub unsafe fn ast_resolve_var_refs(ast: *mut Ast);
 }