put parse ast into type annotation into one function

nac3core/src/toplevel/mod.rs

@@ -3,15 +3,12 @@ use std::ops::{Deref, DerefMut};
 use std::{collections::HashMap, collections::HashSet, sync::Arc};
 
 use super::typecheck::type_inferencer::PrimitiveStore;
-use super::typecheck::typedef::{SharedUnifier, Type, TypeEnum, Unifier};
+use super::typecheck::typedef::{FunSignature, FuncArg, SharedUnifier, Type, TypeEnum, Unifier};
 use crate::symbol_resolver::SymbolResolver;
-use crate::typecheck::typedef::{FunSignature, FuncArg};
 use itertools::{izip, Itertools};
 use parking_lot::{Mutex, RwLock};
 use rustpython_parser::ast::{self, Stmt};
 
-
-
 #[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
 pub struct DefinitionId(pub usize);
 

nac3core/src/toplevel/type_annotation.rs

@@ -19,6 +19,7 @@ pub enum TypeAnnotation {
     TypeVarKind(u32, Type),
     SelfTypeKind(DefinitionId),
 }
+
 pub fn parse_ast_to_type_annotation_kinds<T>(
     resolver: &dyn SymbolResolver,
     top_level_defs: &[Arc<RwLock<TopLevelDef>>],
@@ -26,31 +27,111 @@ pub fn parse_ast_to_type_annotation_kinds<T>(
     primitives: &PrimitiveStore,
     expr: &ast::Expr<T>,
 ) -> Result<TypeAnnotation, String> {
-    let results = vec![
+    match &expr.node {
-        parse_ast_to_concrete_primitive_kind(
+        ast::ExprKind::Name { id, .. } => match id.as_str() {
-            resolver,
+            "int32" => Ok(TypeAnnotation::PrimitiveKind(primitives.int32)),
-            top_level_defs,
+            "int64" => Ok(TypeAnnotation::PrimitiveKind(primitives.int64)),
-            unifier,
+            "float" => Ok(TypeAnnotation::PrimitiveKind(primitives.float)),
-            primitives,
+            "bool" => Ok(TypeAnnotation::PrimitiveKind(primitives.bool)),
-            expr,
+            "None" => Ok(TypeAnnotation::PrimitiveKind(primitives.none)),
-        ),
+            x => {
-        parse_ast_to_concretized_custom_class_kind(
+                if let Some(obj_id) = resolver.get_identifier_def(x) {
-            resolver,
+                    let def = top_level_defs[obj_id.0].read();
-            top_level_defs,
+                    if let TopLevelDef::Class { .. } = &*def {
-            unifier,
+                        Ok(TypeAnnotation::ConcretizedCustomClassKind {
-            primitives,
+                            id: obj_id,
-            expr,
+                            params: vec![],
-        ),
+                        })
-        parse_ast_to_type_variable_kind(resolver, top_level_defs, unifier, primitives, expr),
-        parse_ast_to_virtual_kind(resolver, top_level_defs, unifier, primitives, expr),
-    ];
-    let results = results.iter().filter(|x| x.is_ok()).collect_vec();
-    if results.len() == 1 {
-        results[0].clone()
                     } else {
-        Err("cannot parsed the type annotation without ambiguity".into())
+                        Err("function cannot be used as a type".into())
+                    }
+                } else if let Some(ty) = resolver.get_symbol_type(unifier, primitives, id) {
+                    if let TypeEnum::TVar { id, meta: TypeVarMeta::Generic, range } =
+                        unifier.get_ty(ty).as_ref()
+                    {
+                        // NOTE: always create a new one here
+                        // and later unify if needed
+                        // but record the var_id of the original type var returned by symbol resolver
+                        let range = range.borrow();
+                        let range = range.as_slice();
+                        Ok(TypeAnnotation::TypeVarKind(
+                            *id,
+                            unifier.get_fresh_var_with_range(range).0,
+                        ))
+                    } else {
+                        Err("not a type variable identifier".into())
+                    }
+                } else {
+                    Err("name cannot be parsed as a type annotation".into())
                 }
             }
+        },
+
+        // TODO: subscript or call
+        ast::ExprKind::Subscript { value, slice, .. }
+            if { matches!(&value.node, ast::ExprKind::Name { id, .. } if id == "virtual") } =>
+        {
+            let def = parse_ast_to_type_annotation_kinds(
+                resolver,
+                top_level_defs,
+                unifier,
+                primitives,
+                slice.as_ref(),
+            )?;
+            if !matches!(def, TypeAnnotation::ConcretizedCustomClassKind { .. }) {
+                unreachable!("must be concretized custom class kind in the virtual")
+            }
+            Ok(TypeAnnotation::VirtualKind(def.into()))
+        }
+
+        ast::ExprKind::Subscript { value, slice, .. } => {
+            if let ast::ExprKind::Name { id, .. } = &value.node {
+                if vec!["virtual", "Generic"].contains(&id.as_str()) {
+                    return Err("keywords cannot be class name".into());
+                }
+                let obj_id = resolver
+                    .get_identifier_def(id)
+                    .ok_or_else(|| "unknown class name".to_string())?;
+                let def = top_level_defs[obj_id.0].read();
+                if let TopLevelDef::Class { .. } = &*def {
+                    let param_type_infos = if let ast::ExprKind::Tuple { elts, .. } = &slice.node {
+                        elts.iter()
+                            .map(|v| {
+                                parse_ast_to_type_annotation_kinds(
+                                    resolver,
+                                    top_level_defs,
+                                    unifier,
+                                    primitives,
+                                    v,
+                                )
+                            })
+                            .collect::<Result<Vec<_>, _>>()?
+                    } else {
+                        vec![parse_ast_to_type_annotation_kinds(
+                            resolver,
+                            top_level_defs,
+                            unifier,
+                            primitives,
+                            slice,
+                        )?]
+                    };
+                    // NOTE: allow type var in class generic application list
+                    Ok(TypeAnnotation::ConcretizedCustomClassKind {
+                        id: obj_id,
+                        params: param_type_infos,
+                    })
+                } else {
+                    Err("function cannot be used as a type".into())
+                }
+            } else {
+                Err("unsupported expression type for class name".into())
+            }
+        }
+
+        _ => Err("unsupported expression for type annotation".into()),
+    }
+}
+
 pub fn get_type_from_type_annotation_kinds(
     top_level_defs: &[Arc<RwLock<TopLevelDef>>],
     unifier: &mut Unifier,
@@ -154,157 +235,3 @@ pub fn get_type_from_type_annotation_kinds(
         }
     }
 }
-fn parse_ast_to_concrete_primitive_kind<T>(
-    _resolver: &dyn SymbolResolver,
-    _top_level_defs: &[Arc<RwLock<TopLevelDef>>],
-    _unifier: &mut Unifier,
-    primitives: &PrimitiveStore,
-    expr: &ast::Expr<T>,
-) -> Result<TypeAnnotation, String> {
-    match &expr.node {
-        ast::ExprKind::Name { id, .. } => match id.as_str() {
-            "int32" => Ok(TypeAnnotation::PrimitiveKind(primitives.int32)),
-            "int64" => Ok(TypeAnnotation::PrimitiveKind(primitives.int64)),
-            "float" => Ok(TypeAnnotation::PrimitiveKind(primitives.float)),
-            "bool" => Ok(TypeAnnotation::PrimitiveKind(primitives.bool)),
-            "None" => Ok(TypeAnnotation::PrimitiveKind(primitives.none)),
-            _ => Err("not primitive".into()),
-        },
-        _ => Err("not primitive".into()),
-    }
-}
-pub fn parse_ast_to_concretized_custom_class_kind<T>(
-    resolver: &dyn SymbolResolver,
-    top_level_defs: &[Arc<RwLock<TopLevelDef>>],
-    unifier: &mut Unifier,
-    primitives: &PrimitiveStore,
-    expr: &ast::Expr<T>,
-) -> Result<TypeAnnotation, String> {
-    match &expr.node {
-        ast::ExprKind::Name { id, .. } => match id.as_str() {
-            "int32" | "int64" | "float" | "bool" | "None" => {
-                Err("expect custom class instead of primitives here".into())
-            }
-            x => {
-                let obj_id = resolver
-                    .get_identifier_def(x)
-                    .ok_or_else(|| "unknown class name".to_string())?;
-                let def = top_level_defs[obj_id.0].read();
-                if let TopLevelDef::Class { .. } = &*def {
-                    Ok(TypeAnnotation::ConcretizedCustomClassKind {
-                        id: obj_id,
-                        params: vec![],
-                    })
-                } else {
-                    Err("function cannot be used as a type".into())
-                }
-            }
-        },
-        ast::ExprKind::Subscript { value, slice, .. } => {
-            if let ast::ExprKind::Name { id, .. } = &value.node {
-                if vec!["virtual", "Generic"].contains(&id.as_str()) {
-                    return Err("keywords cannot be class name".into());
-                }
-                let obj_id = resolver
-                    .get_identifier_def(id)
-                    .ok_or_else(|| "unknown class name".to_string())?;
-                let def = top_level_defs[obj_id.0].read();
-                if let TopLevelDef::Class { .. } = &*def {
-                    let param_type_infos =
-                        if let ast::ExprKind::Tuple { elts, .. } = &slice.node {
-                            elts.iter()
-                                .map(|v| {
-                                    parse_ast_to_type_annotation_kinds(
-                                        resolver,
-                                        top_level_defs,
-                                        unifier,
-                                        primitives,
-                                        v,
-                                    )
-                                })
-                                .collect::<Result<Vec<_>, _>>()?
-                        } else {
-                            vec![parse_ast_to_type_annotation_kinds(
-                                resolver,
-                                top_level_defs,
-                                unifier,
-                                primitives,
-                                slice,
-                            )?]
-                        };
-                    // TODO: allow type var in class generic application list
-                    // if param_type_infos
-                    //     .iter()
-                    //     .any(|x| matches!(x, TypeAnnotation::TypeVarKind(..)))
-                    // {
-                    //     return Err(
-                    //         "cannot apply type variable to class generic parameters".into()
-                    //     );
-                    // }
-                    Ok(TypeAnnotation::ConcretizedCustomClassKind {
-                        id: obj_id,
-                        params: param_type_infos,
-                    })
-                } else {
-                    Err("function cannot be used as a type".into())
-                }
-            } else {
-                Err("unsupported expression type for class name".into())
-            }
-        }
-        _ => Err("unsupported expression type for concretized class".into()),
-    }
-}
-pub fn parse_ast_to_virtual_kind<T>(
-    resolver: &dyn SymbolResolver,
-    top_level_defs: &[Arc<RwLock<TopLevelDef>>],
-    unifier: &mut Unifier,
-    primitives: &PrimitiveStore,
-    expr: &ast::Expr<T>,
-) -> Result<TypeAnnotation, String> {
-    match &expr.node {
-        ast::ExprKind::Subscript { value, slice, .. }
-            if { matches!(&value.node, ast::ExprKind::Name { id, .. } if id == "virtual") } =>
-        {
-            let def = parse_ast_to_concretized_custom_class_kind(
-                resolver,
-                top_level_defs,
-                unifier,
-                primitives,
-                slice.as_ref(),
-            )?;
-            if !matches!(def, TypeAnnotation::ConcretizedCustomClassKind { .. }) {
-                unreachable!("must be concretized custom class kind in the virtual")
-            }
-            Ok(TypeAnnotation::VirtualKind(def.into()))
-        }
-        _ => Err("virtual type annotation must be like `virtual[ .. ]`".into()),
-    }
-}
-pub fn parse_ast_to_type_variable_kind<T>(
-    resolver: &dyn SymbolResolver,
-    _top_level_defs: &[Arc<RwLock<TopLevelDef>>],
-    unifier: &mut Unifier,
-    primitives: &PrimitiveStore,
-    expr: &ast::Expr<T>,
-) -> Result<TypeAnnotation, String> {
-    if let ast::ExprKind::Name { id, .. } = &expr.node {
-        let ty = resolver
-            .get_symbol_type(unifier, primitives, id)
-            .ok_or_else(|| "unknown type variable name".to_string())?;
-        if let TypeEnum::TVar { id, meta: TypeVarMeta::Generic, range } =
-            unifier.get_ty(ty).as_ref()
-        {
-            // NOTE: always create a new one here
-            // and later unify if needed
-            // but record the var_id of the original type var returned by symbol resolver
-            let range = range.borrow();
-            let range = range.as_slice();
-            Ok(TypeAnnotation::TypeVarKind(*id, unifier.get_fresh_var_with_range(range).0))
-        } else {
-            Err("not a type variable identifier".into())
-        }
-    } else {
-        Err("unsupported expression for type variable".into())
-    }
-}

nac3core/src/typecheck/type_inferencer/test.rs

@@ -1,7 +1,10 @@
 use super::super::typedef::*;
 use super::*;
 use crate::symbol_resolver::*;
-use crate::{location::Location, toplevel::{DefinitionId, TopLevelDef}};
+use crate::{
+    location::Location,
+    toplevel::{DefinitionId, TopLevelDef},
+};
 use indoc::indoc;
 use itertools::zip;
 use parking_lot::RwLock;