fix: add and check the k code list to the entry. (#642)

* fix: add and check the k code list to the entry.

* fix: make fmt.

kclvm/config/src/path.rs

@@ -12,7 +12,9 @@
 //! The real path of `${my_pkg:KCL_MOD}/xxx/main.k` is `/usr/my_pkg/sub/main.k`.
 use anyhow::Result;
 use pcre2::bytes::Regex;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
+
+use crate::modfile::KCL_FILE_SUFFIX;
 
 #[derive(Clone, Debug, Default)]
 /// [`ModRelativePath`] is a path that is relative to the root package path.
@@ -144,6 +146,21 @@ impl ModRelativePath {
                 },
             ))
     }
+
+    /// [`is_dir`] returns true if the path is a directory.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use kclvm_config::path::ModRelativePath;
+    /// let path = ModRelativePath::new("${name:KCL_MOD}/src/path".to_string());
+    /// assert_eq!(path.is_dir(), true);
+    /// let path = ModRelativePath::new("${name:KCL_MOD}/src/path/main.k".to_string());
+    /// assert_eq!(path.is_dir(), false);
+    /// ```
+    pub fn is_dir(&self) -> bool {
+        Path::new(&self.path).is_dir() || !self.path.ends_with(KCL_FILE_SUFFIX)
+    }
 }
 
 #[cfg(test)]

kclvm/parser/src/entry.rs

@@ -117,6 +117,7 @@ pub struct Entry {
     name: String,
     path: String,
     k_files: Vec<String>,
+    k_code_lists: Vec<Option<String>>,
 }
 
 impl Entry {
@@ -126,6 +127,7 @@ impl Entry {
             name,
             path,
             k_files: vec![],
+            k_code_lists: vec![],
         }
     }
 
@@ -149,10 +151,32 @@ impl Entry {
         self.k_files.extend(k_files);
     }
 
+    /// [`extend_k_files_and_codes`] will extend the k files and k codes of [`Entry`] to the given k file and k code.
+    pub fn extend_k_files_and_codes(
+        &mut self,
+        k_files: Vec<String>,
+        k_codes: &mut VecDeque<String>,
+    ) {
+        for k_file in k_files.iter() {
+            self.k_code_lists.push(k_codes.pop_front());
+            self.k_files.push(k_file.to_string());
+        }
+    }
+
+    /// [`push_k_code`] will push the k code of [`Entry`] to the given k code.
+    pub fn push_k_code(&mut self, k_code: Option<String>) {
+        self.k_code_lists.push(k_code);
+    }
+
     /// [`get_k_files`] will return the k files of [`Entry`].
     pub fn get_k_files(&self) -> &Vec<String> {
         &self.k_files
     }
+
+    /// [`get_k_codes`] will return the k codes of [`Entry`].
+    pub fn get_k_codes(&self) -> &Vec<Option<String>> {
+        &self.k_code_lists
+    }
 }
 
 /// [`get_compile_entries_from_paths`] returns all the [`Entries`] for compilation from the given [`file_paths`].
@@ -251,8 +275,14 @@ pub fn get_compile_entries_from_paths(
         return Err("No input KCL files or paths".to_string());
     }
     let mut result = Entries::default();
