feat: support top-level type annotations

examples/abs.zy

@@ -4,6 +4,4 @@ fn abs(x: i32): i32 {
     return if x < 0 then -x else x
 }
 
-fn main(): void {
-    std.debug.print(abs(-42))
-}
+std.debug.print(abs(-42))

examples/explicit-main.zy

@@ -0,0 +1,7 @@
+const std = import("std")
+
+fn main(): void {
+    std.debug.print("hello from main")
+}
+
+main()

examples/get-readme.zy

@@ -6,10 +6,8 @@ fn getReadme(filename: string): !string {
     return data
 }
 
-fn main(): void {
-    const data = getReadme("README.md") catch err {
-        std.debug.print(err)
-        return
-    }
-    std.debug.print(data)
+const data = getReadme("README.md") catch err {
+    std.debug.print(err)
+    return
 }
+std.debug.print(data)

examples/switch.zy

@@ -9,9 +9,7 @@ fn label(n: i32): string {
     }
 }
 
-fn main(): void {
-    std.debug.print(label(1))
-    std.debug.print(label(2))
-    std.debug.print(label(3))
-    std.debug.print(label(99))
-}
+std.debug.print(label(1))
+std.debug.print(label(2))
+std.debug.print(label(3))
+std.debug.print(label(99))

src/codegen/zig/mod.rs

@@ -2,6 +2,8 @@ use crate::codegen::Backend;
 use crate::parser::*;
 use std::collections::{HashMap, HashSet};
 
+const MAIN_MANGLED: &str = "__zyre_fn_main";
+
 mod stdlib;
 mod tracker;
 
