diff --git a/kclvm/ast/src/ast.rs b/kclvm/ast/src/ast.rs index 5e723ff46..999a75029 100644 --- a/kclvm/ast/src/ast.rs +++ b/kclvm/ast/src/ast.rs @@ -166,21 +166,15 @@ impl TryInto> for Node { /// AST node type T pub type NodeRef = Box>; -#[derive(Serialize, Deserialize, Debug, Clone)] -pub enum ParseMode { - Null, - ParseComments, -} - /// KCL command line argument spec, e.g. `kcl main.k -D name=value` -#[derive(Serialize, Deserialize, Debug, Clone)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] pub struct CmdArgSpec { pub name: String, pub value: String, } /// KCL command line override spec, e.g. `kcl main.k -O pkgpath:path.to.field=field_value` -#[derive(Serialize, Deserialize, Debug, Clone)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] pub struct OverrideSpec { pub pkgpath: String, pub field_path: String, @@ -188,7 +182,7 @@ pub struct OverrideSpec { pub action: OverrideAction, } -#[derive(Serialize, Deserialize, Debug, Clone)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] pub enum OverrideAction { CreateOrUpdate, Delete, diff --git a/kclvm/tools/src/query/mod.rs b/kclvm/tools/src/query/mod.rs index 2420ea5ec..39089db11 100644 --- a/kclvm/tools/src/query/mod.rs +++ b/kclvm/tools/src/query/mod.rs @@ -1,11 +1,11 @@ -//! This package is mainly the implementation of the KCL query tool, mainly including -//! KCL code modification `override` and other implementations. We can call the `override_file` -//! function to modify the file. The main principle is to parse the AST according to the -//! input file name, and according to the ast: :OverrideSpec transforms the nodes in the +//! This package is mainly the implementation of the KCL query tool, mainly including +//! KCL code modification `override` and other implementations. We can call the `override_file` +//! function to modify the file. The main principle is to parse the AST according to the +//! input file name, and according to the ast::OverrideSpec transforms the nodes in the //! AST, recursively modifying or deleting the values of the nodes in the AST. pub mod r#override; #[cfg(test)] mod tests; -pub use r#override::{apply_overrides, override_file, spec_str_to_override}; +pub use r#override::{apply_overrides, override_file, parse_override_spec, split_field_path}; diff --git a/kclvm/tools/src/query/override.rs b/kclvm/tools/src/query/override.rs index 7a9d3de86..e930bf7f1 100644 --- a/kclvm/tools/src/query/override.rs +++ b/kclvm/tools/src/query/override.rs @@ -27,10 +27,44 @@ use crate::printer::print_ast_module; /// /// result: [Result] /// Whether override is successful +/// +/// # Examples +/// +/// ```no_run +/// use kclvm_tools::query::override_file; +/// +/// let result = override_file( +/// "test.k", +/// &["alice.age=18".to_string()], +/// &[] +/// ).unwrap(); +/// ``` +/// +/// - test.k (before override) +/// +/// ```kcl +/// schema Person: +/// age: int +/// +/// alice = Person { +/// age = 10 +/// } +/// ``` +/// +/// - test.k (after override) +/// +/// ```kcl +/// schema Person: +/// age: int +/// +/// alice = Person { +/// age = 18 +/// } +/// ``` pub fn override_file(file: &str, specs: &[String], import_paths: &[String]) -> Result { let overrides = specs .iter() - .map(|s| spec_str_to_override(s)) + .map(|s| parse_override_spec(s)) .filter_map(Result::ok) .collect::>(); let mut module = match parse_file(file, None) { @@ -51,58 +85,81 @@ pub fn override_file(file: &str, specs: &[String], import_paths: &[String]) -> R } /// Override spec string to override structure -pub fn spec_str_to_override(spec: &str) -> Result { - let err = Err(anyhow!("Invalid spec format '{}', expected := or :-", spec)); +/// +/// # Example +/// +/// ``` +/// use kclvm_tools::query::parse_override_spec; +/// use kclvm_ast::ast; +/// +/// let spec = parse_override_spec("alice.age=10").unwrap(); +/// assert_eq!(spec, ast::OverrideSpec { +/// pkgpath: "".to_string(), +/// field_path: "alice.age".to_string(), +/// field_value: "10".to_string(), +/// action: ast::OverrideAction::CreateOrUpdate, +/// }); +/// ``` +pub fn parse_override_spec(spec: &str) -> Result { + let err = Err(invalid_spec_error(spec)); if spec.contains('=') { // Create or update the override value. let split_values = spec.splitn(2, '=').collect::>(); - let paths = split_values[0].splitn(2, ':').collect::>(); - if split_values.len() < 2 || paths.len() > 2 { - err - } else { - let (pkgpath, field_path) = if paths.len() == 1 { - ("".to_string(), paths[0].to_string()) - } else { - (paths[0].to_string(), paths[1].to_string()) - }; - if field_path.is_empty() || split_values[1].is_empty() { - err - } else { - Ok(ast::OverrideSpec { - pkgpath, - field_path, - field_value: split_values[1].to_string(), - action: ast::OverrideAction::CreateOrUpdate, - }) - } - } + let path = split_values + .get(0) + .ok_or_else(|| invalid_spec_error(spec))?; + let field_value = split_values + .get(1) + .ok_or_else(|| invalid_spec_error(spec))?; + let (pkgpath, field_path) = split_field_path(path)?; + Ok(ast::OverrideSpec { + pkgpath, + field_path, + field_value: field_value.to_string(), + action: ast::OverrideAction::CreateOrUpdate, + }) } else if let Some(stripped_spec) = spec.strip_suffix('-') { // Delete the override value. - let paths = stripped_spec.splitn(2, ':').collect::>(); - if paths.len() > 2 { - err - } else { - let (pkgpath, field_path) = if paths.len() == 1 { - ("".to_string(), paths[0].to_string()) - } else { - (paths[0].to_string(), paths[1].to_string()) - }; - if field_path.is_empty() { - err - } else { - Ok(ast::OverrideSpec { - pkgpath, - field_path, - field_value: "".to_string(), - action: ast::OverrideAction::Delete, - }) - } - } + let (pkgpath, field_path) = split_field_path(stripped_spec)?; + Ok(ast::OverrideSpec { + pkgpath, + field_path, + field_value: "".to_string(), + action: ast::OverrideAction::Delete, + }) } else { err } } +/// Get field package path and identifier name from the path. +/// +/// # Example +/// +/// ``` +/// use kclvm_tools::query::split_field_path; +/// +/// let (pkgpath, field) = split_field_path("pkg.to.path:field").unwrap(); +/// assert_eq!(pkgpath, "pkg.to.path"); +/// assert_eq!(field, "field"); +/// ``` +pub fn split_field_path(path: &str) -> Result<(String, String)> { + let err = Err(anyhow!("Invalid field path {:?}", path)); + let paths = path.splitn(2, ':').collect::>(); + let (pkgpath, field_path) = if paths.len() == 1 { + ("".to_string(), paths[0].to_string()) + } else if paths.len() == 2 { + (paths[0].to_string(), paths[1].to_string()) + } else { + return err; + }; + if field_path.is_empty() { + err + } else { + Ok((pkgpath, field_path)) + } +} + /// Apply overrides on the AST program with the override specifications. pub fn apply_overrides( prog: &mut ast::Program, @@ -148,7 +205,7 @@ pub fn apply_override_on_module( let name = path .split('.') .last() - .ok_or_else(|| anyhow!("invalid import path {}", path))?; + .ok_or_else(|| anyhow!("Invalid import path {}", path))?; let import_node = ast::ImportStmt { path: path.to_string(), rawpath: "".to_string(), @@ -161,7 +218,7 @@ pub fn apply_override_on_module( line, 1, line, - 7 + path.len() as u64, + ("import ".len() + path.len()) as u64, )); m.body.insert((line - 1) as usize, import_stmt) } @@ -194,6 +251,12 @@ pub fn apply_override_on_module( } } +/// Get the invalid spec error message. +#[inline] +fn invalid_spec_error(spec: &str) -> anyhow::Error { + anyhow!("Invalid spec format '{}', expected := or :-", spec) +} + /// Build a expression from string fn build_expr_from_string(value: &str) -> ast::NodeRef { if value.is_empty() { @@ -225,7 +288,10 @@ fn fix_multi_assign(m: &mut ast::Module) { transformer.walk_module(m); for (offset, (index, assign_stmt)) in transformer.multi_assign_mapping.iter().enumerate() { let insert_index = index + offset; - let pos = m.body[insert_index].pos().clone(); + let pos = match m.body.get(insert_index) { + Some(stmt) => stmt.pos().clone(), + None => bug!("AST module body index {} out of bound", insert_index), + }; m.body.insert( insert_index, Box::new(ast::Node::node_with_pos( @@ -280,7 +346,14 @@ struct OverrideTransformer { impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer { fn walk_unification_stmt(&mut self, unification_stmt: &'ctx mut ast::UnificationStmt) { - if unification_stmt.target.node.names[0] != self.target_id { + let name = match unification_stmt.target.node.names.get(0) { + Some(name) => name, + None => bug!( + "Invalid AST unification target names {:?}", + unification_stmt.target.node.names + ), + }; + if name != &self.target_id { return; } self.override_target_count = 1; @@ -313,7 +386,7 @@ impl<'ctx> MutSelfMutWalker<'ctx> for OverrideTransformer { if self.override_target_count == 0 { return; } - if !self.find_schema_config_and_repalce(schema_expr) { + if !self.lookup_schema_config_and_repalce(schema_expr) { // Not exist and append an override value when the action is CREATE_OR_UPDATE if let ast::OverrideAction::CreateOrUpdate = self.action { if let ast::Expr::Config(config_expr) = &mut schema_expr.config.node { @@ -441,133 +514,105 @@ impl OverrideTransformer { } #[derive(Debug)] -enum OverrideConfig<'a> { +enum OverrideNode<'a> { Schema(&'a mut ast::SchemaExpr), Config(&'a mut ast::ConfigExpr), } impl OverrideTransformer { - pub(crate) fn find_schema_config_and_repalce( + /// Lookup schema config all fields and replace if it is matched with the override spec, + /// return whether is found a replaced one. + pub(crate) fn lookup_schema_config_and_repalce( &mut self, schema_expr: &mut ast::SchemaExpr, ) -> bool { let (paths, paths_with_id) = self.get_schema_config_field_paths(schema_expr); match paths.iter().position(|r| r == &self.field_path) { Some(pos) => { - let mut config = OverrideConfig::Schema(schema_expr); - self.replace_with_id_path(&mut config, &paths_with_id[pos]); + let mut config = OverrideNode::Schema(schema_expr); + self.replace_node_with_path(&mut config, &paths_with_id[pos]); true } None => false, } } - pub(crate) fn replace_with_id_path(&mut self, config: &mut OverrideConfig, path_with_id: &str) { - if path_with_id.is_empty() { + /// Replace AST node with path including schema_expr and config_expr. + pub(crate) fn replace_node_with_path(&mut self, config: &mut OverrideNode, path: &str) { + // Do not replace empty path parts on the config expression. + if path.is_empty() { return; } - let parts = path_with_id.split('|').collect::>(); - for (i, part) in parts.iter().enumerate() { + // Split a path into multiple parts. + let parts = path.split('|').collect::>(); + for i in 0..parts.len() { match config { - OverrideConfig::Schema(schema_expr) => { + OverrideNode::Schema(schema_expr) => { if let ast::Expr::Config(config_expr) = &mut schema_expr.config.node { - let mut delete_index_set = HashSet::new(); - for (j, item) in config_expr.items.iter_mut().enumerate() { - let path = self.get_path_from_key(&item.node.key); - if &path == part { - match self.action { - ast::OverrideAction::CreateOrUpdate => { - if i == parts.len() - 1 { - self.override_value.set_pos(item.pos()); - item.node.value = self.override_value.clone(); - } - let path_with_id = &parts[i + 1..].join("|"); - match &mut item.node.value.node { - ast::Expr::Schema(schema_expr) => { - let mut config = - OverrideConfig::Schema(schema_expr); - self.replace_with_id_path( - &mut config, - path_with_id, - ); - } - ast::Expr::Config(config_expr) => { - let mut config = - OverrideConfig::Config(config_expr); - self.replace_with_id_path( - &mut config, - path_with_id, - ); - } - _ => {} - } - } - ast::OverrideAction::Delete => { - delete_index_set.insert(j); - } - } - } - } - if !delete_index_set.is_empty() { - let items: Vec<(usize, &ast::NodeRef)> = config_expr - .items - .iter() - .enumerate() - .filter(|(i, _)| !delete_index_set.contains(i)) - .collect(); - config_expr.items = items - .iter() - .map(|(_, item)| { - <&ast::NodeRef>::clone(item).clone() - }) - .collect(); - } + self.replace_config_expr_with_path_parts(config_expr, &parts, i); } } - OverrideConfig::Config(config_expr) => { - let mut delete_index_set = HashSet::new(); - for (j, item) in config_expr.items.iter_mut().enumerate() { - let path = self.get_path_from_key(&item.node.key); - if &path == part { - match self.action { - ast::OverrideAction::CreateOrUpdate => { - if i == parts.len() - 1 && parts.len() == 1 { - self.override_value.set_pos(item.pos()); - item.node.value = self.override_value.clone(); - } - let path_with_id = &parts[i + 1..].join("|"); - match &mut item.node.value.node { - ast::Expr::Schema(schema_expr) => { - let mut config = OverrideConfig::Schema(schema_expr); - self.replace_with_id_path(&mut config, path_with_id); - } - ast::Expr::Config(config_expr) => { - let mut config = OverrideConfig::Config(config_expr); - self.replace_with_id_path(&mut config, path_with_id); - } - _ => {} - } - } - ast::OverrideAction::Delete => { - delete_index_set.insert(j); - } + OverrideNode::Config(config_expr) => { + self.replace_config_expr_with_path_parts(config_expr, &parts, i); + } + } + } + } + + /// Replace AST config expr with one part of path. + pub(crate) fn replace_config_expr_with_path_parts( + &mut self, + config_expr: &mut ast::ConfigExpr, + parts: &[&str], + start_index: usize, + ) { + // Do not replace empty path parts on the config expression. + if parts.is_empty() { + return; + } + let part = parts[start_index]; + let mut delete_index_set = HashSet::new(); + for (j, item) in config_expr.items.iter_mut().enumerate() { + let path = self.get_path_from_key(&item.node.key); + if &path == part { + match self.action { + ast::OverrideAction::CreateOrUpdate => { + // The last part in path. + if start_index == 0 && parts.len() == 1 { + self.override_value.set_pos(item.pos()); + item.node.value = self.override_value.clone(); + } + let path = &parts[start_index + 1..].join("|"); + // Replace value recursively using the path. + match &mut item.node.value.node { + ast::Expr::Schema(schema_expr) => { + let mut config = OverrideNode::Schema(schema_expr); + self.replace_node_with_path(&mut config, path); + } + ast::Expr::Config(config_expr) => { + let mut config = OverrideNode::Config(config_expr); + self.replace_node_with_path(&mut config, path); } + _ => {} } } - if !delete_index_set.is_empty() { - let items: Vec<(usize, &ast::NodeRef)> = config_expr - .items - .iter() - .enumerate() - .filter(|(i, _)| !delete_index_set.contains(i)) - .collect(); - config_expr.items = items - .iter() - .map(|(_, item)| <&ast::NodeRef>::clone(item).clone()) - .collect(); + ast::OverrideAction::Delete => { + delete_index_set.insert(j); } } } } + if !delete_index_set.is_empty() { + let items: Vec<(usize, &ast::NodeRef)> = config_expr + .items + .iter() + .enumerate() + .filter(|(i, _)| !delete_index_set.contains(i)) + .collect(); + config_expr.items = items + .iter() + .map(|(_, item)| <&ast::NodeRef>::clone(item).clone()) + .collect(); + } } } diff --git a/kclvm/tools/src/query/tests.rs b/kclvm/tools/src/query/tests.rs index bd2b5926b..12e918e5a 100644 --- a/kclvm/tools/src/query/tests.rs +++ b/kclvm/tools/src/query/tests.rs @@ -49,7 +49,7 @@ fn test_override_file_config() { ]; let overrides = specs .iter() - .map(|s| spec_str_to_override(s)) + .map(|s| parse_override_spec(s)) .filter_map(Result::ok) .collect::>(); let import_paths = vec![]; @@ -107,9 +107,9 @@ appConfigurationUnification: AppConfiguration { } #[test] -fn test_spec_str_to_override_invalid() { +fn test_parse_override_spec_invalid() { let specs = vec![":a:", "=a=", ":a", "a-1"]; for spec in specs { - assert!(spec_str_to_override(spec).is_err(), "{} test failed", spec); + assert!(parse_override_spec(spec).is_err(), "{} test failed", spec); } }