fix: Disallow yield in defer + fix inference for generators (#106)

* fix(lang): Disallow yield in a defered statement

* fix inference

src/alumina-boot/src/ast/expressions.rs

@@ -1553,7 +1553,6 @@ impl<'ast, 'src> LambdaVisitor<'ast, 'src> {
             is_local: true,
             is_lambda: true,
             is_protocol_fn: false,
-            is_generator: false,
         }));
 
         let bindings = self

src/alumina-boot/src/ast/maker.rs

@@ -394,11 +394,16 @@ impl<'ast> AstItemMaker<'ast> {
         body: Option<tree_sitter::Node<'src>>,
         attributes: &'ast [Attribute],
     ) -> Result<(), AluminaError> {
+        let mut attributes = attributes.to_vec();
+
         let mut parameters: Vec<Parameter<'ast>> = Vec::new();
         let code = scope.code().unwrap();
 
         let is_extern = node.child_by_field(FieldKind::Extern).is_some();
         let is_generator = node.child_by_field(FieldKind::Generator).is_some();
+        if is_generator {
+            attributes.push(Attribute::Generator);
+        }
 
         let has_varargs = node
             .child_by_field(FieldKind::Parameters)
@@ -516,7 +521,7 @@ impl<'ast> AstItemMaker<'ast> {
 
         let result = Item::Function(Function {
             name,
-            attributes,
+            attributes: attributes.alloc_on(self.ast),
             placeholders,
             args: parameters.alloc_on(self.ast),
             return_type,
@@ -526,7 +531,6 @@ impl<'ast> AstItemMaker<'ast> {
             is_local: self.local,
             is_lambda: false,
             is_protocol_fn,
-            is_generator,
         });
 
         item.assign(result);

src/alumina-boot/src/ast/mod.rs

@@ -740,7 +740,6 @@ pub struct Function<'ast> {
     pub span: Option<Span>,
     pub is_local: bool,
     pub is_lambda: bool,
-    pub is_generator: bool,
     pub varargs: bool,
     pub is_protocol_fn: bool,
 }
@@ -836,6 +835,7 @@ pub enum Attribute<'ast> {
     Intrinsic,
     StaticConstructor,
     LinkName(&'ast str),
+    Generator,
 }
 
 impl<'ast> Attribute<'ast> {
@@ -865,6 +865,7 @@ impl<'ast> Attribute<'ast> {
             Attribute::ConstOnly => Attribute::ConstOnly,
             Attribute::NoConst => Attribute::NoConst,
             Attribute::TupleCall => Attribute::TupleCall,
+            Attribute::Generator => Attribute::Generator,
         }
     }
 }

src/alumina-boot/src/common.rs

@@ -186,6 +186,8 @@ pub enum CodeDiagnostic {
     ConstMessage(String),
     #[error("cannot defer inside a defered expression")]
     DeferInDefer,
+    #[error("cannot yield inside a defered expression")]
+    YieldInDefer,
     #[error("`...` expressions can only be used in macros")]
     EtCeteraOutsideOfMacro,
     #[error("`$` identifiers can only be used in macros")]

src/alumina-boot/src/ir/mono.rs

@@ -1378,7 +1378,8 @@ impl<'a, 'ast, 'ir> Monomorphizer<'a, 'ast, 'ir> {
 
         let return_type = child.lower_type_for_value(func.return_type)?;
         let generator_type = func
-            .is_generator
+            .attributes
+            .contains(&ast::Attribute::Generator)
             .then(|| child.monomorphize_lang_item(LangItemKind::Generator, [return_type]))
             .transpose()?
             .map(|i| child.types.named(i));
@@ -1434,7 +1435,7 @@ impl<'a, 'ast, 'ir> Monomorphizer<'a, 'ast, 'ir> {
             return Ok(());
         }
 
-        if func.is_generator {
+        if generator_type.is_some() {
             let mut grandchild = Self::with_replacements(
                 child.mono_ctx,
                 child.replacements.clone(),
@@ -1610,7 +1611,6 @@ impl<'a, 'ast, 'ir> Monomorphizer<'a, 'ast, 'ir> {
                 is_local: fun.is_local,
                 is_lambda: false,
                 is_protocol_fn: false,
-                is_generator: fun.is_generator,
             }));
 
             result.push(ast::AssociatedFn {
@@ -2871,7 +2871,23 @@ impl<'a, 'ast, 'ir> Monomorphizer<'a, 'ast, 'ir> {
         }
 
         if let Some(return_type_hint) = return_type_hint {
-            infer_pairs.push((fun.return_type, return_type_hint));
+            if fun.attributes.contains(&ast::Attribute::Generator) {
+                let item = self
+                    .mono_ctx
+                    .ast
+                    .lang_item(LangItemKind::Generator)
+                    .with_backtrace(&self.diag)?;
+
+                let wrapped = ast::Ty::Generic(
+                    ast::Ty::Item(item).alloc_on(self.mono_ctx.ast),
+                    [fun.return_type].alloc_on(self.mono_ctx.ast),
+                )
+                .alloc_on(self.mono_ctx.ast);
+
+                infer_pairs.push((wrapped, return_type_hint));
+            } else {
+                infer_pairs.push((fun.return_type, return_type_hint));
+            }
         }
 
         let mut type_inferer = TypeInferer::new(
@@ -5353,6 +5369,10 @@ impl<'a, 'ast, 'ir> Monomorphizer<'a, 'ast, 'ir> {
             bail!(self, CodeDiagnostic::YieldOutsideOfGenerator);
         }
 
+        if self.defer_context.as_ref().is_some_and(|d| d.in_defer) {
+            bail!(self, CodeDiagnostic::YieldInDefer);
+        }
+
         let inner = inner
             .map(|inner| self.lower_expr(inner, self.yield_type))
             .transpose()?

sysroot/std/panicking.alu

@@ -306,7 +306,7 @@ mod internal {
 
                 // minicoro does not play well with unwinding, so if we are inside a coroutine,
                 // we suppress the backtrace to avoid segfaults
-                #[cfg(coroutines)]{}
+                #[cfg(coroutines)]
                 let _ = print_panic_message(&panic_info, !runtime::in_coroutine());
                 #[cfg(not(coroutines))]
                 let _ = print_panic_message(&panic_info, true);

tests/diag/yield_in_defer.alu

@@ -0,0 +1,10 @@
+//! exit_code: 1
+fn* gen() {
+    defer {
+        yield  // diag: error(yield_in_defer): "cannot yield inside a defered expression"
+    }
+}
+
+fn main() {
+    gen();
+}

tests/lang/lang.alu

@@ -896,3 +896,13 @@ fn test_tuple_concat() {
     assert_eq!((1, 2).concat(()), (1, 2));
     assert_eq!(().concat(()), ());
 }
+
+
+#[test]
+fn test_generator_inference() {
+    fn* empty<T>(_: Option<T>) -> T {
+    }
+
+    let _gen: std::runtime::Generator<i32> = empty(Option::none());
+    let _gen = empty(Option::some(12));
+}