Skip to content

Commit

Permalink
feat: use Arc<RwLock<Module>> in program to reduce deep copy when com…
Browse files Browse the repository at this point in the history 
…piling

Signed-off-by: he1pa <[email protected]>
  • Loading branch information
@He1pa
He1pa committed Oct 31, 2024
1 parent fcf48f0 commit cfd8b8a
Show file tree
Hide file tree
Showing 36 changed files with 357 additions and 147 deletions.
73 changes: 55 additions & 18 deletions kclvm/ast/src/ast.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,10 @@

use kclvm_utils::path::PathPrefix;
use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
use std::collections::HashMap;
use std::{
collections::HashMap,
sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
};

use compiler_base_span::{Loc, Span};
use std::fmt::Debug;
Expand Down Expand Up @@ -383,63 +386,97 @@ pub struct SerializeProgram {
impl Into<SerializeProgram> for Program {
fn into(self) -> SerializeProgram {
SerializeProgram {
root: self.root,
root: self.root.clone(),
pkgs: self
.pkgs
.iter()
.map(|(name, modules)| {
(
name.clone(),
modules.iter().map(|m| m.as_ref().clone()).collect(),
modules
.iter()
.map(|m| {
self.get_module(m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect(),
)
})
.collect(),
}
}
}

#[derive(Debug, Clone, Default, PartialEq)]
#[derive(Debug, Clone, Default)]
pub struct Program {
pub root: String,
pub pkgs: HashMap<String, Vec<Arc<Module>>>,
pub pkgs: HashMap<String, Vec<String>>,
pub modules: HashMap<String, Arc<RwLock<Module>>>,
}

impl Program {
/// Get main entry files.
pub fn get_main_files(&self) -> Vec<String> {
match self.pkgs.get(crate::MAIN_PKG) {
Some(modules) => modules.iter().map(|m| m.filename.clone()).collect(),
Some(modules) => modules.clone(),
None => vec![],
}
}
/// Get the first module in the main package.
pub fn get_main_package_first_module(&self) -> Option<Arc<Module>> {
pub fn get_main_package_first_module(&self) -> Option<RwLockReadGuard<'_, Module>> {
match self.pkgs.get(crate::MAIN_PKG) {
Some(modules) => modules.first().cloned(),
Some(modules) => match modules.first() {
Some(first_module_path) => self.get_module(&first_module_path).unwrap_or(None),
None => None,
},
None => None,
}
}
/// Get stmt on position
pub fn pos_to_stmt(&self, pos: &Position) -> Option<Node<Stmt>> {
for (_, v) in &self.pkgs {
for m in v {
if m.filename == pos.filename {
return m.pos_to_stmt(pos);
if let Ok(m) = self.get_module(m) {
let m = m?;
if m.filename == pos.filename {
return m.pos_to_stmt(pos);
}
}
}
}
None
}

pub fn get_module(&self, module_name: &str) -> Option<&Module> {
for (_, modules) in &self.pkgs {
for module in modules {
if module.filename == module_name {
return Some(module);
}
}
pub fn get_module(
&self,
module_path: &str,
) -> Result<Option<RwLockReadGuard<'_, Module>>, &str> {
match self.modules.get(module_path) {
Some(module_ref) => match module_ref.read() {
Ok(m) => Ok(Some(m)),
Err(_) => Err("Failed to acquire module lock"),
},
None => Ok(None),
}
None
}

pub fn get_mut_module(
&self,
module_path: &str,
) -> Result<Option<RwLockWriteGuard<'_, Module>>, &str> {
match self.modules.get(module_path) {
Some(module_ref) => match module_ref.write() {
Ok(m) => Ok(Some(m)),
Err(_) => Err("Failed to acquire module lock"),
},
None => Ok(None),
}
}

pub fn get_module_ref(&self, module_path: &str) -> Option<Arc<RwLock<Module>>> {
self.modules.get(module_path).cloned()
}
}

Expand Down
16 changes: 14 additions & 2 deletions kclvm/evaluator/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,13 @@ impl<'ctx> Evaluator<'ctx> {
self.init_scope(kclvm_ast::MAIN_PKG);
let modules: Vec<Module> = modules
.iter()
.map(|arc| arc.clone().as_ref().clone())
.map(|m| {
self.program
.get_module(m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect();
self.compile_ast_modules(&modules);
}
Expand All @@ -171,7 +177,13 @@ impl<'ctx> Evaluator<'ctx> {
self.init_scope(kclvm_ast::MAIN_PKG);
let modules: Vec<Module> = modules
.iter()
.map(|arc| arc.clone().as_ref().clone())
.map(|m| {
self.program
.get_module(m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect();
self.compile_ast_modules(&modules)
} else {
Expand Down
13 changes: 12 additions & 1 deletion kclvm/evaluator/src/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,13 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
self.init_scope(&pkgpath);
let modules: Vec<Module> = modules
.iter()
.map(|arc| arc.clone().as_ref().clone())
.map(|m| {
self.program
.get_module(&m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect();
self.compile_ast_modules(&modules);
self.pop_pkgpath();
Expand Down Expand Up @@ -1158,6 +1164,11 @@ impl<'ctx> Evaluator<'ctx> {
.get(&pkgpath_without_prefix!(frame.pkgpath))
{
if let Some(module) = module_list.get(*index) {
let module = self
.program
.get_module(module)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", module));
if let Some(stmt) = module.body.get(setter.stmt) {
self.push_backtrack_meta(setter);
self.walk_stmt(stmt).expect(INTERNAL_ERROR_MSG);
Expand Down
16 changes: 14 additions & 2 deletions kclvm/evaluator/src/scope.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,13 @@ impl<'ctx> Evaluator<'ctx> {
let modules = self.program.pkgs.get(pkgpath).expect(&msg);
let modules: Vec<Module> = modules
.iter()
.map(|arc| arc.clone().as_ref().clone())
.map(|m| {
self.program
.get_module(m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect();
modules
} else if pkgpath.starts_with(kclvm_runtime::PKG_PATH_PREFIX)
Expand All @@ -56,7 +62,13 @@ impl<'ctx> Evaluator<'ctx> {
.expect(kcl_error::INTERNAL_ERROR_MSG);
let modules: Vec<Module> = modules
.iter()
.map(|arc| arc.clone().as_ref().clone())
.map(|m| {
self.program
.get_module(m)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", m))
.clone()
})
.collect();
modules
} else {
Expand Down
7 changes: 6 additions & 1 deletion kclvm/loader/src/option.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,12 @@ pub fn list_options(opts: &LoadPackageOptions) -> Result<Vec<OptionHelp>> {
for (pkgpath, modules) in &packages.program.pkgs {
extractor.pkgpath = pkgpath.clone();
for module in modules {
extractor.walk_module(module)
let module = packages
.program
.get_module(module)
.expect("Failed to acquire module lock")
.expect(&format!("module {:?} not found in program", module));
extractor.walk_module(&module)
}
}
Ok(extractor.options)
Expand Down
45 changes: 22 additions & 23 deletions kclvm/parser/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ use parser::Parser;
use std::collections::{HashMap, HashSet, VecDeque};
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::Instant;

use kclvm_span::create_session_globals_then;

Expand Down Expand Up @@ -119,7 +118,7 @@ pub fn parse_single_file(filename: &str, code: Option<String>) -> Result<ParseFi
);
let result = loader.load_main()?;
let module = match result.program.get_main_package_first_module() {
Some(module) => module.as_ref().clone(),
Some(module) => module.clone(),
None => ast::Module::default(),
};
let file_graph = match loader.file_graph.read() {
Expand Down Expand Up @@ -318,7 +317,7 @@ pub type KCLModuleCache = Arc<RwLock<ModuleCache>>;
#[derive(Default, Debug)]
pub struct ModuleCache {
/// File ast cache
pub ast_cache: IndexMap<PathBuf, Arc<ast::Module>>,
pub ast_cache: IndexMap<PathBuf, Arc<RwLock<ast::Module>>>,
/// Which pkgs the file belongs to. Sometimes a file is not only contained in the pkg in the file system directory, but may also be in the main package.
pub file_pkg: IndexMap<PathBuf, HashSet<PkgFile>>,
/// File dependency cache
Expand Down Expand Up @@ -683,7 +682,7 @@ pub fn parse_file(
file: PkgFile,
src: Option<String>,
module_cache: KCLModuleCache,
pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
pkgs: &mut HashMap<String, Vec<String>>,
pkgmap: &mut PkgMap,
file_graph: FileGraphCache,
opts: &LoadProgramOptions,
Expand All @@ -696,20 +695,15 @@ pub fn parse_file(
}
.cloned(),
};
let m = Arc::new(parse_file_with_session(
sess.clone(),
file.path.to_str().unwrap(),
src,
)?);

let deps = get_deps(&file, m.as_ref(), pkgs, pkgmap, opts, sess)?;
let m = parse_file_with_session(sess.clone(), file.path.to_str().unwrap(), src)?;
let deps = get_deps(&file, &m, pkgs, pkgmap, opts, sess)?;
let dep_files = deps.keys().map(|f| f.clone()).collect();
pkgmap.extend(deps.clone());
match &mut module_cache.write() {
Ok(module_cache) => {
module_cache
.ast_cache
.insert(file.canonicalize(), m.clone());
.insert(file.canonicalize(), Arc::new(RwLock::new(m)));
match module_cache.file_pkg.get_mut(&file.canonicalize()) {
Some(s) => {
s.insert(file.clone());
Expand Down Expand Up @@ -738,7 +732,7 @@ pub fn parse_file(
pub fn get_deps(
file: &PkgFile,
m: &Module,
modules: &mut HashMap<String, Vec<Arc<Module>>>,
modules: &mut HashMap<String, Vec<String>>,
pkgmap: &mut PkgMap,
opts: &LoadProgramOptions,
sess: ParseSessionRef,
Expand Down Expand Up @@ -791,7 +785,7 @@ pub fn parse_pkg(
sess: ParseSessionRef,
files: Vec<(PkgFile, Option<String>)>,
module_cache: KCLModuleCache,
pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
pkgs: &mut HashMap<String, Vec<String>>,
pkgmap: &mut PkgMap,
file_graph: FileGraphCache,
opts: &LoadProgramOptions,
Expand All @@ -817,7 +811,7 @@ pub fn parse_entry(
sess: ParseSessionRef,
entry: &entry::Entry,
module_cache: KCLModuleCache,
pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
pkgs: &mut HashMap<String, Vec<String>>,
pkgmap: &mut PkgMap,
file_graph: FileGraphCache,
opts: &LoadProgramOptions,
Expand Down Expand Up @@ -877,7 +871,8 @@ pub fn parse_entry(
Ok(m_cache) => match m_cache.ast_cache.get(&file.canonicalize()) {
Some(m) => {
let deps = m_cache.dep_cache.get(&file).cloned().unwrap_or_else(|| {
get_deps(&file, m.as_ref(), pkgs, pkgmap, opts, sess.clone()).unwrap()
get_deps(&file, &m.read().unwrap(), pkgs, pkgmap, opts, sess.clone())
.unwrap()
});
let dep_files: Vec<PkgFile> = deps.keys().map(|f| f.clone()).collect();
pkgmap.extend(deps.clone());
Expand Down Expand Up @@ -933,7 +928,7 @@ pub fn parse_program(
) -> Result<LoadProgramResult> {
let compile_entries = get_compile_entries_from_paths(&paths, &opts)?;
let workdir = compile_entries.get_root_path().to_string();
let mut pkgs: HashMap<String, Vec<Arc<Module>>> = HashMap::new();
let mut pkgs: HashMap<String, Vec<String>> = HashMap::new();
let mut pkgmap = PkgMap::new();
let mut new_files = HashSet::new();

Expand Down Expand Up @@ -981,8 +976,11 @@ pub fn parse_program(
}
Err(e) => return Err(anyhow::anyhow!("Parse program failed: {e}")),
};

let mut modules: HashMap<String, Arc<RwLock<Module>>> = HashMap::new();
for file in files.iter() {
let mut m_ref = match module_cache.read() {
let filename = file.canonicalize().to_str().unwrap().to_string();
let m_ref = match module_cache.read() {
Ok(module_cache) => module_cache
.ast_cache
.get(&file.canonicalize())
Expand All @@ -995,23 +993,24 @@ pub fn parse_program(
};
if new_files.contains(file) {
let pkg = pkgmap.get(file).expect("file not in pkgmap");
let mut m = Arc::make_mut(&mut m_ref);
let mut m = m_ref.write().unwrap();
fix_rel_import_path_with_file(&pkg.pkg_root, &mut m, file, &pkgmap, opts, sess.clone());
}

modules.insert(filename.clone(), m_ref);
match pkgs.get_mut(&file.pkg_path) {
Some(modules) => {
modules.push(m_ref);
Some(pkg_modules) => {
pkg_modules.push(filename);
}
None => {
pkgs.insert(file.pkg_path.clone(), vec![m_ref]);
pkgs.insert(file.pkg_path.clone(), vec![filename]);
}
}
}

let program = ast::Program {
root: workdir,
pkgs,
modules,
};

Ok(LoadProgramResult {
Expand Down
Loading

0 comments on commit cfd8b8a

Please sign in to comment.