@@ -165,15 +167,13 @@ impl ZigBackend {
         // Hoisting: collect export const and their dependencies
         let hoisted = Self::collect_hoisted(program);
 
-        let has_explicit_main = program
-            .iter()
-            .any(|item| matches!(item, TopLevel::FnDecl(f) if f.name == "main"));
-
         // Pass 2: collect body_stmts first (needed for header generation)
         let mut body_stmts: Vec<Stmt> = Vec::new();
         for item in program {
             match item {
-                TopLevel::ConstDecl { name, value, .. } => {
+                TopLevel::ConstDecl {
+                    name, ty, value, ..
+                } => {
                     let is_import = matches!(&value.kind, ExprKind::Import(_));
                     let is_module_alias = if let ExprKind::MemberAccess { obj, .. } = &value.kind {
                         matches!(&obj.kind, ExprKind::Var(m) if import_names.contains(m) || self.aliases.contains_key(m))
@@ -189,7 +189,7 @@ impl ZigBackend {
                         body_stmts.push(Stmt {
                             kind: StmtKind::ConstDecl {
                                 name: name.clone(),
-                                ty: None,
+                                ty: ty.clone(),
                                 value: value.clone(),
                             },
                             span: (0, 0),
@@ -201,9 +201,8 @@ impl ZigBackend {
             }
         }
 
-        // main / implicit main uses std.heap, so std is required
-        let needs_std =
-            self.program_uses_std(program) || has_explicit_main || !body_stmts.is_empty();
+        // implicit main uses std.heap, so std is required
+        let needs_std = self.program_uses_std(program) || !body_stmts.is_empty();
 
         let mut out = String::new();
         if needs_std {
@@ -289,7 +288,7 @@ impl ZigBackend {
             }
         }
 
-        if !body_stmts.is_empty() && !has_explicit_main {
+        if !body_stmts.is_empty() {
             let needs_alloc = self.uses_allocator(&body_stmts);
             out.push_str("pub fn main() !void {\n");
             out.push_str("    try __zyre_runtime.Output.init();\n");
@@ -323,40 +322,41 @@ impl ZigBackend {
         }
     }
 
+    fn is_allocating_call(&self, callee: &Expr) -> bool {
+        if let ExprKind::Var(name) = &callee.kind {
+            self.allocating_fns.contains(name)
+        } else {
+            false
+        }
+    }
+
     fn gen_fn(&mut self, f: &FnDecl) -> String {
         let params: Vec<String> = f
             .params
             .iter()
             .map(|(name, ty)| format!("{}: {}", name, self.gen_type(ty)))
             .collect();
 
-        let (ret, pub_prefix) = if f.name == "main" {
-            ("!void".to_string(), "pub ")
-        } else if f.exported {
-            (self.gen_type(&f.ret), "pub ")
-        } else {
-            (self.gen_type(&f.ret), "")
-        };
+        let pub_prefix = if f.exported { "pub " } else { "" };
+        let ret = self.gen_type(&f.ret);
 
         let needs_alloc = self.allocating_fns.contains(&f.name);
 
-        let params_str = if f.name == "main" {
-            "".to_string()
-        } else if needs_alloc {
+        let params_str = if needs_alloc {
             let mut all = vec!["__zyre_allocator: std.mem.Allocator".to_string()];
             all.extend(params);
             all.join(", ")
         } else {
             params.join(", ")
         };
 
-        let mut out = format!("{}fn {}({}) {} {{\n", pub_prefix, f.name, params_str, ret);
-
-        if f.name == "main" {
-            out.push_str("    try __zyre_runtime.Output.init();\n");
-            out.push_str("    defer __zyre_runtime.Output.restore();\n");
-            out.push_str(&Self::gen_arena_setup(needs_alloc));
-        }
+        // "main" is reserved for the implicit entry point in Zig; mangle user-defined fn main
+        let zig_name = if f.name == "main" {
+            MAIN_MANGLED
+        } else {
+            &f.name
+        };
+        let mut out = format!("{}fn {}({}) {} {{\n", pub_prefix, zig_name, params_str, ret);
 
         for stmt in &f.body {
             out.push_str(&self.gen_stmt(stmt, 1));
@@ -508,13 +508,16 @@ impl ZigBackend {
                         }
                     }
                 }
-                let callee_s = self.gen_expr(callee);
+                // Mangle calls to user-defined "main" to match the renamed fn
+                let callee_s = if matches!(&callee.kind, ExprKind::Var(n) if n == "main") {
+                    MAIN_MANGLED.to_string()
+                } else {
+                    self.gen_expr(callee)
+                };
                 let mut args_str = self.gen_args(args);
                 // Insert __zyre_allocator as the first argument when calling an allocating fn
-                if let ExprKind::Var(name) = &callee.kind {
-                    if self.allocating_fns.contains(name) {
-                        args_str.insert(0, "__zyre_allocator".to_string());
-                    }
+                if self.is_allocating_call(callee) {
+                    args_str.insert(0, "__zyre_allocator".to_string());
                 }
                 format!("{}({})", callee_s, args_str.join(", "))
             }

src/codegen/zig/stdlib/mod.rs

@@ -80,6 +80,10 @@ impl ZigBackend {
         // Check recursively
         match &expr.kind {
             ExprKind::Call { callee, args } => {
+                // Calling a user function that transitively needs an allocator
+                if self.is_allocating_call(callee) {
+                    return true;
+                }
                 self.expr_uses_allocator(callee) || args.iter().any(|a| self.expr_uses_allocator(a))
             }
             ExprKind::MemberAccess { obj, .. } => self.expr_uses_allocator(obj),

src/fmt.rs

@@ -61,13 +61,18 @@ impl Formatter {
         match item {
             TopLevel::ConstDecl {
                 name,
+                ty,
                 value,
                 exported,
                 ..
             } => {
                 let prefix = if *exported { "export " } else { "" };
+                let type_ann = ty
+                    .as_ref()
+                    .map(|t| format!(": {}", fmt_type(t)))
+                    .unwrap_or_default();
                 let val = self.fmt_expr(value);
-                self.line(&format!("{}const {} = {}", prefix, name, val));
+                self.line(&format!("{}const {}{} = {}", prefix, name, type_ann, val));
             }
             TopLevel::FnDecl(f) => self.fmt_fn(f),
             TopLevel::StructDecl {

src/parser.rs

@@ -190,6 +190,7 @@ pub enum TopLevel {
     ConstDecl {
         name: String,
         name_span: Span,
+        ty: Option<TypeExpr>,
         value: Expr,
         exported: bool,
     },
@@ -394,6 +395,12 @@ impl Parser {
             Token::Const => {
                 self.advance();
                 let (name, name_span) = self.expect_ident_spanned();
+                let ty = if self.peek() == &Token::Colon {
+                    self.advance();
+                    Some(self.parse_type())
+                } else {
+                    None
+                };
                 self.expect(Token::Eq);
                 match self.peek().clone() {
                     Token::Import => {
@@ -402,6 +409,7 @@ impl Parser {
                         TopLevel::ConstDecl {
                             name,
                             name_span,
+                            ty,
                             value,
                             exported,
                         }
@@ -454,6 +462,7 @@ impl Parser {
                         TopLevel::ConstDecl {
                             name,
                             name_span,
+                            ty,
                             value,
                             exported,
                         }

src/tests/fixtures/alloc_propagate.zy

@@ -4,9 +4,7 @@ fn readIt(): !string {
     return std.fs.readTextFile("x.txt")?
 }
 
-fn main(): void {
-    const data = readIt() catch _err {
-        return
-    }
-    std.debug.print(data)
+const data = readIt() catch _err {
+    return
 }
+std.debug.print(data)

src/tests/fixtures/array.zy

@@ -1,6 +1,4 @@
 const std = import("std")
 
-fn main(): void {
-    const arr: i32[3] = [1, 2, 3]
-    std.debug.print(arr[0])
-}
+const arr: i32[3] = [1, 2, 3]
+std.debug.print(arr[0])

src/tests/fixtures/enum_decl.zy

@@ -16,6 +16,4 @@ fn describe(d: Direction): void {
     }
 }
 
-fn main(): void {
-    std.debug.print("ok")
-}
+std.debug.print("ok")

src/tests/fixtures/explicit_main.zy

@@ -0,0 +1,7 @@
+const std = import("std")
+
+fn main(): void {
+    std.debug.print("hello from main")
+}
+
+main()

src/tests/fixtures/export_const.zy

@@ -4,6 +4,4 @@ const a = 2
 const b = 3
 export const c = a * b
 
-fn main(): void {
-    std.debug.print(c)
-}
+std.debug.print(c)

src/tests/fixtures/fn_call.zy

@@ -4,6 +4,4 @@ fn add(a: i32, b: i32): i32 {
     return a + b
 }
 
-fn main(): void {
-    std.debug.print(add(1, 2))
-}
+std.debug.print(add(1, 2))

src/tests/fixtures/if_stmt.zy

@@ -1,8 +1,6 @@
 const std = import("std")
 
-fn main(): void {
-    const x = true
-    if x {
-        std.debug.print("yes")
-    }
+const x = true
+if x {
+    std.debug.print("yes")
 }

src/tests/fixtures/print_int.zy

@@ -1,6 +1,4 @@
 const std = import("std")
 
-fn main(): void {
-    const x: i32 = 42
-    std.debug.print(x)
-}
+const x: i32 = 42
+std.debug.print(x)

src/tests/fixtures/print_str.zy

@@ -1,5 +1,3 @@
 const std = import("std")
 
-fn main(): void {
-    std.debug.print("hello")
-}
+std.debug.print("hello")

src/tests/fixtures/read_text_file.zy

@@ -1,8 +1,6 @@
 const std = import("std")
 
-fn main(): void {
-    const data = std.fs.readTextFile("file.txt") catch _err {
-        return
-    }
-    std.debug.print(data)
+const data = std.fs.readTextFile("file.txt") catch _err {
+    return
 }
+std.debug.print(data)

src/tests/fixtures/struct_decl.zy

@@ -9,6 +9,4 @@ fn getX(p: Point): i32 {
     return p.x
 }
 
-fn main(): void {
-    std.debug.print("ok")
-}
+std.debug.print("ok")

src/tests/fixtures/switch_stmt.zy

@@ -8,6 +8,4 @@ fn check(n: i32): void {
     }
 }
 
-fn main(): void {
-    check(1)
-}
+check(1)

src/tests/fixtures/while_loop.zy

@@ -1,9 +1,7 @@
 const std = import("std")
 
-fn main(): void {
-    const x = true
-    while x {
-        std.debug.print("loop")
-        break
-    }
+const x = true
+while x {
+    std.debug.print("loop")
+    break
 }