glyph-types 0.0.1

//! Type inference engine for Glyph

use crate::{Type, TypeError};
use std::collections::HashMap;

/// Type inference context
pub struct InferenceContext {
    /// Variable bindings
    bindings: HashMap<String, Type>,

    /// Type variable substitutions
    substitutions: HashMap<u32, Type>,

    /// Next type variable ID
    next_type_var: u32,
}

impl InferenceContext {
    /// Create a new inference context
    pub fn new() -> Self {
        Self {
            bindings: HashMap::new(),
            substitutions: HashMap::new(),
            next_type_var: 0,
        }
    }

    /// Create a fresh type variable
    pub fn fresh_type_var(&mut self) -> Type {
        let id = self.next_type_var;
        self.next_type_var += 1;
        Type::TypeVar(id)
    }

    /// Add a variable binding
    pub fn bind_variable(&mut self, name: String, typ: Type) {
        self.bindings.insert(name, typ);
    }

    /// Look up a variable's type
    pub fn lookup_variable(&self, name: &str) -> Option<&Type> {
        self.bindings.get(name)
    }

    /// Unify two types
    pub fn unify(&mut self, t1: &Type, t2: &Type) -> Result<(), TypeError> {
        let t1 = self.apply_substitutions(t1);
        let t2 = self.apply_substitutions(t2);

        match (&t1, &t2) {
            // Same types unify
            (Type::Int, Type::Int)
            | (Type::Float, Type::Float)
            | (Type::Str, Type::Str)
            | (Type::Bool, Type::Bool)
            | (Type::Bytes, Type::Bytes)
            | (Type::Unit, Type::Unit) => Ok(()),

            // Unknown unifies with anything
            (Type::Unknown, _) | (_, Type::Unknown) => Ok(()),

            // Type variable unification
            (Type::TypeVar(id), t) | (t, Type::TypeVar(id)) => {
                if let Type::TypeVar(id2) = t {
                    if id == id2 {
                        return Ok(());
                    }
                }
                // Occurs check would go here for full correctness
                self.substitutions.insert(*id, t.clone());
                Ok(())
            }

            // Container types
            (Type::List(t1), Type::List(t2)) => self.unify(t1, t2),
            (Type::Optional(t1), Type::Optional(t2)) => self.unify(t1, t2),
            (Type::Promise(t1), Type::Promise(t2)) => self.unify(t1, t2),

            (Type::Dict(k1, v1), Type::Dict(k2, v2)) => {
                self.unify(k1, k2)?;
                self.unify(v1, v2)
            }

            (Type::Result(ok1, err1), Type::Result(ok2, err2)) => {
                self.unify(ok1, ok2)?;
                self.unify(err1, err2)
            }

            // Function types
            (
                Type::Function {
                    params: p1,
                    return_type: r1,
                },
                Type::Function {
                    params: p2,
                    return_type: r2,
                },
            ) => {
                if p1.len() != p2.len() {
                    return Err(TypeError::ArgumentCountMismatch {
                        expected: p1.len(),
                        found: p2.len(),
                    });
                }

                for ((_, t1), (_, t2)) in p1.iter().zip(p2.iter()) {
                    self.unify(t1, t2)?;
                }

                self.unify(r1, r2)
            }

            // Type mismatch
            _ => Err(TypeError::TypeMismatch {
                expected: t1,
                found: t2,
            }),
        }
    }

    /// Apply substitutions to a type
    pub fn apply_substitutions(&self, typ: &Type) -> Type {
        match typ {
            Type::TypeVar(id) => {
                if let Some(substitution) = self.substitutions.get(id) {
                    self.apply_substitutions(substitution)
                } else {
                    typ.clone()
                }
            }
            Type::List(t) => Type::List(Box::new(self.apply_substitutions(t))),
            Type::Dict(k, v) => Type::Dict(
                Box::new(self.apply_substitutions(k)),
                Box::new(self.apply_substitutions(v)),
            ),
            Type::Optional(t) => Type::Optional(Box::new(self.apply_substitutions(t))),
            Type::Promise(t) => Type::Promise(Box::new(self.apply_substitutions(t))),
            Type::Result(ok, err) => Type::Result(
                Box::new(self.apply_substitutions(ok)),
                Box::new(self.apply_substitutions(err)),
            ),
            Type::Function {
                params,
                return_type,
            } => Type::Function {
                params: params
                    .iter()
                    .map(|(name, t)| (name.clone(), self.apply_substitutions(t)))
                    .collect(),
                return_type: Box::new(self.apply_substitutions(return_type)),
            },
            _ => typ.clone(),
        }
    }
}

impl Default for InferenceContext {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_unify_basic_types() {
        let mut ctx = InferenceContext::new();

        // Same types should unify
        assert!(ctx.unify(&Type::Int, &Type::Int).is_ok());
        assert!(ctx.unify(&Type::Str, &Type::Str).is_ok());

        // Different types should not unify
        assert!(ctx.unify(&Type::Int, &Type::Str).is_err());
    }

    #[test]
    fn test_unify_with_unknown() {
        let mut ctx = InferenceContext::new();

        // Unknown should unify with anything
        assert!(ctx.unify(&Type::Unknown, &Type::Int).is_ok());
        assert!(ctx.unify(&Type::Str, &Type::Unknown).is_ok());
    }

    #[test]
    fn test_type_variable_substitution() {
        let mut ctx = InferenceContext::new();
        let tvar = ctx.fresh_type_var();

        // Unify type variable with concrete type
        assert!(ctx.unify(&tvar, &Type::Int).is_ok());

        // Apply substitutions should resolve the type variable
        let resolved = ctx.apply_substitutions(&tvar);
        assert_eq!(resolved, Type::Int);
    }
}