completing the generic types example

Signed-off-by: Jakub Doka <jakub.doka2@gmail.com>

lang/README.md

@@ -558,6 +558,69 @@ main := fn(): uint {
 }
 ```
 
+#### generic_types
+```hb
+malloc_sys_call := 69
+free_sys_call := 96
+
+malloc := fn(size: uint, align: uint): ?^void return @eca(malloc_sys_call, size, align)
+free := fn(ptr: ^void, size: uint, align: uint): void return @eca(free_sys_call, ptr, size, align)
+
+Vec := fn($Elem: type): type return struct {
+	data: ^Elem,
+	len: uint,
+	cap: uint,
+
+	new := fn(): Self return .{data: @bitcast(0), len: 0, cap: 0}
+
+	deinit := fn(vec: ^Self): void {
+		free(@bitcast(vec.data), vec.cap * @sizeof(Elem), @alignof(Elem));
+		*vec = new()
+		return
+	}
+
+	push := fn(vec: ^Self, value: Elem): ?^Elem {
+		if vec.len == vec.cap {
+			if vec.cap == 0 {
+				vec.cap = 1
+			} else {
+				vec.cap *= 2
+			}
+
+			new_alloc := @as(?^Elem, @bitcast(malloc(vec.cap * @sizeof(Elem), @alignof(Elem))))
+			if new_alloc == null return null
+
+			src_cursor := vec.data
+			dst_cursor := @as(^Elem, new_alloc)
+			end := vec.data + vec.len
+
+			loop if src_cursor == end break else {
+				*dst_cursor = *src_cursor
+				src_cursor += 1
+				dst_cursor += 1
+			}
+
+			if vec.len != 0 {
+				free(@bitcast(vec.data), vec.len * @sizeof(Elem), @alignof(Elem))
+			}
+			vec.data = new_alloc
+		}
+
+		slot := vec.data + vec.len;
+		*slot = value
+		vec.len += 1
+		return slot
+	}
+}
+
+main := fn(): uint {
+	vec := Vec(uint).new()
+	_ = vec.push(69)
+	defer vec.deinit()
+	return *vec.data
+}
+```
+
 #### die
 ```hb
 main := fn(): never {
@@ -584,6 +647,7 @@ main := fn(): uint {
 }
 ```
 
+
 ### Incomplete Examples
 
 #### comptime_pointers
@@ -599,71 +663,6 @@ modify := fn($num: ^uint): void {
 }
 ```
 
-#### generic_types
-```hb
-malloc_sys_call := 69
-free_sys_call := 96
-
-malloc := fn(size: uint, align: uint): ?^void return @eca(malloc_sys_call, size, align)
-free := fn(ptr: ^void, size: uint, align: uint): void return @eca(free_sys_call, ptr, size, align)
-
-Vec := fn($Elem: type): type {
-	return struct {
-		data: ^Elem,
-		len: uint,
-		cap: uint,
-	}
-}
-
-new := fn($Elem: type): Vec(Elem) return Vec(Elem).{data: @bitcast(0), len: 0, cap: 0}
-
-deinit := fn($Elem: type, vec: ^Vec(Elem)): void {
-	free(@bitcast(vec.data), vec.cap * @sizeof(Elem), @alignof(Elem));
-	*vec = new(Elem)
-	return
-}
-
-push := fn($Elem: type, vec: ^Vec(Elem), value: Elem): ?^Elem {
-	if vec.len == vec.cap {
-		if vec.cap == 0 {
-			vec.cap = 1
-		} else {
-			vec.cap *= 2
-		}
-
-		new_alloc := @as(?^Elem, @bitcast(malloc(vec.cap * @sizeof(Elem), @alignof(Elem))))
-		if new_alloc == null return null
-
-		src_cursor := vec.data
-		dst_cursor := @as(^Elem, new_alloc)
-		end := vec.data + vec.len
-
-		loop if src_cursor == end break else {
-			*dst_cursor = *src_cursor
-			src_cursor += 1
-			dst_cursor += 1
-		}
-
-		if vec.len != 0 {
-			free(@bitcast(vec.data), vec.len * @sizeof(Elem), @alignof(Elem))
-		}
-		vec.data = new_alloc
-	}
-
-	slot := vec.data + vec.len;
-	*slot = value
-	vec.len += 1
-	return slot
-}
-
-main := fn(): uint {
-	vec := new(uint)
-	_f := push(uint, &vec, 69)
-	defer deinit(uint, &vec)
-	res := *vec.data
-	return res
-}
-```
 
 #### fb_driver
 ```hb

lang/src/lib.rs

@@ -356,7 +356,7 @@ pub mod ty {
                 Kind::Struct(s) => {
                     let st = &ctx.structs[s];
                     debug_assert_ne!(st.pos, Pos::MAX);
-                    crate::SymKey::Struct(st.file, st.pos, st.captures)
+                    crate::SymKey::Struct(st.file, st.pos, st.captured)
                 }
                 Kind::Enum(e) => {
                     let en = &ctx.enums[e];
@@ -914,7 +914,7 @@ struct Struct {
     file: Module,
     size: Cell<Size>,
     align: Cell<u8>,
-    captures: ty::Tuple,
+    captured: ty::Tuple,
     explicit_alignment: Option<u8>,
     field_start: u32,
     ast: ExprRef,
@@ -1284,6 +1284,22 @@ impl Types {
             _ => None,
         }
     }
+
+    fn captures_of<'a>(
+        &self,
+        ty: ty::Id,
+        file: &'a parser::Ast,
+    ) -> Option<(&'a [Ident], ty::Tuple)> {
+        match ty.expand() {
+            ty::Kind::Struct(s) => {
+                let Expr::Struct { captured, .. } = self.ins.structs[s].ast.get(file) else {
+                    unreachable!()
+                };
+                Some((captured, self.ins.structs[s].captured))
+            }
+            _ => None,
+        }
+    }
 }
 
 struct OptLayout {

lang/src/son.rs

@@ -3681,6 +3681,11 @@ impl<'a> Codegen<'a> {
                     }
                 }
 
+                if ret.is_some() {
+                    self.gen_defers(defer_base);
+                }
+                self.ci.defers.truncate(defer_base);
+
                 for var in self.ci.scope.vars.drain(base..) {
                     var.remove(&mut self.ci.nodes);
                 }
@@ -3689,9 +3694,6 @@ impl<'a> Codegen<'a> {
                     aclass.remove(&mut self.ci.nodes);
                 }
 
-                self.gen_defers(defer_base);
-                self.ci.defers.truncate(defer_base);
-
                 ret
             }
             Expr::Loop { body, .. } => {
@@ -4266,7 +4268,7 @@ impl<'a> Codegen<'a> {
         )
     }
 
-    fn gen_call(&mut self, func: &Expr, args: &[Expr], inline: bool) -> Option<Value> {
+    fn gen_call(&mut self, func: &Expr, args: &[Expr], mut inline: bool) -> Option<Value> {
         let (ty, mut caller) = match *func {
             Expr::Field { target, pos, name } => {
                 match self.gen_field(Ctx::default(), target, pos, name)? {
@@ -4295,6 +4297,8 @@ impl<'a> Codegen<'a> {
             return Value::NEVER;
         };
 
+        inline |= sig.ret == ty::Id::TYPE;
+
         let Func { expr, file, is_inline, parent, .. } = self.tys.ins.funcs[fu];
         let ast = &self.files[file.index()];
         let &Expr::Closure { args: cargs, body, .. } = expr.get(ast) else { unreachable!() };
@@ -5438,6 +5442,13 @@ impl<'a> Codegen<'a> {
                     break Some(f);
                 }
 
+                if let Some((captures, capture_tuple)) = self.tys.captures_of(piter, f)
+                    && let Some(idx) = captures.iter().position(|&cid| Ok(cid) == id)
+                {
+                    debug_assert_eq!(captures.len(), capture_tuple.len());
+                    return self.tys.ins.args[capture_tuple.range().start + idx];
+                }
+
                 piter = match self.tys.parent_of(piter) {
                     Some(p) => p,
                     None => {
@@ -5591,6 +5602,7 @@ impl<'a> Codegen<'a> {
                     .push(Struct {
                         file: sc.file,
                         pos,
+                        captured,
                         name: sc.name.unwrap_or_default(),
                         field_start: self.tys.ins.struct_fields.len() as _,
                         explicit_alignment: packed.then_some(1),

lang/tests/son_tests_generic_types.txt

@@ -1,6 +1,130 @@
-test.hb:60:22: somehow this was not found
+deinit:
-    defer deinit(uint, &vec)
+    ADDI64 r254, r254, -40d
-                        ^
+    ST r31, r254, 0a, 40h
-test.hb:60:21: expected argument vec to be of type ^[Struct0]{data: ^uint, len: uint, cap: uint}, got ^never
+    CP r32, r2
-    defer deinit(uint, &vec)
+    LD r33, r32, 16a, 8h
-                       ^
+    LI64 r34, 8d
+    MUL64 r33, r33, r34
+    LD r35, r32, 0a, 8h
+    CP r2, r35
+    CP r3, r33
+    CP r4, r34
+    JAL r31, r0, :free
+    CP r1, r32
+    JAL r31, r0, :new
+    LD r31, r254, 0a, 40h
+    ADDI64 r254, r254, 40d
+    JALA r0, r31, 0a
+free:
+    CP r13, r2
+    CP r14, r3
+    CP r15, r4
+    LRA r16, r0, :free_sys_call
+    LD r16, r16, 0a, 8h
+    CP r2, r16
+    CP r3, r13
+    CP r4, r14
+    CP r5, r15
+    ECA 
+    JALA r0, r31, 0a
+main:
+    ADDI64 r254, r254, -56d
+    ST r31, r254, 24a, 32h
+    ADDI64 r32, r254, 0d
+    CP r1, r32
+    JAL r31, r0, :new
+    LI64 r33, 69d
+    CP r2, r32
+    CP r3, r33
+    JAL r31, r0, :push
+    CP r33, r1
+    LD r34, r254, 0a, 8h
+    LD r33, r34, 0a, 8h
+    CP r2, r32
+    JAL r31, r0, :deinit
+    CP r1, r33
+    LD r31, r254, 24a, 32h
+    ADDI64 r254, r254, 56d
+    JALA r0, r31, 0a
+malloc:
+    CP r13, r2
+    CP r14, r3
+    LRA r15, r0, :malloc_sys_call
+    LD r15, r15, 0a, 8h
+    CP r2, r15
+    CP r3, r13
+    CP r4, r14
+    ECA 
+    CP r13, r1
+    CP r1, r13
+    JALA r0, r31, 0a
+new:
+    ADDI64 r254, r254, -24d
+    CP r14, r1
+    ADDI64 r13, r254, 0d
+    ST r0, r254, 0a, 8h
+    ST r0, r254, 8a, 8h
+    ST r0, r254, 16a, 8h
+    BMC r13, r14, 24h
+    ADDI64 r254, r254, 24d
+    JALA r0, r31, 0a
+push:
+    ADDI64 r254, r254, -88d
+    ST r31, r254, 0a, 88h
+    CP r38, r2
+    CP r39, r3
+    LI64 r37, 1d
+    LD r33, r38, 8a, 8h
+    LD r32, r38, 16a, 8h
+    JNE r32, r33, :0
+    JNE r32, r0, :1
+    CP r32, r37
+    JMP :2
+ 1: MULI64 r32, r32, 2d
+ 2: LI64 r40, 8d
+    MUL64 r34, r32, r40
+    CP r2, r34
+    CP r3, r40
+    JAL r31, r0, :malloc
+    CP r35, r1
+    ST r32, r38, 16a, 8h
+    JNE r35, r0, :3
+    CP r1, r0
+    JMP :4
+ 3: MULI64 r33, r33, 8d
+    LD r32, r38, 0a, 8h
+    ADD64 r41, r32, r33
+    CP r34, r35
+ 7: LD r33, r38, 0a, 8h
+    LD r36, r38, 8a, 8h
+    JNE r41, r32, :5
+    JEQ r36, r0, :6
+    MUL64 r32, r36, r40
+    CP r2, r33
+    CP r3, r32
+    CP r4, r40
+    JAL r31, r0, :free
+    JMP :6
+ 6: ST r35, r38, 0a, 8h
+    JMP :0
+ 5: ADDI64 r36, r34, 8d
+    ADDI64 r33, r32, 8d
+    LD r32, r32, 0a, 8h
+    ST r32, r34, 0a, 8h
+    CP r34, r36
+    CP r32, r33
+    JMP :7
+ 0: LD r32, r38, 8a, 8h
+    MULI64 r33, r32, 8d
+    LD r34, r38, 0a, 8h
+    ADD64 r33, r34, r33
+    ST r39, r33, 0a, 8h
+    ADD64 r32, r32, r37
+    ST r32, r38, 8a, 8h
+    CP r1, r33
+ 4: LD r31, r254, 0a, 88h
+    ADDI64 r254, r254, 88d
+    JALA r0, r31, 0a
+code size: 923
+ret: 69
+status: Ok(())