ty.rs

use std::rc::Rc;

use crate::resolver::Resolver;
use crate::ty::parser::parse_type_str;
use crate::ty::{assignable_to, is_upper_bound, Attr, DictType, SchemaType, Type, TypeKind};
use indexmap::IndexMap;
use kclvm_ast::ast;
use kclvm_ast::pos::GetPos;
use kclvm_error::diagnostic::Range;
use kclvm_error::*;

use super::node::ResolvedResult;

fn ty_str_to_pkgpath(ty_str: &str) -> &str {
    let splits: Vec<&str> = ty_str.rsplitn(2, '.').collect();
    let len = splits.len();
    splits[len - 1]
}

pub fn ty_str_replace_pkgpath(ty_str: &str, pkgpath: &str) -> String {
    let pkgpath = format!("@{}", pkgpath);
    if ty_str.contains('.') && ty_str_to_pkgpath(ty_str) == pkgpath {
        ty_str.replacen(&format!("{}.", pkgpath), "", 1)
    } else {
        ty_str.to_string()
    }
}

impl<'ctx> Resolver<'ctx> {
    #[inline]
    pub fn any_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.any.clone()
    }
    #[inline]
    pub fn int_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.int.clone()
    }
    #[inline]
    pub fn float_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.float.clone()
    }
    #[inline]
    pub fn bool_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.bool.clone()
    }
    #[inline]
    pub fn str_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.str.clone()
    }
    #[inline]
    pub fn none_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.none.clone()
    }
    #[inline]
    pub fn void_ty(&self) -> Rc<Type> {
        self.ctx.ty_ctx.builtin_types.void.clone()
    }
    /// Parse the type string with the scope, if parse_ty returns a Named type(schema type or type alias),
    /// found it from the scope.
    pub fn parse_ty_with_scope(&mut self, ty: &ast::Type, range: Range) -> ResolvedResult {
        let ty: Rc<Type> = Rc::new(ty.clone().into());
        // If a named type, find it from scope to get the specific type
        let ret_ty = self.upgrade_named_ty_with_scope(ty.clone(), &range);
        self.add_type_alias(
            &ty.into_type_annotation_str(),
            &ret_ty.into_type_annotation_str(),
        );
        ret_ty
    }

    pub fn parse_ty_str_with_scope(&mut self, ty_str: &str, range: Range) -> ResolvedResult {
        let ty: Rc<Type> = parse_type_str(ty_str);
        // If a named type, find it from scope to get the specific type
        let ret_ty = self.upgrade_named_ty_with_scope(ty, &range);
        self.add_type_alias(ty_str, &ret_ty.into_type_annotation_str());
        ret_ty
    }

    /// The given expression must be the expected type.
    #[inline]
    pub fn must_be_type(&mut self, expr: &'ctx ast::NodeRef<ast::Expr>, expected_ty: Rc<Type>) {
        let ty = self.expr(expr);
        self.must_assignable_to(ty, expected_ty, expr.get_span_pos(), None);
    }

    /// Must assignable to the expected type.
    #[inline]
    pub fn must_assignable_to(
        &mut self,
        ty: Rc<Type>,
        expected_ty: Rc<Type>,
        range: Range,
        expected_pos: Option<Range>,
    ) {
        if !self.check_type(ty.clone(), expected_ty.clone(), &range) {
            let mut msgs = vec![Message {
                range,
                style: Style::LineAndColumn,
                message: format!("expected {}, got {}", expected_ty.ty_str(), ty.ty_str(),),
                note: None,
                suggested_replacement: None,
            }];

            if let Some(expected_pos) = expected_pos {
                msgs.push(Message {
                    range: expected_pos,
                    style: Style::LineAndColumn,
                    message: format!(
                        "variable is defined here, its type is {}, but got {}",
                        expected_ty.ty_str(),
                        ty.ty_str(),
                    ),
                    note: None,
                    suggested_replacement: None,
                });
            }
            self.handler.add_error(ErrorKind::TypeError, &msgs);
        }
    }

    /// Check the type assignment statement between type annotation and target.
    pub fn check_assignment_type_annotation(
        &mut self,
        assign_stmt: &kclvm_ast::ast::AssignStmt,
        value_ty: Rc<Type>,
    ) {
        if assign_stmt.type_annotation.is_none() {
            return;
        }
        for target in &assign_stmt.targets {
            if target.node.names.is_empty() {
                continue;
            }
            let name = &target.node.names[0].node;
            // If the assignment statement has type annotation, check the type of value and the type annotation of target

            if let Some(ty_annotation) = &assign_stmt.ty {
                let annotation_ty =
                    self.parse_ty_with_scope(&ty_annotation.node, ty_annotation.get_span_pos());
                // If the target defined in the scope, check the type of value and the type annotation of target
                let target_ty = if let Some(obj) = self.scope.borrow().elems.get(name) {
                    let obj = obj.borrow();
                    if obj.ty.is_any() {
                        annotation_ty
                    } else {
                        if !is_upper_bound(annotation_ty.clone(), obj.ty.clone()) {
                            self.handler.add_error(
                                ErrorKind::TypeError,
                                &[
                                    Message {
                                        range: target.get_span_pos(),
                                        style: Style::LineAndColumn,
                                        message: format!(
                                            "can not change the type of '{}' to {}",
                                            name,
                                            annotation_ty.ty_str()
                                        ),
                                        note: None,
                                        suggested_replacement: None,
                                    },
                                    Message {
                                        range: obj.get_span_pos(),
                                        style: Style::LineAndColumn,
                                        message: format!("expected {}", obj.ty.ty_str()),
                                        note: None,
                                        suggested_replacement: None,
                                    },
                                ],
                            );
                        }
                        obj.ty.clone()
                    }
                } else {
                    annotation_ty
                };

                self.set_type_to_scope(name, target_ty.clone(), target.get_span_pos());

                // Check the type of value and the type annotation of target
                self.must_assignable_to(value_ty.clone(), target_ty, target.get_span_pos(), None)
            }
        }
    }

    /// The check type main function, returns a boolean result.
    #[inline]
    pub fn check_type(&mut self, ty: Rc<Type>, expected_ty: Rc<Type>, range: &Range) -> bool {
        match (&ty.kind, &expected_ty.kind) {
            (TypeKind::List(item_ty), TypeKind::List(expected_item_ty)) => {
                self.check_type(item_ty.clone(), expected_item_ty.clone(), range)
            }
            (
                TypeKind::Dict(DictType { key_ty, val_ty, .. }),
                TypeKind::Dict(DictType {
                    key_ty: expected_key_ty,
                    val_ty: expected_val_ty,
                    ..
                }),
            ) => {
                self.check_type(key_ty.clone(), expected_key_ty.clone(), range)
                    && self.check_type(val_ty.clone(), expected_val_ty.clone(), range)
            }
            (TypeKind::Dict(DictType { key_ty, val_ty, .. }), TypeKind::Schema(schema_ty)) => {
                self.dict_assignable_to_schema(key_ty.clone(), val_ty.clone(), schema_ty, range)
            }
            (TypeKind::Union(types), _) => types
                .iter()
                .all(|ty| self.check_type(ty.clone(), expected_ty.clone(), range)),
            (_, TypeKind::Union(types)) => types
                .iter()
                .any(|expected_ty| self.check_type(ty.clone(), expected_ty.clone(), range)),
            _ => assignable_to(ty, expected_ty),
        }
    }

    /// Judge a dict can be converted to schema in compile time
    /// Do relaxed schema check key and value type check.
    pub fn dict_assignable_to_schema(
        &mut self,
        key_ty: Rc<Type>,
        val_ty: Rc<Type>,
        schema_ty: &SchemaType,
        range: &Range,
    ) -> bool {
        if let Some(index_signature) = &schema_ty.index_signature {
            if !assignable_to(val_ty.clone(), index_signature.val_ty.clone()) {
                self.handler.add_type_error(
                    &format!(
                        "expected schema index signature value type {}, got {}",
                        index_signature.val_ty.ty_str(),
                        val_ty.ty_str()
                    ),
                    range.clone(),
                );
            }
            if index_signature.any_other {
                return assignable_to(key_ty, index_signature.key_ty.clone())
                    && assignable_to(val_ty, index_signature.val_ty.clone());
            }
            true
        } else {
            true
        }
    }

    fn upgrade_named_ty_with_scope(&mut self, ty: Rc<Type>, range: &Range) -> ResolvedResult {
        match &ty.kind {
            TypeKind::List(item_ty) => {
                Type::list_ref(self.upgrade_named_ty_with_scope(item_ty.clone(), range))
            }
            TypeKind::Dict(DictType {
                key_ty,
                val_ty,
                attrs,
            }) => Type::dict_ref_with_attrs(
                self.upgrade_named_ty_with_scope(key_ty.clone(), range),
                self.upgrade_named_ty_with_scope(val_ty.clone(), range),
                attrs
                    .into_iter()
                    .map(|(key, attr)| {
                        (
                            key.to_string(),
                            Attr {
                                ty: self.upgrade_named_ty_with_scope(val_ty.clone(), range),
                                range: attr.range.clone(),
                            },
                        )
                    })
                    .collect(),
            ),
            TypeKind::Union(types) => Type::union_ref(
                &types
                    .iter()
                    .map(|ty| self.upgrade_named_ty_with_scope(ty.clone(), range))
                    .collect::<Vec<Rc<Type>>>(),
            ),
            TypeKind::Named(ty_str) => {
                let ty_str = ty_str_replace_pkgpath(ty_str, &self.ctx.pkgpath);
                let names: Vec<&str> = if ty_str.starts_with('@') {
                    let names: Vec<&str> = ty_str.rsplitn(2, '.').collect();
                    names.iter().rev().cloned().collect()
                } else {
                    ty_str.split('.').collect()
                };
                if names.is_empty() {
                    self.handler
                        .add_compile_error("missing type annotation", range.clone());
                    return self.any_ty();
                }
                let mut pkgpath = "".to_string();
                let name = names[0];
                if names.len() > 1 && !self.ctx.local_vars.contains(&name.to_string()) {
                    if let Some(mapping) = self.ctx.import_names.get(&self.ctx.filename) {
                        pkgpath = mapping
                            .get(name)
                            .map_or("".to_string(), |pkgpath| pkgpath.to_string());
                    }
                }
                self.ctx.l_value = false;
                self.resolve_var(
                    &names.iter().map(|n| n.to_string()).collect::<Vec<String>>(),
                    &pkgpath,
                    range.clone(),
                )
            }
            _ => ty.clone(),
        }
    }

    pub fn add_type_alias(&mut self, name: &str, alias: &str) {
        if alias.starts_with('@') {
            if name == &alias[1..] {
                return;
            }
        } else if name == alias {
            return;
        }
        match self.ctx.type_alias_mapping.get_mut(&self.ctx.pkgpath) {
            Some(mapping) => {
                mapping.insert(name.to_string(), alias.to_string());
            }
            None => {
                let mut mapping = IndexMap::default();
                mapping.insert(name.to_string(), alias.to_string());
                self.ctx
                    .type_alias_mapping
                    .insert(self.ctx.pkgpath.clone(), mapping);
            }
        }
    }
}