feat: add base64 and net member function types (kcl-lang#695)

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,

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<IndexMap<String, Type>> = 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<IndexMap<String, Type>> = 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())
+}

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()),

kclvm/sema/src/resolver/test_data/system_package.k

@@ -0,0 +1,4 @@
+import base64
+
+base64_encode = base64.encode
+base64_decode = base64.decode

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());
+}