-    for s in file_paths {
+    let mut k_code_queue = VecDeque::from(opts.k_code_list.clone());
+    for (i, s) in file_paths.iter().enumerate() {
         let path = ModRelativePath::from(s.to_string());
+
+        if path.is_dir() && opts.k_code_list.len() > i {
+            return Err("Invalid code list".to_string());
+        }
+
         // If the path is a [`ModRelativePath`] with preffix '${<package_name>:KCL_MOD}',
         // calculate the real path and the package name.
         if let Some((pkg_name, pkg_path)) = path
@@ -269,8 +299,11 @@ pub fn get_compile_entries_from_paths(
                 .canonicalize_by_root_path(pkg_path)
                 .map_err(|err| err.to_string())?;
             if let Some(root) = get_pkg_root(&s) {
-                let mut entry = Entry::new(pkg_name.clone(), root.clone());
-                entry.extend_k_files(get_main_files_from_pkg_path(&s, &root, &pkg_name, opts)?);
+                let mut entry: Entry = Entry::new(pkg_name.clone(), root.clone());
+                entry.extend_k_files_and_codes(
+                    get_main_files_from_pkg_path(&s, &root, &pkg_name, opts)?,
+                    &mut k_code_queue,
+                );
                 result.push_entry(entry);
                 continue;
             }
@@ -282,17 +315,17 @@ pub fn get_compile_entries_from_paths(
                 .is_none()
         {
             // Push it into `result`, and deal it later.
-            result.push(kclvm_ast::MAIN_PKG.to_string(), path.get_path());
+            let mut entry = Entry::new(kclvm_ast::MAIN_PKG.to_string(), path.get_path());
+            entry.push_k_code(k_code_queue.pop_front());
+            result.push_entry(entry);
             continue;
         } else if let Some(root) = get_pkg_root(s) {
             // If the path is a normal path.
             let mut entry: Entry = Entry::new(kclvm_ast::MAIN_PKG.to_string(), root.clone());
-            entry.extend_k_files(get_main_files_from_pkg_path(
-                &s,
-                &root,
-                &kclvm_ast::MAIN_PKG.to_string(),
-                opts,
-            )?);
+            entry.extend_k_files_and_codes(
+                get_main_files_from_pkg_path(&s, &root, &kclvm_ast::MAIN_PKG.to_string(), opts)?,
+                &mut k_code_queue,
+            );
             result.push_entry(entry);
         }
     }
@@ -304,12 +337,10 @@ pub fn get_compile_entries_from_paths(
     {
         let mut entry = Entry::new(kclvm_ast::MAIN_PKG.to_string(), "".to_string());
         for s in file_paths {
-            entry.extend_k_files(get_main_files_from_pkg_path(
-                s,
-                "",
-                &kclvm_ast::MAIN_PKG.to_string(),
-                opts,
-            )?);
+            entry.extend_k_files_and_codes(
+                get_main_files_from_pkg_path(s, "", &kclvm_ast::MAIN_PKG.to_string(), opts)?,
+                &mut k_code_queue,
+            );
         }
         result.push_entry(entry);
     }
@@ -432,7 +463,7 @@ fn get_main_files_from_pkg_path(
 }
 
 /// Get file list in the directory.
-fn get_dir_files(dir: &str) -> Result<Vec<String>, String> {
+pub fn get_dir_files(dir: &str) -> Result<Vec<String>, String> {
     if !std::path::Path::new(dir).exists() {
         return Ok(Vec::new());
     }

kclvm/parser/src/lib.rs

@@ -291,22 +291,15 @@ impl Loader {
             // Get files from options with root.
             // let k_files = self.get_main_files_from_pkg(entry.path(), entry.name())?;
             let k_files = entry.get_k_files();
+            let maybe_k_codes = entry.get_k_codes();
+
             // load module
 
             for (i, filename) in k_files.iter().enumerate() {
-                if i < self.opts.k_code_list.len() {
-                    let mut m = parse_file_with_session(
-                        self.sess.clone(),
-                        filename,
-                        Some(self.opts.k_code_list[i].clone()),
-                    )?;
-                    self.fix_rel_import_path(entry.path(), &mut m);
-                    pkg_files.push(m);
-                } else {
-                    let mut m = parse_file_with_session(self.sess.clone(), filename, None)?;
-                    self.fix_rel_import_path(entry.path(), &mut m);
-                    pkg_files.push(m);
-                }
+                let mut m =
+                    parse_file_with_session(self.sess.clone(), filename, maybe_k_codes[i].clone())?;
+                self.fix_rel_import_path(entry.path(), &mut m);
+                pkg_files.push(m);
             }
 
             // Insert an empty vec to determine whether there is a circular import.

kclvm/parser/src/testdata/test_k_code_list/main.k

@@ -0,0 +1 @@
+test = "test"

kclvm/parser/src/testdata/test_k_code_list/main1.k

@@ -0,0 +1 @@
+test1 = "test1"

kclvm/parser/src/tests.rs

@@ -497,3 +497,20 @@ fn test_get_compile_entries_from_paths() {
         kcl1_path.canonicalize().unwrap().to_str().unwrap()
     );
 }
+
+#[test]
+fn test_dir_with_k_code_list() {
+    let sm = SourceMap::new(FilePathMapping::empty());
+    let sess = Arc::new(ParseSession::with_source_map(Arc::new(sm)));
+    let testpath = PathBuf::from("./src/testdata/test_k_code_list")
+        .canonicalize()
+        .unwrap();
+
+    let mut opts = LoadProgramOptions::default();
+    opts.k_code_list = vec!["test_code = 1".to_string()];
+
+    match load_program(sess.clone(), &[&testpath.display().to_string()], Some(opts)) {
+        Ok(_) => panic!("unreachable code"),
+        Err(err) => assert_eq!(err, "Invalid code list"),
+    }
+}

kclvm/tools/src/LSP/src/tests.rs

@@ -1,6 +1,7 @@
 use std::env;
 use std::path::PathBuf;
 use std::process::Command;
+use std::sync::Arc;
 
 use indexmap::IndexSet;
 use kclvm_ast::ast::Program;
@@ -19,6 +20,7 @@ use lsp_types::MarkedString;
 use lsp_types::SymbolKind;
 use lsp_types::Url;
 use lsp_types::{Position, Range, TextDocumentContentChangeEvent};
+use parking_lot::RwLock;
 
 use crate::completion::KCLCompletionItem;
 use crate::document_symbol::document_symbol;
@@ -38,8 +40,11 @@ fn compile_test_file(testfile: &str) -> (String, Program, ProgramScope, IndexSet
 
     let file = test_file.to_str().unwrap().to_string();
 
-    let (program, prog_scope, diags) =
-        parse_param_and_compile(Param { file: file.clone() }, None).unwrap();
+    let (program, prog_scope, diags) = parse_param_and_compile(
+        Param { file: file.clone() },
+        Some(Arc::new(RwLock::new(Default::default()))),
+    )
+    .unwrap();
     (file, program, prog_scope, diags)
 }
 

kclvm/tools/src/LSP/src/util.rs

@@ -1,5 +1,5 @@
 use std::cell::RefCell;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::rc::Rc;
 use std::{fs, sync::Arc};
 
@@ -15,6 +15,7 @@ use kclvm_driver::kpm_metadata::fetch_metadata;
 use kclvm_driver::{get_kcl_files, lookup_compile_unit};
 use kclvm_error::Diagnostic;
 use kclvm_error::Position as KCLPos;
+use kclvm_parser::entry::get_dir_files;
 use kclvm_parser::{load_program, ParseSession};
 use kclvm_sema::resolver::scope::Scope;
 use kclvm_sema::resolver::{resolve_program, scope::ProgramScope};
@@ -114,10 +115,24 @@ fn load_files_code_from_vfs(files: &[&str], vfs: Arc<RwLock<Vfs>>) -> anyhow::Re
             }
             None => {
                 // In order to ensure that k_file corresponds to k_code, load the code from the file system if not exist
-                res.push(
-                    fs::read_to_string(path)
-                        .map_err(|_| anyhow::anyhow!("can't convert file to url: {}", file))?,
-                );
+                let p: &Path = path.as_ref();
+                if p.is_file() {
+                    res.push(
+                        fs::read_to_string(path)
+                            .map_err(|_| anyhow::anyhow!("can't convert file to url: {}", file))?,
+                    );
+                } else if p.is_dir() {
+                    let k_files = get_dir_files(p.to_str().unwrap())
+                        .map_err(|_| anyhow::anyhow!("can't get dir files: {} ", file))?;
+                    for k_file in k_files {
+                        let k_file_path = Path::new(k_file.as_str());
+                        res.push(
+                            fs::read_to_string(k_file_path).map_err(|_| {
+                                anyhow::anyhow!("can't convert file to url: {}", file)
+                            })?,
+                        );
+                    }
+                }
             }
         }
     }