Preserve tensor type metadata and alias targets in Grammar2

darmie · darmie · commit cd79201be98f · 2026-02-26T21:33:16.000Z
diff --git a/crates/zyn_peg/src/runtime2/interpreter.rs b/crates/zyn_peg/src/runtime2/interpreter.rs
@@ -18,7 +18,8 @@ use zyntax_typed_ast::typed_ast::{
 };
 use zyntax_typed_ast::{
     type_registry::{
-        CallingConvention, ConstValue, Mutability, NullabilityKind, PrimitiveType, Type, Visibility,
+        CallingConvention, ConstValue, Mutability, NullabilityKind, PrimitiveType, Type,
+        TypeMetadata, Visibility,
     },
     typed_node, ParameterKind, Span, TypedAnnotation, TypedAnnotationArg, TypedAnnotationValue,
     TypedBlock, TypedCall, TypedDeclaration, TypedExpression, TypedExtern, TypedExternStruct,
@@ -1147,12 +1148,14 @@ impl<'g> GrammarInterpreter<'g> {
         span: Span,
     ) -> Result<ParsedValue, String> {
         let declarations = self.get_field_as_decl_list("declarations", fields, state)?;
+        // Preserve any placeholder or inferred type entries created while parsing.
+        let type_registry = state.type_registry().clone();
 
         Ok(ParsedValue::Program(Box::new(TypedProgram {
             declarations,
             span,
             source_files: vec![],
-            type_registry: zyntax_typed_ast::type_registry::TypeRegistry::new(),
+            type_registry,
         })))
     }
 
@@ -1162,13 +1165,68 @@ impl<'g> GrammarInterpreter<'g> {
         variant: &str,
         fields: &[(String, ExprIR)],
         state: &mut ParserState<'a>,
-        _span: Span,
+        span: Span,
     ) -> Result<ParsedValue, String> {
         let ty = match variant {
             "Unit" => Type::Primitive(PrimitiveType::Unit),
             "Named" => {
                 let name = self.get_field_as_interned("name", fields, state)?;
-                Type::Unresolved(name)
+                let mut type_args = self
+                    .get_field_as_type_list_optional("type_args", fields, state)?
+                    .unwrap_or_default();
+                let mut const_args = self.get_field_as_const_list("const_args", fields, state)?;
+
+                // Tensor shape/dtype sugar from tensor[...] syntax.
+                let tensor_items =
+                    self.get_field_as_interned_list("tensor_items", fields, state)?;
+                if !tensor_items.is_empty() {
+                    let mut tokens: Vec<String> = tensor_items
+                        .into_iter()
+                        .map(|item| {
+                            item.resolve_global().or_else(|| {
+                                state
+                                    .builder()
+                                    .arena()
+                                    .resolve_string(item)
+                                    .map(|s| s.to_string())
+                            })
+                        })
+                        .collect::<Option<Vec<String>>>()
+                        .unwrap_or_default();
+
+                    if let Some(last) = tokens.last() {
+                        if let Some(dtype_prim) = self.primitive_type_from_name(last) {
+                            type_args.push(Type::Primitive(dtype_prim));
+                            tokens.pop();
+                        }
+                    }
+
+                    for token in tokens {
+                        if let Ok(v) = token.parse::<i64>() {
+                            const_args.push(ConstValue::Int(v));
+                        } else {
+                            const_args.push(ConstValue::Variable(
+                                zyntax_typed_ast::InternedString::new_global(&token),
+                            ));
+                        }
+                    }
+                }
+
+                let type_id = if let Some(type_def) = state.type_registry().get_type_by_name(name) {
+                    type_def.id
+                } else {
+                    state
+                        .type_registry()
+                        .register_atomic_type(name, TypeMetadata::default(), span)
+                };
+
+                Type::Named {
+                    id: type_id,
+                    type_args,
+                    const_args,
+                    variance: vec![],
+                    nullability: NullabilityKind::NonNull,
+                }
             }
             "Primitive" => {
                 // Parse primitive type from name
@@ -1180,21 +1238,9 @@ impl<'g> GrammarInterpreter<'g> {
                 if name_str == "type" {
                     return Ok(ParsedValue::Type(Type::Any));
                 }
-                let prim = match name_str.as_str() {
-                    "i8" => PrimitiveType::I8,
-                    "i16" => PrimitiveType::I16,
-                    "i32" => PrimitiveType::I32,
-                    "i64" => PrimitiveType::I64,
-                    "u8" => PrimitiveType::U8,
-                    "u16" => PrimitiveType::U16,
-                    "u32" => PrimitiveType::U32,
-                    "u64" => PrimitiveType::U64,
-                    "f32" => PrimitiveType::F32,
-                    "f64" => PrimitiveType::F64,
-                    "bool" => PrimitiveType::Bool,
-                    "void" => PrimitiveType::Unit,
-                    _ => return Err(format!("unknown primitive type: {}", name_str)),
-                };
+                let prim = self
+                    .primitive_type_from_name(&name_str)
+                    .ok_or_else(|| format!("unknown primitive type: {}", name_str))?;
                 Type::Primitive(prim)
             }
             "Pointer" => {
@@ -1834,6 +1880,83 @@ impl<'g> GrammarInterpreter<'g> {
         }
     }
 
+    fn get_field_as_const_list<'a>(
+        &self,
+        name: &str,
+        fields: &[(String, ExprIR)],
+        state: &mut ParserState<'a>,
+    ) -> Result<Vec<ConstValue>, String> {
+        let Some(expr) = self.get_field(name, fields) else {
+            return Ok(vec![]);
+        };
+
+        let val = self.eval_expr(expr, state)?;
+        match val {
+            ParsedValue::List(items) => items
+                .into_iter()
+                .map(|item| self.parsed_value_to_const(item, state))
+                .collect(),
+            ParsedValue::Optional(None) | ParsedValue::None => Ok(vec![]),
+            ParsedValue::Optional(Some(inner)) => match *inner {
+                ParsedValue::List(items) => items
+                    .into_iter()
+                    .map(|item| self.parsed_value_to_const(item, state))
+                    .collect(),
+                single => Ok(vec![self.parsed_value_to_const(single, state)?]),
+            },
+            single => Ok(vec![self.parsed_value_to_const(single, state)?]),
+        }
+    }
+
+    fn parsed_value_to_const<'a>(
+        &self,
+        value: ParsedValue,
+        state: &mut ParserState<'a>,
+    ) -> Result<ConstValue, String> {
+        match value {
+            ParsedValue::Int(i) => Ok(ConstValue::Int(i)),
+            ParsedValue::Bool(b) => Ok(ConstValue::Bool(b)),
+            ParsedValue::Text(s) => Ok(ConstValue::String(state.intern(&s))),
+            ParsedValue::Interned(s) => Ok(ConstValue::Variable(s)),
+            ParsedValue::Literal(TypedLiteral::Integer(i)) => Ok(ConstValue::Int(i as i64)),
+            ParsedValue::Literal(TypedLiteral::Bool(b)) => Ok(ConstValue::Bool(b)),
+            ParsedValue::Literal(TypedLiteral::String(s)) => Ok(ConstValue::String(s)),
+            ParsedValue::Literal(TypedLiteral::Char(c)) => Ok(ConstValue::Char(c)),
+            ParsedValue::Optional(Some(inner)) => self.parsed_value_to_const(*inner, state),
+            ParsedValue::Optional(None) | ParsedValue::None => {
+                Err("cannot convert empty optional to const value".to_string())
+            }
+            other => Err(format!(
+                "cannot convert value to const argument: {:?}",
+                other
+            )),
+        }
+    }
+
+    fn primitive_type_from_name(&self, name: &str) -> Option<PrimitiveType> {
+        match name {
+            "i8" => Some(PrimitiveType::I8),
+            "i16" => Some(PrimitiveType::I16),
+            "i32" => Some(PrimitiveType::I32),
+            "i64" => Some(PrimitiveType::I64),
+            "i128" => Some(PrimitiveType::I128),
+            "u8" => Some(PrimitiveType::U8),
+            "u16" => Some(PrimitiveType::U16),
+            "u32" => Some(PrimitiveType::U32),
+            "u64" => Some(PrimitiveType::U64),
+            "u128" => Some(PrimitiveType::U128),
+            "f32" => Some(PrimitiveType::F32),
+            "f64" => Some(PrimitiveType::F64),
+            "bool" => Some(PrimitiveType::Bool),
+            "char" => Some(PrimitiveType::Char),
+            "str" | "String" => Some(PrimitiveType::String),
+            "isize" => Some(PrimitiveType::ISize),
+            "usize" => Some(PrimitiveType::USize),
+            "void" => Some(PrimitiveType::Unit),
+            _ => None,
+        }
+    }
+
     /// Execute a helper function call
     fn execute_helper_call<'a>(
         &self,
diff --git a/crates/zynml/ml.zyn b/crates/zynml/ml.zyn
@@ -764,9 +764,10 @@ type_opaque = { "@" ~ "opaque" ~ "(" ~ string_literal ~ ")" ~ "type" ~ name:iden
   }
 
 // Type alias: type Scalar = f32
-type_alias = { "type" ~ name:identifier ~ "=" ~ type_expr }
+type_alias = { "type" ~ name:identifier ~ "=" ~ target:type_expr }
   -> TypedDeclaration::TypeAlias {
       name: intern(name),
+      target: target,
   }
 
 // Struct type: type Point = { x: f32, y: f32 }
@@ -1367,6 +1368,7 @@ optional_type = { "?" ~ inner:type_expr_non_optional }
   -> Type::Named {
       name: intern("Option"),
       type_args: [inner],
+      const_args: [],
   }
 
 // Type expression without the optional prefix (to avoid infinite recursion)
@@ -1378,9 +1380,12 @@ type_expr_non_optional = { ty:fn_type | ty:tuple_type | ty:tensor_type | ty:gene
 //   tensor[2, 3]
 //   tensor[batch, seq, hidden]
 //   tensor[1, 3, 224, 224, f32]
-tensor_type = { "tensor" ~ "[" ~ tensor_type_items? ~ "]" }
+tensor_type = { "tensor" ~ "[" ~ items:tensor_type_items? ~ "]" }
   -> Type::Named {
       name: intern("Tensor"),
+      tensor_items: items,
+      type_args: [],
+      const_args: [],
   }
 
 tensor_type_items = { first:tensor_type_item ~ rest:tensor_type_item_comma* }
@@ -1389,7 +1394,7 @@ tensor_type_items = { first:tensor_type_item ~ rest:tensor_type_item_comma* }
 tensor_type_item_comma = { "," ~ item:tensor_type_item }
   -> item
 
-tensor_type_item = { identifier | "?" | ASCII_DIGIT+ }
+tensor_type_item = @{ identifier | "?" | ASCII_DIGIT+ }
   -> text()
 
 // Primitive types must be matched BEFORE simple_type to avoid treating them as named types
@@ -1453,16 +1458,22 @@ prim_unit = { "()" }
 simple_type = { name:identifier }
   -> Type::Named {
       name: intern(name),
+      type_args: [],
+      const_args: [],
   }
 
 array_type = { "Array" ~ "<" ~ elem:type_expr ~ ">" }
   -> Type::Named {
       name: intern("Array"),
+      type_args: [elem],
+      const_args: [],
   }
 
-generic_type = { name:identifier ~ "<" ~ type_args ~ ">" }
+generic_type = { name:identifier ~ "<" ~ args:type_args ~ ">" }
   -> Type::Named {
       name: intern(name),
+      type_args: args,
+      const_args: [],
   }
 
 // Type arguments: can be regular types or associated type bindings (Item=T)
@@ -1505,6 +1516,8 @@ where_bounds = { bound:where_bound ~ ("," ~ bound:where_bound)* }
 where_bound = { name:identifier ~ ":" ~ bounds:where_trait_bounds }
   -> Type::Named {
       name: intern(name),
+      type_args: [],
+      const_args: [],
   }
 
 // Multiple trait bounds separated by +: Clone + Debug + Iterator<Item=T>
@@ -1519,12 +1532,16 @@ where_trait_bound = { bound:trait_bound_generic | bound:trait_bound_fn | bound:t
 trait_bound_generic = { name:identifier ~ "<" ~ args:type_args ~ ">" }
   -> Type::Named {
       name: intern(name),
+      type_args: args,
+      const_args: [],
   }
 
 // Simple trait bound: Clone, Debug
 trait_bound_simple = { name:identifier }
   -> Type::Named {
       name: intern(name),
+      type_args: [],
+      const_args: [],
   }
 
 // Function type as trait bound: (T) => U
diff --git a/crates/zynml/tests/e2e_tests.rs b/crates/zynml/tests/e2e_tests.rs
@@ -188,13 +188,60 @@ mod type_system {
     #[test]
     fn test_parse_type_alias_tensor() {
         let grammar = get_grammar();
-        // Current grammar doesn't support tensor[shape, dtype] syntax in type position
-        let result = grammar.parse_to_json("type Embedding = Tensor");
+        let result = grammar.parse_to_json("type Embedding = tensor[2, 3, f32]");
         assert!(
             result.is_ok(),
             "Should parse tensor type alias: {:?}",
             result.err()
         );
+
+        let parsed: serde_json::Value =
+            serde_json::from_str(&result.unwrap()).expect("JSON should deserialize");
+        let named = &parsed["declarations"][0]["node"]["TypeAlias"]["target"]["Named"];
+        let const_args = named["const_args"]
+            .as_array()
+            .expect("const_args should be an array");
+        assert_eq!(
+            const_args.len(),
+            2,
+            "tensor shape should have two dimensions"
+        );
+        assert_eq!(const_args[0]["Int"], 2);
+        assert_eq!(const_args[1]["Int"], 3);
+        assert_eq!(
+            named["type_args"][0]["Primitive"], "F32",
+            "tensor dtype should map to Primitive::F32"
+        );
+    }
+
+    #[test]
+    fn test_parse_type_alias_tensor_symbolic_shape() {
+        let grammar = get_grammar();
+        let result = grammar.parse_to_json("type HiddenState = tensor[batch, seq, hidden]");
+        assert!(
+            result.is_ok(),
+            "Should parse tensor type alias with symbolic shape: {:?}",
+            result.err()
+        );
+
+        let parsed: serde_json::Value =
+            serde_json::from_str(&result.unwrap()).expect("JSON should deserialize");
+        let named = &parsed["declarations"][0]["node"]["TypeAlias"]["target"]["Named"];
+        let const_args = named["const_args"]
+            .as_array()
+            .expect("const_args should be an array");
+        assert_eq!(const_args.len(), 3);
+        assert_eq!(const_args[0]["Variable"], "batch");
+        assert_eq!(const_args[1]["Variable"], "seq");
+        assert_eq!(const_args[2]["Variable"], "hidden");
+        assert_eq!(
+            named["type_args"]
+                .as_array()
+                .expect("type_args should be an array")
+                .len(),
+            0,
+            "symbolic tensor shape should not imply dtype when omitted"
+        );
     }
 
     // --- Struct Definitions ---
diff --git a/crates/zyntax_embed/src/runtime.rs b/crates/zyntax_embed/src/runtime.rs
@@ -1669,6 +1669,21 @@ impl ZyntaxRuntime {
                     );
                 }
             }
+            Type::Named { id, type_args, .. } => {
+                for type_arg in type_args {
+                    Self::resolve_in_type(type_arg, type_registry);
+                }
+
+                // Canonicalize Named IDs by name in case this ID came from a placeholder
+                // registry entry created before imports/externs were merged.
+                if let Some(type_def) = type_registry.get_type_by_id(*id) {
+                    if let Some(canonical) = type_registry.get_type_by_name(type_def.name) {
+                        if canonical.id != *id {
+                            *id = canonical.id;
+                        }
+                    }
+                }
+            }
             // Recursively resolve nested types
             Type::Reference { ty: inner, .. } => {
                 Self::resolve_in_type(inner, type_registry);
