From 3759ff944a4ed8f625b88e7d39c989d585173b64 Mon Sep 17 00:00:00 2001 From: peefy Date: Mon, 4 Sep 2023 13:03:04 +0800 Subject: [PATCH] feat: add base64 and net member function types --- kclvm/sema/src/builtin/string.rs | 2 +- kclvm/sema/src/builtin/system_module.rs | 100 +++++++++++++++++- kclvm/sema/src/resolver/attr.rs | 7 +- .../src/resolver/test_data/system_package.k | 4 + kclvm/sema/src/resolver/tests.rs | 32 ++++++ 5 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 kclvm/sema/src/resolver/test_data/system_package.k diff --git a/kclvm/sema/src/builtin/string.rs b/kclvm/sema/src/builtin/string.rs index 32903bf52..2e636abeb 100644 --- a/kclvm/sema/src/builtin/string.rs +++ b/kclvm/sema/src/builtin/string.rs @@ -18,7 +18,7 @@ macro_rules! register_string_member { register_string_member! { capitalize => Type::function( Some(Rc::new(Type::STR)), - Rc::new(Type::ANY), + Rc::new(Type::STR), &[], r#"Return a copy of the string with its first character capitalized and the rest lowercased."#, false, diff --git a/kclvm/sema/src/builtin/system_module.rs b/kclvm/sema/src/builtin/system_module.rs index 86da72ea2..accb49dc2 100644 --- a/kclvm/sema/src/builtin/system_module.rs +++ b/kclvm/sema/src/builtin/system_module.rs @@ -1,7 +1,62 @@ -// Copyright 2021 The KCL Authors. All rights reserved. +// Copyright The KCL Authors. All rights reserved. + +use std::rc::Rc; + +use crate::ty::{Parameter, Type, TypeRef}; +use indexmap::IndexMap; +use once_cell::sync::Lazy; pub const BASE64: &str = "base64"; pub const BASE64_FUNCTION_NAMES: [&str; 2] = ["encode", "decode"]; +macro_rules! register_base64_member { + ($($name:ident => $ty:expr)*) => ( + pub const BASE64_FUNCTION_TYPES: Lazy> = Lazy::new(|| { + let mut builtin_mapping = IndexMap::default(); + $( builtin_mapping.insert(stringify!($name).to_string(), $ty); )* + builtin_mapping + }); + ) +} +register_base64_member! { + encode => Type::function( + None, + Type::str_ref(), + &[ + Parameter { + name: "value".to_string(), + ty: Type::str_ref(), + has_default: false, + }, + Parameter { + name: "encoding".to_string(), + ty: Type::str_ref(), + has_default: true, + }, + ], + r#"Encode the string `value` using the codec registered for encoding."#, + false, + None, + ) + decode => Type::function( + None, + Type::str_ref(), + &[ + Parameter { + name: "value".to_string(), + ty: Type::str_ref(), + has_default: false, + }, + Parameter { + name: "encoding".to_string(), + ty: Type::str_ref(), + has_default: true, + }, + ], + r#"Decode the string `value` using the codec registered for encoding."#, + false, + None, + ) +} pub const NET: &str = "net"; pub const NET_FUNCTION_NAMES: [&str; 16] = [ @@ -22,6 +77,32 @@ pub const NET_FUNCTION_NAMES: [&str; 16] = [ "is_global_unicast_IP", "is_unspecified_IP", ]; +macro_rules! register_net_member { + ($($name:ident => $ty:expr)*) => ( + pub const NET_FUNCTION_TYPES: Lazy> = Lazy::new(|| { + let mut builtin_mapping = IndexMap::default(); + $( builtin_mapping.insert(stringify!($name).to_string(), $ty); )* + builtin_mapping + }); + ) +} +// TODO: add more system package types. +register_net_member! { + split_host_port => Type::function( + None, + Type::list_ref(Type::str_ref()), + &[ + Parameter { + name: "ip_end_point".to_string(), + ty: Type::str_ref(), + has_default: false, + }, + ], + r#"Split the `host` and `port` from the `ip_end_point`."#, + false, + None, + ) +} pub const MANIFESTS: &str = "manifests"; pub const MANIFESTS_FUNCTION_NAMES: [&str; 1] = ["yaml_stream"]; @@ -134,3 +215,20 @@ pub fn get_system_module_members(name: &str) -> Vec<&str> { _ => bug!("invalid system module name '{}'", name), } } + +/// Get the system package member function type, if not found, return the any type. +pub fn get_system_member_function_ty(name: &str, func: &str) -> TypeRef { + // TODO: add more system package types. + let optional_ty = match name { + BASE64 => { + let types = BASE64_FUNCTION_TYPES; + types.get(func).cloned() + } + NET => { + let types = NET_FUNCTION_TYPES; + types.get(func).cloned() + } + _ => None, + }; + optional_ty.map(|ty| Rc::new(ty)).unwrap_or(Type::any_ref()) +} diff --git a/kclvm/sema/src/resolver/attr.rs b/kclvm/sema/src/resolver/attr.rs index de30b2c0b..3e22f4d3c 100644 --- a/kclvm/sema/src/resolver/attr.rs +++ b/kclvm/sema/src/resolver/attr.rs @@ -1,7 +1,7 @@ use std::rc::Rc; use crate::builtin::system_module::{get_system_module_members, UNITS, UNITS_NUMBER_MULTIPLIER}; -use crate::builtin::STRING_MEMBER_FUNCTIONS; +use crate::builtin::{get_system_member_function_ty, STRING_MEMBER_FUNCTIONS}; use crate::resolver::Resolver; use crate::ty::{DictType, ModuleKind, Type, TypeKind}; use kclvm_error::diagnostic::Range; @@ -100,7 +100,10 @@ impl<'ctx> Resolver<'ctx> { (true, Rc::new(Type::number_multiplier_non_lit_ty())) } else { let members = get_system_module_members(&module_ty.pkgpath); - (members.contains(&attr), self.any_ty()) + ( + members.contains(&attr), + get_system_member_function_ty(&module_ty.pkgpath, attr), + ) } } ModuleKind::Plugin => (true, self.any_ty()), diff --git a/kclvm/sema/src/resolver/test_data/system_package.k b/kclvm/sema/src/resolver/test_data/system_package.k new file mode 100644 index 000000000..2873bc47e --- /dev/null +++ b/kclvm/sema/src/resolver/test_data/system_package.k @@ -0,0 +1,4 @@ +import base64 + +base64_encode = base64.encode +base64_decode = base64.decode diff --git a/kclvm/sema/src/resolver/tests.rs b/kclvm/sema/src/resolver/tests.rs index 6db600e02..0bf7678ec 100644 --- a/kclvm/sema/src/resolver/tests.rs +++ b/kclvm/sema/src/resolver/tests.rs @@ -436,3 +436,35 @@ fn test_pkg_scope() { assert!(pkg_scope.contains_pos(&pos)); } + +#[test] +fn test_system_package() { + let sess = Arc::new(ParseSession::default()); + let mut program = load_program( + sess.clone(), + &["./src/resolver/test_data/system_package.k"], + None, + ) + .unwrap(); + let scope = resolve_program(&mut program); + let main_scope = scope + .scope_map + .get(kclvm_runtime::MAIN_PKG_PATH) + .unwrap() + .borrow_mut() + .clone(); + + assert!(main_scope.lookup("base64").unwrap().borrow().ty.is_module()); + assert!(main_scope + .lookup("base64_encode") + .unwrap() + .borrow() + .ty + .is_func()); + assert!(main_scope + .lookup("base64_decode") + .unwrap() + .borrow() + .ty + .is_func()); +}