|
| 1 | +use indexmap::IndexMap; |
| 2 | +use std::{collections::HashMap, path::Path, path::PathBuf, sync::mpsc::channel}; |
| 3 | +use threadpool::ThreadPool; |
| 4 | + |
| 5 | +use kclvm_ast::ast::{self, Program}; |
| 6 | +use kclvm_compiler::codegen::{llvm::emit_code, EmitOptions}; |
| 7 | +use kclvm_config::cache::{load_pkg_cache, save_pkg_cache, CacheOption}; |
| 8 | +use kclvm_sema::resolver::resolve_program; |
| 9 | +use kclvm_sema::resolver::scope::ProgramScope; |
| 10 | + |
| 11 | +use crate::{ |
| 12 | + command::Command, |
| 13 | + runner::{ExecProgramArgs, KclvmRunner, KclvmRunnerOptions}, |
| 14 | +}; |
| 15 | + |
| 16 | +const LL_FILE: &str = "_a.out"; |
| 17 | + |
| 18 | +/// Evaluator is used to resolve kcl ast, generate dylibs and execute. |
| 19 | +pub struct Evaluator {} |
| 20 | + |
| 21 | +impl Evaluator { |
| 22 | + pub fn new() -> Self { |
| 23 | + Self {} |
| 24 | + } |
| 25 | + |
| 26 | + /// Take ast.program as input,and resolve ast, generate and link dylib, execute. |
| 27 | + pub fn eval( |
| 28 | + &self, |
| 29 | + mut program: Program, |
| 30 | + plugin_agent: u64, |
| 31 | + args: &ExecProgramArgs, |
| 32 | + ) -> Result<String, String> { |
| 33 | + // resolve ast |
| 34 | + let scope = resolve_program(&mut program); |
| 35 | + scope.check_scope_diagnostics(); |
| 36 | + |
| 37 | + // generate dylibs |
| 38 | + let dylib_paths = self.gen_dylibs(program, scope, plugin_agent); |
| 39 | + |
| 40 | + // link dylibs |
| 41 | + let dylib_path = self.link_all_dylibs(dylib_paths, plugin_agent); |
| 42 | + |
| 43 | + // execute |
| 44 | + self.run_dylib(dylib_path, plugin_agent, args) |
| 45 | + } |
| 46 | + |
| 47 | + /// Generate the dylibs and return file paths. |
| 48 | + fn gen_dylibs( |
| 49 | + &self, |
| 50 | + program: ast::Program, |
| 51 | + scope: ProgramScope, |
| 52 | + plugin_agent: u64, |
| 53 | + ) -> Vec<String> { |
| 54 | + // gen bc or ll_file |
| 55 | + let path = std::path::Path::new(LL_FILE); |
| 56 | + if path.exists() { |
| 57 | + std::fs::remove_file(path).unwrap(); |
| 58 | + } |
| 59 | + for entry in glob::glob(&format!("{}*.ll", LL_FILE)).unwrap() { |
| 60 | + match entry { |
| 61 | + Ok(path) => { |
| 62 | + if path.exists() { |
| 63 | + std::fs::remove_file(path).unwrap(); |
| 64 | + } |
| 65 | + } |
| 66 | + Err(e) => println!("{:?}", e), |
| 67 | + }; |
| 68 | + } |
| 69 | + |
| 70 | + let cache_dir = Path::new(&program.root) |
| 71 | + .join(".kclvm") |
| 72 | + .join("cache") |
| 73 | + .join(kclvm_version::get_full_version()); |
| 74 | + if !cache_dir.exists() { |
| 75 | + std::fs::create_dir_all(&cache_dir).unwrap(); |
| 76 | + } |
| 77 | + let mut compile_progs: IndexMap< |
| 78 | + String, |
| 79 | + ( |
| 80 | + ast::Program, |
| 81 | + IndexMap<String, IndexMap<String, String>>, |
| 82 | + PathBuf, |
| 83 | + ), |
| 84 | + > = IndexMap::default(); |
| 85 | + for (pkgpath, modules) in program.pkgs { |
| 86 | + let mut pkgs = HashMap::new(); |
| 87 | + pkgs.insert(pkgpath.clone(), modules); |
| 88 | + let compile_prog = ast::Program { |
| 89 | + root: program.root.clone(), |
| 90 | + main: program.main.clone(), |
| 91 | + pkgs, |
| 92 | + cmd_args: vec![], |
| 93 | + cmd_overrides: vec![], |
| 94 | + }; |
| 95 | + compile_progs.insert( |
| 96 | + pkgpath, |
| 97 | + (compile_prog, scope.import_names.clone(), cache_dir.clone()), |
| 98 | + ); |
| 99 | + } |
| 100 | + let pool = ThreadPool::new(4); |
| 101 | + let (tx, rx) = channel(); |
| 102 | + let prog_count = compile_progs.len(); |
| 103 | + for (pkgpath, (compile_prog, import_names, cache_dir)) in compile_progs { |
| 104 | + let tx = tx.clone(); |
| 105 | + pool.execute(move || { |
| 106 | + let root = &compile_prog.root; |
| 107 | + let is_main_pkg = pkgpath == kclvm_ast::MAIN_PKG; |
| 108 | + let file = if is_main_pkg { |
| 109 | + PathBuf::from(&pkgpath) |
| 110 | + } else { |
| 111 | + cache_dir.join(&pkgpath) |
| 112 | + }; |
| 113 | + let ll_file = file.to_str().unwrap(); |
| 114 | + let ll_path = format!("{}.ll", ll_file); |
| 115 | + let dylib_path = format!("{}{}", ll_file, Command::get_lib_suffix()); |
| 116 | + let mut ll_path_lock = |
| 117 | + fslock::LockFile::open(&format!("{}.lock", ll_path)).unwrap(); |
| 118 | + ll_path_lock.lock().unwrap(); |
| 119 | + if Path::new(&ll_path).exists() { |
| 120 | + std::fs::remove_file(&ll_path).unwrap(); |
| 121 | + } |
| 122 | + let dylib_path = if is_main_pkg { |
| 123 | + emit_code( |
| 124 | + &compile_prog, |
| 125 | + import_names, |
| 126 | + &EmitOptions { |
| 127 | + from_path: None, |
| 128 | + emit_path: Some(&ll_file), |
| 129 | + no_link: true, |
| 130 | + }, |
| 131 | + ) |
| 132 | + .expect("Compile KCL to LLVM error"); |
| 133 | + let mut cmd = Command::new(plugin_agent); |
| 134 | + cmd.run_clang_single(&ll_path, &dylib_path) |
| 135 | + } else { |
| 136 | + // If AST module has been modified, ignore the dylib cache |
| 137 | + let dylib_relative_path: Option<String> = |
| 138 | + load_pkg_cache(root, &pkgpath, CacheOption::default()); |
| 139 | + match dylib_relative_path { |
| 140 | + Some(dylib_relative_path) => { |
| 141 | + if dylib_relative_path.starts_with('.') { |
| 142 | + dylib_relative_path.replacen(".", root, 1) |
| 143 | + } else { |
| 144 | + dylib_relative_path |
| 145 | + } |
| 146 | + } |
| 147 | + None => { |
| 148 | + emit_code( |
| 149 | + &compile_prog, |
| 150 | + import_names, |
| 151 | + &EmitOptions { |
| 152 | + from_path: None, |
| 153 | + emit_path: Some(&ll_file), |
| 154 | + no_link: true, |
| 155 | + }, |
| 156 | + ) |
| 157 | + .expect("Compile KCL to LLVM error"); |
| 158 | + let mut cmd = Command::new(plugin_agent); |
| 159 | + let dylib_path = cmd.run_clang_single(&ll_path, &dylib_path); |
| 160 | + let dylib_relative_path = dylib_path.replacen(root, ".", 1); |
| 161 | + |
| 162 | + save_pkg_cache( |
| 163 | + root, |
| 164 | + &pkgpath, |
| 165 | + dylib_relative_path, |
| 166 | + CacheOption::default(), |
| 167 | + ); |
| 168 | + dylib_path |
| 169 | + } |
| 170 | + } |
| 171 | + }; |
| 172 | + if Path::new(&ll_path).exists() { |
| 173 | + std::fs::remove_file(&ll_path).unwrap(); |
| 174 | + } |
| 175 | + ll_path_lock.unlock().unwrap(); |
| 176 | + tx.send(dylib_path) |
| 177 | + .expect("channel will be there waiting for the pool"); |
| 178 | + }); |
| 179 | + } |
| 180 | + rx.iter().take(prog_count).collect::<Vec<String>>() |
| 181 | + } |
| 182 | + |
| 183 | + /// Link the dylibs generated by method "gen_bc_or_ll_file". |
| 184 | + fn link_all_dylibs(&self, dylib_paths: Vec<String>, plugin_agent: u64) -> String { |
| 185 | + let mut cmd = Command::new(plugin_agent); |
| 186 | + cmd.link_dylibs(&dylib_paths, "") |
| 187 | + } |
| 188 | + |
| 189 | + /// Execute the dylibs linked by method "link_all_dylibs". |
| 190 | + fn run_dylib( |
| 191 | + &self, |
| 192 | + dylib_path: String, |
| 193 | + plugin_agent: u64, |
| 194 | + args: &ExecProgramArgs, |
| 195 | + ) -> Result<String, String> { |
| 196 | + let runner = KclvmRunner::new( |
| 197 | + dylib_path.as_str(), |
| 198 | + Some(KclvmRunnerOptions { |
| 199 | + plugin_agent_ptr: plugin_agent, |
| 200 | + }), |
| 201 | + ); |
| 202 | + runner.run(&args) |
| 203 | + } |
| 204 | +} |
0 commit comments