feat: add type annotation support for function parameters

- Update fn_param grammar rule to support optional type annotations (name: Type)
- Add alloc_handle and store_type methods to AstHostFunctions trait
- Default to Type::Any for parameters without type annotations (for dynamic typing)
- Type inference will resolve actual types based on use-case before Cranelift

This enables both dynamic typing (fn add(x, y)) and static typing (fn add(x: Int, y: Int))
in the same grammar, supporting gradual typing in ZynML.

Author: darmie

crates/zyn_peg/src/runtime.rs

@@ -368,6 +368,12 @@ pub trait AstHostFunctions {
         None // Default implementation returns None
     }
 
+    /// Allocate a new node handle
+    fn alloc_handle(&mut self) -> NodeHandle;
+
+    /// Store a type associated with a handle (for type resolution)
+    fn store_type(&mut self, handle: NodeHandle, ty: zyntax_typed_ast::type_registry::Type);
+
     /// Create a method call expression
     fn create_method_call(&mut self, receiver: NodeHandle, method: &str, args: Vec<NodeHandle>) -> NodeHandle;
 
@@ -1127,13 +1133,6 @@ impl TypedAstBuilder {
         }
     }
 
-    /// Allocate a new handle ID
-    fn alloc_handle(&mut self) -> NodeHandle {
-        let handle = NodeHandle(self.next_id);
-        self.next_id += 1;
-        handle
-    }
-
     /// Store an expression and return its handle
     fn store_expr(&mut self, expr: TypedNode<TypedExpression>) -> NodeHandle {
         let handle = self.alloc_handle();
@@ -1293,6 +1292,16 @@ impl AstHostFunctions for TypedAstBuilder {
             .unwrap_or_else(|e| format!(r#"{{"declarations": [], "error": "{}"}}"#, e))
     }
 
+    fn alloc_handle(&mut self) -> NodeHandle {
+        let handle = NodeHandle(self.next_id);
+        self.next_id += 1;
+        handle
+    }
+
+    fn store_type(&mut self, handle: NodeHandle, ty: zyntax_typed_ast::type_registry::Type) {
+        self.types.insert(handle, ty);
+    }
+
     fn create_function(
         &mut self,
         name: &str,
@@ -4025,9 +4034,15 @@ impl<'a, H: AstHostFunctions> CommandInterpreter<'a, H> {
                     }
                     _ => "arg".to_string(),
                 };
+                // If type is not provided, default to Any (dynamic box pointer)
                 let ty = match args.get("type") {
                     Some(RuntimeValue::Node(h)) => *h,
-                    _ => self.host.create_primitive_type("i32"),
+                    _ => {
+                        // Create a handle for Type::Any
+                        let handle = self.host.alloc_handle();
+                        self.host.store_type(handle, zyntax_typed_ast::type_registry::Type::Any);
+                        handle
+                    }
                 };
                 let handle = self.host.create_param(&name, ty);
                 Ok(RuntimeValue::Node(handle))

crates/zynml/ml.zyn

@@ -430,7 +430,7 @@ impl_block = { "impl" ~ identifier ~ type_args_angle? ~ "for" ~ identifier ~ "{"
       ]
   }
 
-impl_item = { impl_assoc_type | impl_method }
+impl_item = { impl_assoc_type | impl_method_return | impl_method_params | impl_method_simple }
   -> TypedDeclaration {
       "get_child": { "index": 0 }
   }
@@ -446,9 +446,46 @@ impl_assoc_type = { "type" ~ identifier ~ "=" ~ type_expr }
       ]
   }
 
-impl_method = { fn_def }
+// Impl methods can't have their own type params (they use the impl block's)
+// So we define specific impl method rules without type params
+
+// fn add(self, rhs: T) -> ReturnType { ... }
+impl_method_return = { "fn" ~ identifier ~ "(" ~ fn_params? ~ ")" ~ "->" ~ type_expr ~ block }
   -> TypedDeclaration {
-      "get_child": { "index": 0 }
+      "commands": [
+          { "define": "function", "args": {
+              "name": "$1",
+              "params": "$2",
+              "return_type": "$3",
+              "body": "$4"
+          }}
+      ]
+  }
+
+// fn process(self, other: T) { ... }
+impl_method_params = { "fn" ~ identifier ~ "(" ~ fn_params ~ ")" ~ block }
+  -> TypedDeclaration {
+      "commands": [
+          { "define": "function", "args": {
+              "name": "$1",
+              "params": "$2",
+              "return_type": "void",
+              "body": "$3"
+          }}
+      ]
+  }
+
+// fn simple() { ... }
+impl_method_simple = { "fn" ~ identifier ~ "(" ~ ")" ~ block }
+  -> TypedDeclaration {
+      "commands": [
+          { "define": "function", "args": {
+              "name": "$1",
+              "params": [],
+              "return_type": "void",
+              "body": "$2"
+          }}
+      ]
   }
 
 // Type arguments in angle brackets (for trait impl)
@@ -580,12 +617,16 @@ fn_params = { fn_param ~ ("," ~ fn_param)* }
       "get_all_children": true
   }
 
-// Parameter is just identifier for now (dynamic typing)
-fn_param = { identifier }
+// Parameter with optional type annotation
+// Supports: x, x: Int, self, self: Tensor
+fn_param = { identifier ~ (":" ~ type_expr)? }
   -> TypedParameter {
-      "get_text": true,
-      "define": "param",
-      "args": { "name": "$result", "type": { "define": "primitive_type", "args": { "name": "i64" } } }
+      "commands": [
+          { "define": "param", "args": {
+              "name": "$1",
+              "type": "$2"
+          }}
+      ]
   }
 
 // ============================================================================