feat(compiler): implement operator overloading with trait dispatch

Add Type::Extern to TypedAST for opaque/FFI types that converts to
HirType::Ptr(Opaque) at HIR level. Implement trait dispatch in SSA
lowering to transform binary ops (a + b) into method calls ($Type$add).

- Add @types directive to ZynPEG for declaring opaque types and returns
- Update ZynML grammar with comprehensive type declarations
- SSA lowering detects Type::Extern and dispatches to trait methods
- Add exhaustive pattern matching for Type::Extern in constraint solver

Author: darmie

crates/compiler/src/ssa.rs

@@ -2956,6 +2956,11 @@ impl SsaBuilder {
                     packed: false,
                 })
             },
+            Type::Extern { name, .. } => {
+                // Extern/Opaque types are pointers to opaque structs at the HIR level
+                // The name is used for trait dispatch (e.g., $Tensor -> $Tensor$add)
+                HirType::Ptr(Box::new(HirType::Opaque(*name)))
+            },
             _ => HirType::I64, // Default for complex types
         }
     }
@@ -3044,7 +3049,8 @@ impl SsaBuilder {
         let type_name = type_name.unwrap();
 
         // Construct the method symbol name: $TypeName$method
-        let method_symbol = format!("${}${}", type_name, method_name);
+        // Note: type_name already includes the $ prefix (e.g., "$Tensor")
+        let method_symbol = format!("{}${}", type_name, method_name);
         log::debug!("[SSA] Operator trait dispatch: {} for type {}", method_symbol, type_name);
 
         // Translate arguments
@@ -3078,9 +3084,33 @@ impl SsaBuilder {
     /// Check if a type should use trait dispatch for operators
     fn is_trait_dispatchable_type(&self, ty: &Type) -> bool {
         match ty {
-            // Named types - check if the HIR conversion results in an opaque type
+            // Extern types - ZRTL-backed opaque types, always use trait dispatch
+            Type::Extern { name, .. } => {
+                let name_str = name.resolve_global()
+                    .unwrap_or_else(|| {
+                        let arena = self.arena.lock().unwrap();
+                        arena.resolve_string(*name)
+                            .map(|s| s.to_string())
+                            .unwrap_or_default()
+                    });
+                log::debug!("[trait_dispatch] Type::Extern name='{}' -> dispatchable", name_str);
+                true
+            }
+            // Named types - check if it's a ZRTL-backed type (name starts with $)
             Type::Named { id, .. } => {
-                // Convert to HIR type and check if it's opaque
+                // Check the type name from registry
+                if let Some(type_def) = self.type_registry.get_type_by_id(*id) {
+                    let name = type_def.name.resolve_global()
+                        .unwrap_or_else(|| {
+                            let arena = self.arena.lock().unwrap();
+                            arena.resolve_string(type_def.name)
+                                .map(|s| s.to_string())
+                                .unwrap_or_default()
+                        });
+                    // ZRTL opaque types start with $
+                    return name.starts_with('$');
+                }
+                // Also check HIR conversion
                 let hir_ty = self.convert_type(ty);
                 matches!(hir_ty, HirType::Opaque(_))
             }
@@ -3091,9 +3121,21 @@ impl SsaBuilder {
         }
     }
 
-    /// Get the symbol prefix for a type (e.g., "Tensor" for $Tensor$add)
+    /// Get the symbol prefix for a type (e.g., "$Tensor" for $Tensor$add)
     fn get_type_symbol_prefix(&self, ty: &Type) -> Option<String> {
         match ty {
+            // Extern types have the name directly
+            Type::Extern { name, .. } => {
+                let name_str = name.resolve_global()
+                    .unwrap_or_else(|| {
+                        let arena = self.arena.lock().unwrap();
+                        arena.resolve_string(*name)
+                            .map(|s| s.to_string())
+                            .unwrap_or_default()
+                    });
+                log::debug!("[trait_dispatch] Type::Extern prefix: '{}'", name_str);
+                Some(name_str)
+            }
             Type::Named { id, .. } => {
                 if let Some(type_def) = self.type_registry.get_type_by_id(*id) {
                     // Use the type name
@@ -3116,8 +3158,7 @@ impl SsaBuilder {
                                     .map(|s| s.to_string())
                                     .unwrap_or_default()
                             });
-                        // Remove $ prefix if present
-                        Some(name_str.trim_start_matches('$').to_string())
+                        Some(name_str)
                     } else {
                         None
                     }

crates/typed_ast/src/constraint_solver.rs

@@ -247,7 +247,8 @@ impl Substitution {
             | Type::Never
             | Type::Any
             | Type::Error
-            | Type::HigherKinded { .. } => ty.clone(),
+            | Type::HigherKinded { .. }
+            | Type::Extern { .. } => ty.clone(),
             Type::Nullable(inner_ty) => {
                 // For nullable types, substitute the inner type and maintain nullability
                 let substituted_inner = self.apply(inner_ty);
@@ -2639,6 +2640,10 @@ impl ConstraintSolver {
                 associated_types,
                 super_traits,
             } => todo!(),
+            Type::Extern { name, .. } => {
+                // Format extern/opaque type by name
+                name.resolve_global().unwrap_or_else(|| format!("extern_{}", name.symbol().to_usize()))
+            }
         }
     }
 
@@ -3394,7 +3399,8 @@ impl ConstraintSolver {
             | Type::Any
             | Type::Error
             | Type::SelfType
-            | Type::HigherKinded { .. } => ty.clone(),
+            | Type::HigherKinded { .. }
+            | Type::Extern { .. } => ty.clone(),
             Type::Nullable(inner_ty) => {
                 // For nullable types, substitute the inner type and maintain nullability
                 let substituted_inner = self.resolve_associated_types(inner_ty, receiver_type);
@@ -3452,6 +3458,7 @@ impl ConstraintSolver {
                 associated_types,
                 super_traits,
             } => todo!(),
+            Type::Extern { .. } => ty.clone(),
         }
     }
 
@@ -3563,7 +3570,8 @@ impl ConstraintSolver {
             | Type::Any
             | Type::Error
             | Type::Associated { .. }
-            | Type::HigherKinded { .. } => ty.clone(),
+            | Type::HigherKinded { .. }
+            | Type::Extern { .. } => ty.clone(),
             Type::Nullable(inner_ty) => {
                 // For nullable types, substitute the inner type and maintain nullability
                 let substituted_inner = self.substitute_self_type(inner_ty, receiver_type);

crates/typed_ast/src/type_registry.rs

@@ -429,6 +429,20 @@ pub enum Type {
         associated_types: Vec<(InternedString, Type)>,
         super_traits: Vec<Type>,
     },
+
+    /// External/Opaque type (for FFI, ZRTL plugins)
+    /// At the HIR level, this becomes a pointer to an opaque type
+    Extern {
+        name: InternedString,
+        layout: Option<ExternLayout>,
+    },
+}
+
+/// Layout information for extern types (optional)
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct ExternLayout {
+    pub size: usize,
+    pub align: usize,
 }
 
 /// Type variable for inference

crates/zyn_peg/src/ast.rs

@@ -1,7 +1,7 @@
 //! AST types for parsed .zyn grammar files
 
 use pest::iterators::Pair;
-use crate::{Rule, ZynGrammar, LanguageInfo, Imports, ContextVar, TypeHelpers, BuiltinMappings, RuleDef, RuleModifier, ActionBlock, ActionField};
+use crate::{Rule, ZynGrammar, LanguageInfo, Imports, ContextVar, TypeHelpers, BuiltinMappings, TypeDeclarations, RuleDef, RuleModifier, ActionBlock, ActionField};
 
 /// Build a ZynGrammar from parsed pest pairs
 pub fn build_grammar(pairs: pest::iterators::Pairs<Rule>) -> Result<ZynGrammar, String> {
@@ -39,6 +39,7 @@ fn process_top_level(grammar: &mut ZynGrammar, pair: Pair<Rule>) -> Result<(), S
         Rule::context_directive => grammar.context = build_context(pair)?,
         Rule::type_helpers_directive => grammar.type_helpers = build_type_helpers(pair)?,
         Rule::builtin_directive => grammar.builtins = build_builtins(pair)?,
+        Rule::types_directive => grammar.types = build_types(pair)?,
         Rule::EOI => {}
         _ => {}
     }
@@ -53,6 +54,7 @@ fn process_directive(grammar: &mut ZynGrammar, pair: Pair<Rule>) -> Result<(), S
             Rule::context_directive => grammar.context = build_context(inner)?,
             Rule::type_helpers_directive => grammar.type_helpers = build_type_helpers(inner)?,
             Rule::builtin_directive => grammar.builtins = build_builtins(inner)?,
+            Rule::types_directive => grammar.types = build_types(inner)?,
             Rule::error_messages_directive => {
                 // TODO: Parse error messages
             }
@@ -205,6 +207,45 @@ fn build_builtins(pair: Pair<Rule>) -> Result<BuiltinMappings, String> {
     Ok(builtins)
 }
 
+/// Parse @types { opaque: [$Tensor, $Audio], returns: { tensor: $Tensor } } directive
+///
+/// Declares:
+/// - opaque: List of opaque type names that are pointer types backed by ZRTL plugins
+/// - returns: Map of function name -> return type for proper type tracking
+fn build_types(pair: Pair<Rule>) -> Result<TypeDeclarations, String> {
+    let mut types = TypeDeclarations::default();
+
+    for inner in pair.into_inner() {
+        if inner.as_rule() == Rule::types_def {
+            let def_text = inner.as_str().trim();
+
+            if def_text.starts_with("opaque") {
+                // Parse opaque type list: opaque: [$Tensor, $Audio]
+                for type_inner in inner.into_inner() {
+                    if type_inner.as_rule() == Rule::type_name {
+                        let type_name = type_inner.as_str().to_string();
+                        types.opaque_types.push(type_name);
+                    }
+                }
+            } else if def_text.starts_with("returns") {
+                // Parse return type mapping: returns: { tensor: $Tensor, audio_load: $Audio }
+                for return_inner in inner.into_inner() {
+                    if return_inner.as_rule() == Rule::return_def {
+                        let mut parts = return_inner.into_inner();
+                        if let (Some(fn_name), Some(type_name)) = (parts.next(), parts.next()) {
+                            let fn_name_str = fn_name.as_str().to_string();
+                            let type_name_str = type_name.as_str().to_string();
+                            types.function_returns.insert(fn_name_str, type_name_str);
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    Ok(types)
+}
+
 fn build_rule_def(pair: Pair<Rule>) -> Result<RuleDef, String> {
     let mut name = String::new();
     let mut modifier = None;

crates/zyn_peg/src/lib.rs

@@ -96,6 +96,21 @@ pub struct BuiltinMappings {
     pub operators: std::collections::HashMap<String, Vec<String>>,
 }
 
+/// Type declarations from @types directive
+/// Declares opaque types and function return types for proper type tracking
+#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
+pub struct TypeDeclarations {
+    /// List of opaque type names (ZRTL-backed types)
+    /// e.g., ["$Tensor", "$Audio"] - these are pointer types backed by plugins
+    #[serde(default)]
+    pub opaque_types: Vec<String>,
+    /// Map of function name -> return type
+    /// e.g., "tensor" -> "$Tensor", "audio_load" -> "$Audio"
+    /// Used during lowering to assign correct types to call expressions
+    #[serde(default)]
+    pub function_returns: std::collections::HashMap<String, String>,
+}
+
 /// A parsed .zyn grammar file
 #[derive(Debug, Clone, Default)]
 pub struct ZynGrammar {
@@ -105,6 +120,8 @@ pub struct ZynGrammar {
     pub type_helpers: TypeHelpers,
     /// Built-in function mappings from @builtin directive
     pub builtins: BuiltinMappings,
+    /// Type declarations from @types directive
+    pub types: TypeDeclarations,
     pub rules: Vec<RuleDef>,
 }
 

crates/zyn_peg/src/runtime.rs

@@ -42,7 +42,7 @@ use std::collections::HashMap;
 use std::path::Path;
 
 use crate::error::{Result, ZynPegError};
-use crate::{ZynGrammar, BuiltinMappings};
+use crate::{ZynGrammar, BuiltinMappings, TypeDeclarations};
 
 // Re-export types from typed_ast for host function implementations
 pub use zyntax_typed_ast::{
@@ -88,6 +88,10 @@ pub struct ZpegMetadata {
     /// - Operators: "$*" -> "vec_dot" (x * y -> vec_dot(x, y))
     #[serde(default)]
     pub builtins: BuiltinMappings,
+    /// Type declarations for opaque types and function return types
+    /// Used for proper type tracking during lowering for operator trait dispatch
+    #[serde(default)]
+    pub types: TypeDeclarations,
 }
 
 /// Commands for a single grammar rule
@@ -302,23 +306,42 @@ pub trait AstHostFunctions {
     /// Create a function call expression
     fn create_call(&mut self, callee: NodeHandle, args: Vec<NodeHandle>) -> NodeHandle;
 
-    /// Create a function call expression with builtin resolution
+    /// Create a function call expression with explicit return type
+    /// Used when we know the return type from @types directive
+    fn create_call_with_return_type(&mut self, callee: NodeHandle, args: Vec<NodeHandle>, return_type: Option<&str>) -> NodeHandle {
+        // Default implementation ignores return type
+        self.create_call(callee, args)
+    }
+
+    /// Create a function call expression with builtin and type resolution
     /// If the callee is a simple identifier matching a builtin function, use the runtime symbol instead
+    /// Also looks up return types from @types directive for proper opaque type tracking
     fn create_call_with_builtin_resolution(
         &mut self,
         callee: NodeHandle,
         args: Vec<NodeHandle>,
         builtins: &BuiltinMappings,
+        types: &TypeDeclarations,
     ) -> NodeHandle {
         // Check if callee is an identifier that matches a builtin function
         if let Some(name) = self.get_identifier_name(callee) {
             log::trace!("[builtin resolution] callee name='{}', checking {} builtins", name, builtins.functions.len());
+
+            // Look up return type from @types directive
+            let return_type = types.function_returns.get(&name).map(|s| s.as_str());
+            if return_type.is_some() {
+                log::trace!("[builtin resolution] found return type for '{}': {:?}", name, return_type);
+            }
+
             if let Some(symbol) = builtins.functions.get(&name) {
                 log::trace!("[builtin resolution] found builtin '{}' -> '{}'", name, symbol);
                 // Create a new identifier with the runtime symbol name
                 let resolved_callee = self.create_identifier(symbol);
-                return self.create_call(resolved_callee, args);
+                return self.create_call_with_return_type(resolved_callee, args, return_type);
             }
+
+            // Not a builtin, but might have return type info
+            return self.create_call_with_return_type(callee, args, return_type);
         } else {
             log::trace!("[builtin resolution] callee is not an identifier");
         }
@@ -543,6 +566,7 @@ impl ZpegCompiler {
             entry_point: grammar.language.entry_point.clone(),
             zpeg_version: env!("CARGO_PKG_VERSION").to_string(),
             builtins: grammar.builtins.clone(),
+            types: grammar.types.clone(),
         };
 
         // Generate pest grammar
@@ -1483,6 +1507,10 @@ impl AstHostFunctions for TypedAstBuilder {
     }
 
     fn create_call(&mut self, callee: NodeHandle, args: Vec<NodeHandle>) -> NodeHandle {
+        self.create_call_with_return_type(callee, args, None)
+    }
+
+    fn create_call_with_return_type(&mut self, callee: NodeHandle, args: Vec<NodeHandle>, return_type: Option<&str>) -> NodeHandle {
         let span = self.default_span();
 
         let callee_expr = self.get_expr(callee)
@@ -1492,7 +1520,35 @@ impl AstHostFunctions for TypedAstBuilder {
             .filter_map(|h| self.get_expr(*h))
             .collect();
 
-        let expr = self.inner.call_positional(callee_expr, arg_exprs, Type::Primitive(PrimitiveType::I32), span);
+        // Convert return type string to Type
+        // Opaque types (starting with $) become Extern types which are pointers at the HIR level
+        let ty = match return_type {
+            Some(type_name) if type_name.starts_with('$') => {
+                // This is an opaque type - create an Extern type
+                // The type name without $ is used as the extern type name
+                log::trace!("[create_call] opaque return type: {}", type_name);
+                Type::Extern {
+                    name: InternedString::new_global(type_name),
+                    layout: None,
+                }
+            }
+            Some(type_name) => {
+                // Regular named type - try to parse it
+                log::trace!("[create_call] named return type: {}", type_name);
+                match type_name {
+                    "i32" | "I32" => Type::Primitive(PrimitiveType::I32),
+                    "i64" | "I64" => Type::Primitive(PrimitiveType::I64),
+                    "f32" | "F32" => Type::Primitive(PrimitiveType::F32),
+                    "f64" | "F64" => Type::Primitive(PrimitiveType::F64),
+                    "bool" | "Bool" => Type::Primitive(PrimitiveType::Bool),
+                    "void" | "Void" | "()" => Type::Primitive(PrimitiveType::Unit),
+                    _ => Type::Primitive(PrimitiveType::I32), // Default for unknown types
+                }
+            }
+            None => Type::Primitive(PrimitiveType::I32), // Default when no return type specified
+        };
+
+        let expr = self.inner.call_positional(callee_expr, arg_exprs, ty, span);
         self.store_expr(expr)
     }
 
@@ -3066,6 +3122,7 @@ impl<'a, H: AstHostFunctions> CommandInterpreter<'a, H> {
                                         base_h.clone(),
                                         call_args,
                                         &self.module.metadata.builtins,
+                                        &self.module.metadata.types,
                                     );
                                     result = RuntimeValue::Node(new_node);
                                 } else if op_info.starts_with("field:") {
@@ -3983,6 +4040,7 @@ impl<'a, H: AstHostFunctions> CommandInterpreter<'a, H> {
                     callee,
                     call_args,
                     &self.module.metadata.builtins,
+                    &self.module.metadata.types,
                 );
                 Ok(RuntimeValue::Node(handle))
             }