kore/
types.rs

1//! KORE Type System - Rust-like with effect tracking
2
3use crate::ast::*;
4use crate::effects::EffectSet;
5use crate::span::Span;
6use crate::error::{KoreError, KoreResult};
7use std::collections::HashMap;
8
9/// Type-checked AST node
10#[derive(Debug, Clone)]
11pub struct TypedProgram {
12    pub items: Vec<TypedItem>,
13}
14
15// Comptime blocks should be empty/removed by now if fully evaluated, or we check them
16#[derive(Debug, Clone)]
17pub struct TypedComptime {
18    pub ast: Block,
19}
20
21#[derive(Debug, Clone)]
22pub struct TypedConst {
23    pub ast: Const,
24    pub ty: ResolvedType,
25}
26
27#[derive(Debug, Clone)]
28pub struct TypedUse {
29    pub ast: Use,
30}
31
32#[derive(Debug, Clone)]
33pub enum TypedItem {
34    Function(TypedFunction),
35    Component(TypedComponent),
36    Shader(TypedShader),
37    Actor(TypedActor),
38    Struct(TypedStruct),
39    Enum(TypedEnum),
40    Comptime(TypedComptime),
41    Const(TypedConst),
42    Macro(TypedMacro),
43    Use(TypedUse),
44    Impl(TypedImpl),
45    Test(TypedTest),
46}
47
48#[derive(Debug, Clone)]
49pub struct TypedTest {
50    pub ast: TestDef,
51}
52
53#[derive(Debug, Clone)]
54pub struct TypedImpl {
55    pub ast: Impl,
56}
57
58#[derive(Debug, Clone)]
59pub struct TypedMacro {
60    pub ast: MacroDef,
61}
62
63#[derive(Debug, Clone)]
64pub struct TypedActor {
65    pub ast: Actor,
66    pub state_types: HashMap<String, ResolvedType>,
67}
68
69#[derive(Debug, Clone)]
70pub struct TypedFunction {
71    pub ast: Function,
72    pub resolved_type: ResolvedType,
73    pub effects: EffectSet,
74}
75
76#[derive(Debug, Clone)]
77pub struct TypedComponent {
78    pub ast: Component,
79    pub prop_types: HashMap<String, ResolvedType>,
80}
81
82#[derive(Debug, Clone)]
83pub struct TypedShader {
84    pub ast: Shader,
85    pub input_types: Vec<ResolvedType>,
86    pub output_type: ResolvedType,
87}
88
89#[derive(Debug, Clone)]
90pub struct TypedStruct {
91    pub ast: Struct,
92    pub field_types: HashMap<String, ResolvedType>,
93}
94
95#[derive(Debug, Clone)]
96pub struct TypedEnum {
97    pub ast: Enum,
98    pub variant_payload_types: HashMap<String, Vec<ResolvedType>>,
99}
100
101/// Fully resolved type
102#[derive(Debug, Clone, PartialEq)]
103pub enum ResolvedType {
104    Unit,
105    Bool,
106    Int(IntSize),
107    Float(FloatSize),
108    String,
109    Char,
110    Array(Box<ResolvedType>, usize),
111    Slice(Box<ResolvedType>),
112    Tuple(Vec<ResolvedType>),
113    Option(Box<ResolvedType>),
114    Result(Box<ResolvedType>, Box<ResolvedType>),
115    Ref { mutable: bool, inner: Box<ResolvedType> },
116    Function { params: Vec<ResolvedType>, ret: Box<ResolvedType>, effects: EffectSet },
117    Struct(String, HashMap<String, ResolvedType>),
118    Enum(String, Vec<(String, ResolvedType)>),
119    Generic(String),
120    Never,
121    Unknown,
122}
123
124#[derive(Debug, Clone, Copy, PartialEq, Eq)]
125pub enum IntSize { I8, I16, I32, I64, I128, Isize, U8, U16, U32, U64, U128, Usize }
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
128pub enum FloatSize { F32, F64 }
129
130/// Type environment for checking
131pub struct TypeEnv {
132    scopes: Vec<HashMap<String, ResolvedType>>,
133    types: HashMap<String, ResolvedType>,
134}
135
136impl TypeEnv {
137    pub fn new() -> Self {
138        let mut env = Self { scopes: vec![HashMap::new()], types: HashMap::new() };
139        // Built-in types
140        env.types.insert("Int".into(), ResolvedType::Int(IntSize::I64));
141        env.types.insert("Float".into(), ResolvedType::Float(FloatSize::F64));
142        env.types.insert("Bool".into(), ResolvedType::Bool);
143        env.types.insert("String".into(), ResolvedType::String);
144        env.types.insert("Vec2".into(), ResolvedType::Tuple(vec![
145            ResolvedType::Float(FloatSize::F32),
146            ResolvedType::Float(FloatSize::F32),
147        ]));
148        env.types.insert("Vec3".into(), ResolvedType::Tuple(vec![
149            ResolvedType::Float(FloatSize::F32),
150            ResolvedType::Float(FloatSize::F32),
151            ResolvedType::Float(FloatSize::F32),
152        ]));
153        env
154    }
155
156    pub fn push_scope(&mut self) { self.scopes.push(HashMap::new()); }
157    pub fn pop_scope(&mut self) { self.scopes.pop(); }
158    
159    pub fn define(&mut self, name: String, ty: ResolvedType) {
160        if let Some(scope) = self.scopes.last_mut() {
161            scope.insert(name, ty);
162        }
163    }
164    
165    pub fn lookup(&self, name: &str) -> Option<&ResolvedType> {
166        for scope in self.scopes.iter().rev() {
167            if let Some(ty) = scope.get(name) { return Some(ty); }
168        }
169        self.types.get(name)
170    }
171}
172
173/// Main type checking entry point
174pub fn check(program: &Program) -> KoreResult<TypedProgram> {
175    let mut env = TypeEnv::new();
176    let mut typed_items = Vec::new();
177    
178    for item in &program.items {
179        typed_items.push(check_item(&mut env, item)?);
180    }
181    
182    Ok(TypedProgram { items: typed_items })
183}
184
185fn check_item(env: &mut TypeEnv, item: &Item) -> KoreResult<TypedItem> {
186    match item {
187        Item::Function(f) => Ok(TypedItem::Function(check_function(env, f)?)),
188        Item::Struct(s) => Ok(TypedItem::Struct(check_struct(env, s)?)),
189        Item::Enum(e) => Ok(TypedItem::Enum(check_enum(env, e)?)),
190        Item::Component(c) => Ok(TypedItem::Component(check_component(env, c)?)),
191        Item::Shader(s) => Ok(TypedItem::Shader(check_shader(env, s)?)),
192        Item::Actor(a) => Ok(TypedItem::Actor(check_actor(env, a)?)),
193        Item::Comptime(b) => Ok(TypedItem::Comptime(TypedComptime { ast: b.body.clone() })),
194        Item::Const(c) => Ok(TypedItem::Const(check_const(env, c)?)),
195        Item::Macro(m) => Ok(TypedItem::Macro(TypedMacro { ast: m.clone() })),
196        Item::Use(u) => Ok(TypedItem::Use(TypedUse { ast: u.clone() })),
197        Item::Impl(i) => Ok(TypedItem::Impl(TypedImpl { ast: i.clone() })),
198        Item::Test(t) => Ok(TypedItem::Test(TypedTest { ast: t.clone() })),
199        _ => {
200            // For now, ignore other items or provide dummy implementation
201            // Since we are running in interpreter mode mostly, types are just for checking.
202            // But we shouldn't fail hard if we encounter valid syntax.
203            // Let's return a dummy TypedItem if possible, or just skip it?
204            // TypedItem enum doesn't have a variant for others.
205            // We should probably just return a dummy function or error with a better message?
206            // Or better, expand TypedItem to include other items.
207            // For now, let's just error if it's something we really don't support yet,
208            // but for simple scripts, we likely only need the above.
209            // If the script has top-level stmts, they are wrapped in main function.
210            // So we are good.
211            Err(KoreError::type_error("Item type not yet supported in type checker", item_span(item)))
212        }
213    }
214}
215
216fn check_const(_env: &mut TypeEnv, c: &Const) -> KoreResult<TypedConst> {
217    let ty = resolve_type(&c.ty)?;
218    // TODO: Check if value matches type
219    Ok(TypedConst { ast: c.clone(), ty })
220}
221
222fn check_actor(_env: &mut TypeEnv, a: &Actor) -> KoreResult<TypedActor> {
223    let mut state_types = HashMap::new();
224    for s in &a.state {
225        state_types.insert(s.name.clone(), resolve_type(&s.ty)?);
226    }
227    Ok(TypedActor { ast: a.clone(), state_types })
228}
229
230fn check_function(env: &mut TypeEnv, f: &Function) -> KoreResult<TypedFunction> {
231    env.push_scope();
232    let mut param_types = Vec::new();
233    for p in &f.params {
234        let ty = resolve_type(&p.ty)?;
235        env.define(p.name.clone(), ty.clone());
236        param_types.push(ty);
237    }
238    let ret = f.return_type.as_ref().map(|t| resolve_type(t)).transpose()?.unwrap_or(ResolvedType::Unit);
239    let effects = EffectSet::from(f.effects.clone());
240    env.pop_scope();
241    
242    Ok(TypedFunction {
243        ast: f.clone(),
244        resolved_type: ResolvedType::Function { params: param_types, ret: Box::new(ret), effects: effects.clone() },
245        effects,
246    })
247}
248
249fn check_struct(_env: &mut TypeEnv, s: &Struct) -> KoreResult<TypedStruct> {
250    let mut fields = HashMap::new();
251    for f in &s.fields {
252        fields.insert(f.name.clone(), resolve_type(&f.ty)?);
253    }
254    Ok(TypedStruct { ast: s.clone(), field_types: fields })
255}
256
257fn check_enum(_env: &mut TypeEnv, e: &Enum) -> KoreResult<TypedEnum> {
258    let mut variant_payload_types: HashMap<String, Vec<ResolvedType>> = HashMap::new();
259
260    for v in &e.variants {
261        let payload_types = match &v.fields {
262            VariantFields::Unit => Vec::new(),
263            VariantFields::Tuple(items) => items.iter().map(resolve_type).collect::<Result<Vec<_>, _>>()?,
264            VariantFields::Struct(fields) => fields.iter().map(|f| resolve_type(&f.ty)).collect::<Result<Vec<_>, _>>()?,
265        };
266        variant_payload_types.insert(v.name.clone(), payload_types);
267    }
268
269    Ok(TypedEnum {
270        ast: e.clone(),
271        variant_payload_types,
272    })
273}
274
275fn check_component(_env: &mut TypeEnv, c: &Component) -> KoreResult<TypedComponent> {
276    let mut props = HashMap::new();
277    for p in &c.props {
278        props.insert(p.name.clone(), resolve_type(&p.ty)?);
279    }
280    Ok(TypedComponent { ast: c.clone(), prop_types: props })
281}
282
283fn check_shader(_env: &mut TypeEnv, s: &Shader) -> KoreResult<TypedShader> {
284    let inputs: Vec<_> = s.inputs.iter().map(|p| resolve_type(&p.ty)).collect::<Result<_, _>>()?;
285    let output = resolve_type(&s.outputs)?;
286    Ok(TypedShader { ast: s.clone(), input_types: inputs, output_type: output })
287}
288
289pub fn resolve_type(ty: &Type) -> KoreResult<ResolvedType> {
290    match ty {
291        Type::Named { name, .. } => match name.as_str() {
292            "Int" => Ok(ResolvedType::Int(IntSize::I64)),
293            "Float" => Ok(ResolvedType::Float(FloatSize::F64)),
294            "Bool" => Ok(ResolvedType::Bool),
295            "String" => Ok(ResolvedType::String),
296            _ => {
297                // Check if this is a generic type parameter (single uppercase letter or _T style)
298                if name.len() == 1 && name.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) {
299                    Ok(ResolvedType::Generic(name.clone()))
300                } else if name.starts_with('_') && name.len() > 1 {
301                    // _T, _Item, etc are also generic
302                    Ok(ResolvedType::Generic(name.clone()))
303                } else {
304                    // Assume it's a struct
305                    Ok(ResolvedType::Struct(name.clone(), HashMap::new()))
306                }
307            }
308        },
309        Type::Unit(_) => Ok(ResolvedType::Unit),
310        Type::Never(_) => Ok(ResolvedType::Never),
311        Type::Tuple(inner, _) => Ok(ResolvedType::Tuple(inner.iter().map(resolve_type).collect::<Result<_, _>>()?)),
312        Type::Function { params, return_type, effects, .. } => {
313            let resolved_params = params.iter().map(resolve_type).collect::<Result<Vec<_>, _>>()?;
314            let resolved_ret = resolve_type(return_type)?;
315            Ok(ResolvedType::Function {
316                params: resolved_params,
317                ret: Box::new(resolved_ret),
318                effects: EffectSet::from(effects.clone()),
319            })
320        }
321        _ => Ok(ResolvedType::Unknown),
322    }
323}
324
325fn item_span(item: &Item) -> Span {
326    match item {
327        Item::Function(f) => f.span,
328        Item::Struct(s) => s.span,
329        Item::Enum(e) => e.span,
330        Item::Component(c) => c.span,
331        Item::Shader(s) => s.span,
332        Item::Actor(a) => a.span,
333        Item::Comptime(b) => b.span,
334        Item::Const(c) => c.span,
335        Item::Macro(m) => m.span,
336        Item::Use(u) => u.span,
337        Item::Impl(i) => i.span,
338        Item::Test(t) => t.span,
339        _ => Span::new(0, 0),
340    }
341}
342
343impl From<Vec<crate::effects::Effect>> for EffectSet {
344    fn from(v: Vec<crate::effects::Effect>) -> Self {
345        let mut s = EffectSet::new();
346        for e in v { s.effects.insert(e); }
347        s
348    }
349}
350
kore/types.rs

kore/
types.rs
