diff --git a/Cargo.lock b/Cargo.lock index 931d31e99..c3c354d17 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2731,49 +2731,66 @@ dependencies = [ "portable-pty", "rayon", "regex", + "serde", + "serde_json", "similar", "smallvec", "syntect", "tokio", "tokio-util", "tree-sitter", + "tree-sitter-astro-next", "tree-sitter-bash", "tree-sitter-c", "tree-sitter-c-sharp", "tree-sitter-clojure", + "tree-sitter-cmake", "tree-sitter-cpp", "tree-sitter-css", "tree-sitter-diff", + "tree-sitter-dockerfile-updated", "tree-sitter-elixir", + "tree-sitter-erlang", "tree-sitter-glimmer", "tree-sitter-go", + "tree-sitter-graphql", "tree-sitter-haskell", "tree-sitter-hcl", "tree-sitter-html", + "tree-sitter-ini", "tree-sitter-java", "tree-sitter-javascript", "tree-sitter-json", "tree-sitter-julia", + "tree-sitter-just", "tree-sitter-kotlin-sg", "tree-sitter-lua", "tree-sitter-make", "tree-sitter-md", "tree-sitter-nix", "tree-sitter-objc", + "tree-sitter-ocaml", "tree-sitter-odin", + "tree-sitter-perl-next", "tree-sitter-php", + "tree-sitter-powershell", + "tree-sitter-proto", "tree-sitter-python", + "tree-sitter-r", "tree-sitter-regex", "tree-sitter-ruby", "tree-sitter-rust", "tree-sitter-scala", + "tree-sitter-sequel", "tree-sitter-solidity", "tree-sitter-starlark", + "tree-sitter-svelte-next", "tree-sitter-swift", "tree-sitter-tlaplus", "tree-sitter-toml-ng", "tree-sitter-typescript", "tree-sitter-verilog", + "tree-sitter-vue-next", "tree-sitter-xml", "tree-sitter-yaml", "tree-sitter-zig", @@ -3836,6 +3853,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-astro-next" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "794a4a59fc2d88e49b4bc41fef9522d77184a36f4e68bbaf545cd1eb2364c46e" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-bash" version = "0.25.1" @@ -3877,6 +3904,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-cmake" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c1b35d1dd7396d24b3e826bb0f975b915ec7e9125b989d5e9d24ebb6a08509a" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-cpp" version = "0.23.4" @@ -3907,6 +3944,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-dockerfile-updated" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e61eb704364be179ab9be89e2891c8e77c0041eca5891a5c04164a21f6df7a0" +dependencies = [ + "cc", + "tree-sitter", +] + [[package]] name = "tree-sitter-elixir" version = "0.3.5" @@ -3917,6 +3964,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-erlang" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2091cce4eda19c03d77928c608ac6617445a6a25691dde1e93ac0102467a6be" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-glimmer" version = "0.0.1" @@ -3937,6 +3994,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-graphql" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efedc4cac157161cc23a0adc4553a2cedc908e1cd754b6cd033a919bb81ce5d6" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-haskell" version = "0.23.1" @@ -3967,6 +4034,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-ini" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "387f79682cd53b7c0a5777c96e601a02b9965a787984ef86dbb8952bdab2d62f" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-java" version = "0.23.5" @@ -4007,6 +4084,17 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-just" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd46794f2c270f43d73b2f2c8c256992471f3493c903604da0886df984771920" +dependencies = [ + "cc", + "tree-sitter", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-kotlin-sg" version = "0.4.0" @@ -4073,6 +4161,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-ocaml" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d19db582b3855f56b5f9ec484170fbfb9ee60b938ec7720d76d2ee788e8b640" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-odin" version = "1.3.0" @@ -4083,6 +4181,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-perl-next" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97b02c571f4d63663ade506b1f81dd9f78e37a0014268c521b6e75df0603d911" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-php" version = "0.24.2" @@ -4093,6 +4201,26 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-powershell" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ff5d1c86d8b1625585fd18e05fff47e38c36564e45ec4f38fd0de4ecbf08e2c" +dependencies = [ + "cc", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-proto" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e410ccb5fa3cbd6bf7b8e512ecf7ad9d5254395b822bfe9f751b50fa978f31c" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-python" version = "0.25.0" @@ -4103,6 +4231,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-r" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "429133cbda9f8a46e03ef3aae6abb6c3d22875f8585cad472138101bfd517255" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-regex" version = "0.25.0" @@ -4143,6 +4281,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-sequel" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d198ad3c319c02e43c21efa1ec796b837afcb96ffaef1a40c1978fbdcec7d17" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-solidity" version = "1.2.13" @@ -4163,6 +4311,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-svelte-next" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f88190d0743e897c3e148a7e241aba0a8844b8afe816943851426e3f7b9753" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-swift" version = "0.7.1" @@ -4213,6 +4371,16 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-vue-next" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9ac89acaa8165aaabfa8ae9683aa423b0d2bf42c815dae5ca031a36a75525f1" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "tree-sitter-xml" version = "0.7.0" diff --git a/crates/pi-natives/Cargo.toml b/crates/pi-natives/Cargo.toml index 70e313f12..d7f50d3eb 100644 --- a/crates/pi-natives/Cargo.toml +++ b/crates/pi-natives/Cargo.toml @@ -40,22 +40,29 @@ ast-grep-core = { version = "0.39", default-features = false, features = [ "tree-sitter", ] } tree-sitter = "0.25" +tree-sitter-astro = { version = "0.1.1", package = "tree-sitter-astro-next" } tree-sitter-bash = "0.25" tree-sitter-c = "0.24" tree-sitter-clojure = "0.1" +tree-sitter-cmake = "0.7.1" tree-sitter-c-sharp = "0.23" tree-sitter-cpp = "0.23" tree-sitter-css = "0.25" tree-sitter-diff = "0.1" +tree-sitter-dockerfile = { version = "0.2.0", package = "tree-sitter-dockerfile-updated" } tree-sitter-elixir = "0.3" +tree-sitter-erlang = "0.15.0" tree-sitter-go = "0.25" +tree-sitter-graphql = "0.1.0" tree-sitter-glimmer = "0.0.1" tree-sitter-haskell = "0.23" tree-sitter-hcl = "1.1" tree-sitter-html = "0.23" +tree-sitter-ini = "1.4.0" tree-sitter-java = "0.23" tree-sitter-javascript = "0.25" tree-sitter-json = "0.23" +tree-sitter-just = "0.2.0" tree-sitter-julia = "0.23" tree-sitter-kotlin = { version = "0.4", package = "tree-sitter-kotlin-sg" } tree-sitter-lua = "0.2" @@ -63,20 +70,28 @@ tree-sitter-make = "1.1" tree-sitter-md = "0.5" tree-sitter-nix = "0.3" tree-sitter-objc = "3.0" +tree-sitter-ocaml = "0.24.2" tree-sitter-odin = "1.3" +tree-sitter-perl = { version = "0.1.0", package = "tree-sitter-perl-next" } tree-sitter-php = "0.24" +tree-sitter-powershell = "0.26.3" +tree-sitter-proto = "0.4.0" tree-sitter-python = "0.25" +tree-sitter-r = "1.2.0" tree-sitter-regex = "0.25" tree-sitter-ruby = "0.23" tree-sitter-rust = "0.24" tree-sitter-scala = "0.24" tree-sitter-solidity = "1.2" +tree-sitter-sql = { version = "0.3.11", package = "tree-sitter-sequel" } tree-sitter-starlark = "1.3" +tree-sitter-svelte = { version = "0.1.1", package = "tree-sitter-svelte-next" } tree-sitter-swift = "0.7" tree-sitter-toml-ng = "0.7" tree-sitter-tlaplus = "1.5" tree-sitter-typescript = "0.23" tree-sitter-verilog = "1.0" +tree-sitter-vue = { version = "0.1.0", package = "tree-sitter-vue-next" } tree-sitter-xml = "0.7" tree-sitter-yaml = "0.7" tree-sitter-zig = "1.1" @@ -111,6 +126,8 @@ regex = "1" similar = "3.0.0" path-clean = "1" pathdiff = "0.2" +serde = "1.0" +serde_json = { version = "1.0", features = ["preserve_order"] } [target.'cfg(unix)'.dependencies] libc = "0.2" diff --git a/crates/pi-natives/src/chunk/ast_astro.rs b/crates/pi-natives/src/chunk/ast_astro.rs new file mode 100644 index 000000000..53a854fba --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_astro.rs @@ -0,0 +1,142 @@ +//! Language-specific chunk classifiers for Astro. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct AstroClassifier; + +impl LangClassifier for AstroClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_astro_node(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_astro_node(node, source) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_astro_node(node, source) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + kind == "document" + } +} + +fn classify_astro_node<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "frontmatter" => Some(classify_frontmatter(node, source)), + "frontmatter_js_block" => Some(group_candidate(node, "code", source)), + "element" => classify_element(node, source), + "script_element" => Some(classify_script_element(node, source)), + "style_element" => Some(classify_style_element(node, source)), + "html_interpolation" => Some(classify_html_interpolation(node, source)), + "attribute_interpolation" => Some(classify_attribute_interpolation(node, source)), + "attribute_js_expr" => Some(group_candidate(node, "expr", source)), + "text" => Some(group_candidate(node, "text", source)), + _ => None, + } +} + +fn classify_frontmatter<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + make_container_chunk( + node, + "frontmatter".to_string(), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["frontmatter_js_block"]), + ) +} + +fn classify_element<'t>(node: Node<'t>, source: &str) -> Option> { + let tag_name = extract_tag_name(node, source)?; + let prefix = if is_component_name(tag_name.as_str()) { + "component" + } else { + "tag" + }; + Some(make_container_chunk( + node, + format!("{prefix}_{tag_name}"), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + )) +} + +fn classify_script_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "is:inline", source) { + "script_inline" + } else { + "script" + }; + // The Astro grammar exposes script bodies as `raw_text`, not nested JS AST. + make_named_chunk(node, name.to_string(), source, None) +} + +fn classify_style_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "define:vars", source) { + "style_vars" + } else if has_attribute(node, "is:global", source) { + "style_global" + } else { + "style" + }; + // The Astro grammar exposes style bodies as `raw_text`, so the section itself + // is the truthful chunk boundary. + make_named_chunk(node, name.to_string(), source, None) +} + +fn classify_html_interpolation<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = child_by_kind(node, &["permissible_text"]) + .and_then(|expr| sanitize_identifier(node_text(source, expr.start_byte(), expr.end_byte()))) + .map_or_else(|| "expr".to_string(), |expr| format!("expr_{expr}")); + + if let Some(nested_element) = + child_by_kind(node, &["element", "script_element", "style_element"]) + { + make_container_chunk( + node, + name, + source, + Some(recurse_self(nested_element, ChunkContext::ClassBody)), + ) + } else { + make_named_chunk(node, name, source, None) + } +} + +fn classify_attribute_interpolation<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = child_by_kind(node, &["attribute_js_expr"]) + .and_then(|expr| sanitize_identifier(node_text(source, expr.start_byte(), expr.end_byte()))) + .map_or_else(|| "attr_expr".to_string(), |expr| format!("attr_expr_{expr}")); + make_named_chunk(node, name, source, None) +} + +fn extract_tag_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["start_tag", "self_closing_tag"]) + .and_then(|tag| child_by_kind(tag, &["tag_name"])) + .and_then(|tag_name| { + sanitize_identifier(node_text(source, tag_name.start_byte(), tag_name.end_byte())) + }) +} + +fn has_attribute(node: Node<'_>, name: &str, source: &str) -> bool { + child_by_kind(node, &["start_tag", "self_closing_tag"]) + .into_iter() + .flat_map(named_children) + .filter(|child| child.kind() == "attribute") + .filter_map(|attr| extract_attribute_name(attr, source)) + .any(|attr_name| attr_name == name) +} + +fn extract_attribute_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["attribute_name"]).map(|name| { + node_text(source, name.start_byte(), name.end_byte()) + .trim() + .to_string() + }) +} + +fn is_component_name(tag_name: &str) -> bool { + tag_name.chars().next().is_some_and(char::is_uppercase) +} diff --git a/crates/pi-natives/src/chunk/ast_cmake.rs b/crates/pi-natives/src/chunk/ast_cmake.rs new file mode 100644 index 000000000..dd62d9c18 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_cmake.rs @@ -0,0 +1,147 @@ +//! CMake-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct CMakeClassifier; + +fn child_text<'a>(source: &'a str, node: Node<'_>) -> &'a str { + node_text(source, node.start_byte(), node.end_byte()) +} + +fn first_named_child(node: Node<'_>) -> Option> { + named_children(node).into_iter().next() +} + +fn first_named_child_of_kind<'t>(node: Node<'t>, kind: &str) -> Option> { + named_children(node) + .into_iter() + .find(|child| child.kind() == kind) +} + +fn command_name(node: Node<'_>, source: &str) -> Option { + first_named_child(node).and_then(|child| sanitize_identifier(child_text(source, child))) +} + +fn argument_nodes(node: Node<'_>) -> Vec> { + first_named_child_of_kind(node, "argument_list") + .map(named_children) + .unwrap_or_default() + .into_iter() + .filter(|child| child.kind() == "argument") + .collect() +} + +fn nth_argument_name(node: Node<'_>, index: usize, source: &str) -> Option { + argument_nodes(node) + .into_iter() + .nth(index) + .and_then(|arg| sanitize_identifier(child_text(source, arg))) +} + +fn classify_definition<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "function_def" => { + let header = first_named_child_of_kind(node, "function_command")?; + let name = nth_argument_name(header, 0, source).unwrap_or_else(|| "anonymous".to_string()); + Some(make_container_chunk( + node, + format!("fn_{name}"), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["body"]), + )) + }, + "macro_def" => { + let header = first_named_child_of_kind(node, "macro_command")?; + let name = nth_argument_name(header, 0, source).unwrap_or_else(|| "anonymous".to_string()); + Some(make_container_chunk( + node, + format!("macro_{name}"), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["body"]), + )) + }, + "if_condition" => Some(make_container_chunk( + node, + "if".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )), + "foreach_loop" | "while_loop" => Some(make_container_chunk( + node, + "loop".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["body"]), + )), + _ => None, + } +} + +fn classify_command<'t>(node: Node<'t>, source: &str) -> Option> { + if node.kind() != "normal_command" { + return None; + } + + let command = command_name(node, source)?; + Some(match command.as_str() { + "cmake_minimum_required" => make_named_chunk(node, "version_gate".to_string(), source, None), + "project" => { + let name = nth_argument_name(node, 0, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("project_{name}"), source, None) + }, + "include" | "find_package" => group_candidate(node, "imports", source), + "option" => { + let name = nth_argument_name(node, 0, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("option_{name}"), source, None) + }, + "set" => { + let name = nth_argument_name(node, 0, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("var_{name}"), source, None) + }, + "add_library" | "add_executable" | "add_custom_target" => { + let name = nth_argument_name(node, 0, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("target_{name}"), source, None) + }, + "install" | "export" => group_candidate(node, "install", source), + other => group_candidate(node, &format!("cmd_{other}"), source), + }) +} + +fn classify_if_child<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "if_command" => Some(group_candidate(node, "cond", source)), + "elseif_command" => Some(positional_candidate(node, "elif", source)), + "else_command" => Some(positional_candidate(node, "else", source)), + "body" => Some(make_container_chunk( + node, + "block".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )), + _ => None, + } +} + +impl LangClassifier for CMakeClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_definition(node, source).or_else(|| classify_command(node, source)) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_definition(node, source) + .or_else(|| classify_if_child(node, source)) + .or_else(|| classify_command(node, source)) + } + + fn is_trivia(&self, kind: &str) -> bool { + matches!( + kind, + "endif_command" + | "endforeach_command" + | "endwhile_command" + | "endfunction_command" + | "endmacro_command" + ) + } +} diff --git a/crates/pi-natives/src/chunk/ast_dockerfile.rs b/crates/pi-natives/src/chunk/ast_dockerfile.rs new file mode 100644 index 000000000..e3d5f0821 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_dockerfile.rs @@ -0,0 +1,110 @@ +//! Chunk classifier for Dockerfile syntax. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct DockerfileClassifier; + +fn child_text<'a>(source: &'a str, node: Node<'_>) -> &'a str { + node_text(source, node.start_byte(), node.end_byte()) +} + +fn first_named_child(node: Node<'_>) -> Option> { + named_children(node).into_iter().next() +} + +fn first_named_child_of_kind<'t>(node: Node<'t>, kind: &str) -> Option> { + named_children(node) + .into_iter() + .find(|child| child.kind() == kind) +} + +fn extract_stage_name(node: Node<'_>, source: &str) -> Option { + if let Some(alias) = child_by_kind(node, &["image_alias"]) { + return sanitize_identifier(child_text(source, alias)); + } + + child_by_kind(node, &["image_spec"]).and_then(|image| { + let image_name = child_by_kind(image, &["image_name"]).unwrap_or(image); + sanitize_identifier(child_text(source, image_name)) + }) +} + +fn extract_pair_key(node: Node<'_>, pair_kind: &str, source: &str) -> Option { + first_named_child_of_kind(node, pair_kind) + .and_then(first_named_child) + .and_then(|key| sanitize_identifier(unquote_text(child_text(source, key)).as_str())) +} + +fn extract_arg_name(node: Node<'_>, source: &str) -> Option { + first_named_child(node).and_then(|name| sanitize_identifier(child_text(source, name))) +} + +fn recurse_command(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &[], &["shell_command", "json_string_array"]) +} + +fn classify_command_instruction<'t>(node: Node<'t>, source: &str) -> Option> { + let chunk_name = match node.kind() { + "run_instruction" => "run", + "cmd_instruction" => "cmd", + "entrypoint_instruction" => "entrypoint", + _ => return None, + }; + + Some(make_container_chunk(node, chunk_name.to_string(), source, recurse_command(node))) +} + +impl LangClassifier for DockerfileClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "from_instruction" => { + let name = extract_stage_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + Some(make_named_chunk(node, format!("stage_{name}"), source, None)) + }, + "arg_instruction" => { + let name = extract_arg_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + Some(make_named_chunk(node, format!("arg_{name}"), source, None)) + }, + "env_instruction" => { + let name = extract_pair_key(node, "env_pair", source) + .unwrap_or_else(|| "anonymous".to_string()); + Some(make_named_chunk(node, format!("env_{name}"), source, None)) + }, + "label_instruction" => { + let name = extract_pair_key(node, "label_pair", source) + .unwrap_or_else(|| "anonymous".to_string()); + Some(make_named_chunk(node, format!("label_{name}"), source, None)) + }, + "run_instruction" | "cmd_instruction" | "entrypoint_instruction" => { + classify_command_instruction(node, source) + }, + "healthcheck_instruction" => Some(make_container_chunk( + node, + "healthcheck".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["cmd_instruction"]), + )), + "copy_instruction" => Some(group_candidate(node, "copy", source)), + "add_instruction" => Some(group_candidate(node, "add", source)), + "workdir_instruction" => Some(group_candidate(node, "workdir", source)), + "expose_instruction" => Some(group_candidate(node, "expose", source)), + "user_instruction" => Some(group_candidate(node, "user", source)), + _ => None, + } + } + + fn classify_class<'t>(&self, _node: Node<'t>, _source: &str) -> Option> { + None + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "cmd_instruction" => classify_command_instruction(node, source), + "shell_command" => Some(group_candidate(node, "shell", source)), + "json_string_array" => Some(group_candidate(node, "argv", source)), + _ => None, + } + } +} diff --git a/crates/pi-natives/src/chunk/ast_erlang.rs b/crates/pi-natives/src/chunk/ast_erlang.rs new file mode 100644 index 000000000..cc4476fef --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_erlang.rs @@ -0,0 +1,144 @@ +//! Language-specific chunk classifier for Erlang. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct ErlangClassifier; + +fn find_named_descendant_by_kind<'t>(node: Node<'t>, kinds: &[&str]) -> Option> { + if kinds.iter().any(|kind| node.kind() == *kind) { + return Some(node); + } + + for child in named_children(node) { + if let Some(found) = find_named_descendant_by_kind(child, kinds) { + return Some(found); + } + } + + None +} + +fn named_text(node: Node<'_>, source: &str) -> Option { + sanitize_identifier(node_text(source, node.start_byte(), node.end_byte())) +} + +fn erlang_name(node: Node<'_>, source: &str) -> Option { + let name_node = match node.kind() { + "module_attribute" | "record_decl" | "record_field" => { + child_by_field_or_kind(node, &["name"], &["atom"]) + }, + "type_alias" => node + .child_by_field_name("name") + .and_then(|name| find_named_descendant_by_kind(name, &["atom"])), + "spec" => child_by_field_or_kind(node, &["fun"], &["atom"]), + "pp_define" => node + .child_by_field_name("lhs") + .and_then(|lhs| find_named_descendant_by_kind(lhs, &["var"])), + "fun_decl" => node + .child_by_field_name("clause") + .and_then(|clause| child_by_field_or_kind(clause, &["name"], &["atom"])), + "function_clause" => child_by_field_or_kind(node, &["name"], &["atom"]), + _ => child_by_kind(node, &["atom", "var"]), + }?; + + named_text(name_node, source) +} + +fn recurse_clause_body(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &["body"], &["clause_body"]) +} + +impl LangClassifier for ErlangClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "module_attribute" => { + make_named_chunk(node, format!("mod_{}", erlang_name(node, source)?), source, None) + }, + "export_attribute" | "export_type_attribute" => group_candidate(node, "exports", source), + "import_attribute" => group_candidate(node, "imports", source), + "pp_include" | "pp_include_lib" => group_candidate(node, "includes", source), + "pp_define" => { + make_named_chunk(node, format!("macro_{}", erlang_name(node, source)?), source, None) + }, + "record_decl" => make_container_chunk( + node, + format!("record_{}", erlang_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ), + "type_alias" => { + make_named_chunk(node, format!("type_{}", erlang_name(node, source)?), source, None) + }, + // The Erlang grammar exposes each top-level clause as its own `fun_decl`. + // Keep that shape instead of inventing a synthetic merged function node. + "fun_decl" => make_named_chunk( + node, + format!("fn_{}", erlang_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "spec" => return None, + _ => return None, + }) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "record_field" => { + make_named_chunk(node, format!("field_{}", erlang_name(node, source)?), source, None) + }, + _ => return None, + }) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "function_clause" => make_named_chunk( + node, + format!("clause_{}", erlang_name(node, source)?), + source, + recurse_clause_body(node), + ), + "fun_clause" | "cr_clause" => { + make_named_chunk(node, "clause".to_string(), source, recurse_clause_body(node)) + }, + "receive_after" => { + make_named_chunk(node, "after".to_string(), source, recurse_clause_body(node)) + }, + "catch_clause" => { + make_named_chunk(node, "catch".to_string(), source, recurse_clause_body(node)) + }, + "receive_expr" => make_named_chunk( + node, + "receive".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "case_expr" => make_named_chunk( + node, + "case".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "try_expr" => make_named_chunk( + node, + "try".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "anonymous_fun" => make_named_chunk( + node, + "fn_anonymous".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + _ => return None, + }) + } + + fn is_absorbable_attr(&self, kind: &str) -> bool { + kind == "spec" + } +} diff --git a/crates/pi-natives/src/chunk/ast_go.rs b/crates/pi-natives/src/chunk/ast_go.rs index b7c5f8fd9..7e28e290b 100644 --- a/crates/pi-natives/src/chunk/ast_go.rs +++ b/crates/pi-natives/src/chunk/ast_go.rs @@ -1,13 +1,6 @@ -//! Language-specific chunk classifier for Go. - -use std::collections::HashMap; - use tree_sitter::Node; -use super::{ - classify::LangClassifier, common::*, rename_chunk_subtree, sort_chunk_children_by_position, - types::ChunkNode, -}; +use super::{classify::LangClassifier, common::*}; pub struct GoClassifier; @@ -146,16 +139,6 @@ impl LangClassifier for GoClassifier { _ => None, } } - - fn post_process( - &self, - chunks: &mut Vec, - root_children: &mut Vec, - source: &str, - ) { - reparent_receiver_methods(chunks, root_children, source); - reparent_new_type_constructors(chunks, root_children, source); - } } /// Classify Go function-level nodes (reused for top-level control flow @@ -228,142 +211,3 @@ fn recurse_type_spec(node: Node<'_>) -> Option> { .unwrap_or(container); Some(RecurseSpec { node: body, context: ChunkContext::ClassBody }) } - -/// Move receiver methods (`fn_X`) under their corresponding `type_Y` chunk. -/// -/// For `func (s *Server) Start()`, extracts receiver type `Server` and -/// reparents the method chunk under `type_Server`. -fn reparent_receiver_methods( - chunks: &mut [ChunkNode], - root_children: &mut Vec, - source: &str, -) { - // Build map: type name -> chunk path for root-level type chunks. - let type_paths: HashMap = chunks - .iter() - .filter(|c| c.parent_path.as_deref() == Some("") && c.path.starts_with("type_")) - .filter_map(|c| { - c.path - .strip_prefix("type_") - .map(|name| (name.to_string(), c.path.clone())) - }) - .collect(); - - // Collect root-level function paths that may be receiver methods. - let receiver_methods: Vec = root_children - .iter() - .filter(|p| p.starts_with("fn_")) - .cloned() - .collect(); - - for method_path in receiver_methods { - let Some(method_idx) = chunks.iter().position(|c| c.path == method_path) else { - continue; - }; - let Some(receiver_type) = extract_receiver_type_name(&chunks[method_idx], source) else { - continue; - }; - let Some(type_path) = type_paths.get(receiver_type.as_str()) else { - continue; - }; - - let method_name = chunks[method_idx].name.clone(); - let new_path = format!("{type_path}.{method_name}"); - - rename_chunk_subtree(chunks, &method_path, &new_path, type_path); - root_children.retain(|p| p != &method_path); - - if let Some(type_idx) = chunks.iter().position(|c| c.path == *type_path) { - let type_chunk = &mut chunks[type_idx]; - type_chunk.leaf = false; - if !type_chunk.children.iter().any(|child| child == &new_path) { - type_chunk.children.push(new_path); - } - } - } - - sort_chunk_children_by_position(chunks); -} - -/// Move `func NewTypeName(...)` constructors under `type_TypeName` so tree -/// order matches source order (avoids root-level `fn_NewX` appearing after -/// nested methods with a lower file line). -fn reparent_new_type_constructors( - chunks: &mut [ChunkNode], - root_children: &mut Vec, - _source: &str, -) { - let type_paths: HashMap = chunks - .iter() - .filter(|c| c.parent_path.as_deref() == Some("") && c.path.starts_with("type_")) - .filter_map(|c| { - c.path - .strip_prefix("type_") - .map(|name| (name.to_string(), c.path.clone())) - }) - .collect(); - - let constructors: Vec = root_children - .iter() - .filter(|p| { - p.starts_with("fn_") - && constructor_suffix_after_new(p) - .is_some_and(|tail| type_paths.contains_key(tail.as_str())) - }) - .cloned() - .collect(); - - for fn_path in constructors { - let Some(fn_idx) = chunks.iter().position(|c| c.path == fn_path) else { - continue; - }; - let Some(tail) = constructor_suffix_after_new(&fn_path) else { - continue; - }; - let Some(type_path) = type_paths.get(tail.as_str()) else { - continue; - }; - let fn_name = chunks[fn_idx].name.clone(); - let new_path = format!("{type_path}.{fn_name}"); - - rename_chunk_subtree(chunks, &fn_path, &new_path, type_path); - root_children.retain(|p| p != &fn_path); - - if let Some(type_idx) = chunks.iter().position(|c| c.path == *type_path) { - let type_chunk = &mut chunks[type_idx]; - type_chunk.leaf = false; - if !type_chunk.children.iter().any(|child| child == &new_path) { - type_chunk.children.push(new_path); - } - } - } - - sort_chunk_children_by_position(chunks); -} - -/// `fn_NewServer` + type `Server` -> `Some("Server")`; `fn_Start` -> None. -fn constructor_suffix_after_new(fn_path: &str) -> Option { - let name = fn_path.strip_prefix("fn_")?; - let tail = name.strip_prefix("New")?; - if tail.is_empty() { - return None; - } - Some(tail.to_string()) -} - -/// Extract the receiver type name from a Go method's header. -/// -/// `func (s *Server) Start()` -> `Some("Server")` -/// `func (s Server) Stop()` -> `Some("Server")` -fn extract_receiver_type_name(chunk: &ChunkNode, source: &str) -> Option { - let header = normalized_header(source, chunk.start_byte as usize, chunk.end_byte as usize); - let receiver = header - .strip_prefix("func")? - .trim_start() - .strip_prefix('(')? - .split(')') - .next()? - .trim(); - let receiver_type = receiver.split_whitespace().last()?; - sanitize_identifier(receiver_type) -} diff --git a/crates/pi-natives/src/chunk/ast_graphql.rs b/crates/pi-natives/src/chunk/ast_graphql.rs new file mode 100644 index 000000000..febea2755 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_graphql.rs @@ -0,0 +1,254 @@ +//! GraphQL-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct GraphqlClassifier; + +impl LangClassifier for GraphqlClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_graphql_root(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_graphql_class(node, source) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_graphql_function(node, source) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + matches!( + kind, + "document" + | "definition" + | "type_system_definition" + | "type_definition" + | "executable_definition" + ) + } + + fn is_trivia(&self, kind: &str) -> bool { + matches!(kind, "comma") + } +} + +fn classify_graphql_root<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "schema_definition" => Some(make_container_chunk( + node, + "schema".to_string(), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + )), + "directive_definition" => Some(make_container_chunk( + node, + format!( + "directive_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["arguments_definition"]), + )), + "scalar_type_definition" => Some(make_named_chunk( + node, + format!( + "scalar_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + )), + "object_type_definition" => Some(make_container_chunk( + node, + format!( + "type_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["fields_definition"]), + )), + "interface_type_definition" => Some(make_container_chunk( + node, + format!( + "interface_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["fields_definition"]), + )), + "union_type_definition" => Some(make_named_chunk( + node, + format!( + "union_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + )), + "enum_type_definition" => Some(make_container_chunk( + node, + format!( + "enum_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["enum_values_definition"]), + )), + "input_object_type_definition" => Some(make_container_chunk( + node, + format!( + "input_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["input_fields_definition"]), + )), + "operation_definition" => Some(make_container_chunk( + node, + extract_graphql_operation_chunk_name(node, source), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["selection_set"]), + )), + "fragment_definition" => Some(make_container_chunk( + node, + format!( + "fragment_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["selection_set"]), + )), + _ => None, + } +} + +fn classify_graphql_class<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "root_operation_type_definition" => Some(make_named_chunk( + node, + format!( + "root_{}", + extract_graphql_operation_type(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + )), + "field_definition" => { + let name = extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + let recurse = recurse_into(node, ChunkContext::ClassBody, &[], &["arguments_definition"]); + Some(match recurse { + Some(recurse) => { + make_container_chunk(node, format!("field_{name}"), source, Some(recurse)) + }, + None => make_named_chunk(node, format!("field_{name}"), source, None), + }) + }, + "input_value_definition" => Some(classify_graphql_input_value(node, source)), + "enum_value_definition" => Some(make_named_chunk( + node, + format!( + "value_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + )), + _ => None, + } +} + +fn classify_graphql_function<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "selection" => classify_graphql_selection(node, source), + _ => None, + } +} + +fn classify_graphql_selection<'t>(node: Node<'t>, source: &str) -> Option> { + let child = first_named_child(node)?; + match child.kind() { + "field" => { + let name = extract_graphql_name(child, source).unwrap_or_else(|| "anonymous".to_string()); + let recurse = recurse_into(child, ChunkContext::FunctionBody, &[], &["selection_set"]); + Some(match recurse { + Some(recurse) => make_container_chunk_from( + node, + child, + format!("field_{name}"), + source, + Some(recurse), + ), + None => make_named_chunk_from(node, child, format!("field_{name}"), source, None), + }) + }, + "fragment_spread" => Some(make_named_chunk_from( + node, + child, + format!( + "spread_{}", + extract_graphql_name(child, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + )), + "inline_fragment" => Some(make_container_chunk_from( + node, + child, + "inline_fragment".to_string(), + source, + recurse_into(child, ChunkContext::FunctionBody, &[], &["selection_set"]), + )), + _ => None, + } +} + +fn classify_graphql_input_value<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let prefix = match node.parent().map(|parent| parent.kind()) { + Some("input_fields_definition") => "field", + _ => "arg", + }; + make_named_chunk( + node, + format!( + "{prefix}_{}", + extract_graphql_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + ) +} + +fn extract_graphql_name(node: Node<'_>, source: &str) -> Option { + find_graphql_name_node(node) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) +} + +fn find_graphql_name_node(node: Node<'_>) -> Option> { + match node.kind() { + "name" | "fragment_name" => Some(node), + _ => named_children(node) + .into_iter() + .find_map(find_graphql_name_node), + } +} + +fn extract_graphql_operation_type(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["operation_type"]) + .and_then(|kind| sanitize_identifier(node_text(source, kind.start_byte(), kind.end_byte()))) +} + +fn extract_graphql_operation_chunk_name(node: Node<'_>, source: &str) -> String { + let operation = + extract_graphql_operation_type(node, source).unwrap_or_else(|| "operation".to_string()); + match extract_graphql_name(node, source) { + Some(name) => format!("{operation}_{name}"), + None => operation, + } +} + +fn first_named_child(node: Node<'_>) -> Option> { + (0..node.named_child_count()).find_map(|index| node.named_child(index)) +} diff --git a/crates/pi-natives/src/chunk/ast_ini.rs b/crates/pi-natives/src/chunk/ast_ini.rs new file mode 100644 index 000000000..dad78cdbf --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_ini.rs @@ -0,0 +1,64 @@ +//! Chunk classifier for INI. +//! +//! The tree-sitter INI grammar is intentionally flat: a document contains +//! root-level `setting` nodes and `section` containers, and a section contains +//! only its own `setting` children. Mirror that structure directly instead of +//! inventing deeper hierarchy. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct IniClassifier; + +impl LangClassifier for IniClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_ini_root(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_ini_class(node, source) + } +} + +fn classify_ini_root<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "section" => make_container_chunk( + node, + format!("section_{}", ini_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ), + // INI permits settings before any section header; keep them as first-class + // chunks instead of forcing them under a synthetic container. + "setting" => make_named_chunk(node, format!("key_{}", ini_name(node, source)?), source, None), + _ => return None, + }) +} + +fn classify_ini_class<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "setting" => make_named_chunk(node, format!("key_{}", ini_name(node, source)?), source, None), + _ => return None, + }) +} + +fn ini_name(node: Node<'_>, source: &str) -> Option { + find_named_text(node, source, &["section_name", "setting_name", "text"]).and_then(|text| { + sanitize_identifier(text.trim().trim_start_matches('[').trim_end_matches(']')) + }) +} + +fn find_named_text<'a>(node: Node<'_>, source: &'a str, kinds: &[&str]) -> Option<&'a str> { + if kinds.iter().any(|kind| node.kind() == *kind) { + return Some(node_text(source, node.start_byte(), node.end_byte())); + } + + for child in named_children(node) { + if let Some(text) = find_named_text(child, source, kinds) { + return Some(text); + } + } + + None +} diff --git a/crates/pi-natives/src/chunk/ast_ipynb.rs b/crates/pi-natives/src/chunk/ast_ipynb.rs new file mode 100644 index 000000000..fccafd29c --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_ipynb.rs @@ -0,0 +1,912 @@ +//! Jupyter notebook (`.ipynb`) chunker. +//! +//! Notebooks are JSON documents whose `cells` array carries the code/markdown +//! that users actually edit. This module parses the JSON, extracts each cell +//! as its own source fragment, and assembles a *virtual source* — the +//! concatenation of all cell bodies with one-line marker headers — that the +//! rest of the chunk pipeline can treat like any other text file. +//! +//! The per-cell sub-chunks come from recursively running [`build_chunk_tree`] +//! on the individual cell sources with their appropriate language (code cells +//! use the notebook's kernel language, defaulting to Python; markdown cells +//! use `markdown`; raw cells fall through to the blank-line fallback). Their +//! byte and line offsets are shifted into the virtual source and then +//! rewrapped under `cell_` parent chunks so the resulting tree looks to +//! edit.rs exactly like a normal multi-symbol file. +//! +//! On write-back, [`notebook_to_json`] walks the (possibly edited) virtual +//! source, splits it at the cell markers, updates each cell's `source` field +//! in a preserved [`NotebookContext`], and serializes the whole notebook back +//! to JSON. Cell metadata (`metadata`, `outputs`, `execution_count`, `id`, +//! attachments, etc.) is preserved verbatim. + +use std::sync::Arc; + +use serde::Serialize; +use serde_json::{Map, Value}; + +use crate::chunk::{ + build_chunk_tree, chunk_checksum, line_start_offsets, + types::{ChunkNode, ChunkTree}, +}; + +/// Marker line prefix placed before every cell body in the virtual source. +/// +/// Format: `# %%% oh-my-pi cell_ []` +/// +/// The leading `#` makes the marker a valid comment in Python and most other +/// code languages, and the `oh-my-pi` tag makes accidental collision with +/// user content vanishingly unlikely. Markdown cells get the same marker — +/// `#` in markdown is a heading, but the marker line itself is stripped +/// before the markdown chunker parses the cell body (see +/// [`build_cell_sub_tree`]). +const MARKER_PREFIX: &str = "# %%% oh-my-pi cell_"; + +/// Returns true if `line` is a cell marker; parses the cell index and type. +fn parse_marker_line(line: &str) -> Option<(usize, &str)> { + let rest = line.strip_prefix(MARKER_PREFIX)?; + let (num_str, after_num) = rest.split_once(' ')?; + let index: usize = num_str.parse().ok()?; + let cell_type = after_num.strip_prefix('[')?.strip_suffix(']')?; + Some((index, cell_type)) +} + +fn format_marker(index: usize, cell_type: &str) -> String { + format!("{MARKER_PREFIX}{index} [{cell_type}]") +} + +/// Metadata for a single notebook cell. Everything except the joined `source` +/// is preserved verbatim for JSON round-tripping. +#[derive(Clone)] +pub struct NotebookCell { + pub cell_type: String, + pub source: String, + pub metadata: Value, + pub outputs: Option, + pub execution_count: Option, + /// Additional fields on the cell object (`id`, `attachments`, …) so we + /// emit the same keys we consumed. + pub other: Map, + /// Whether the original cell used `"source"` as a string (`true`) or an + /// array of lines (`false`). Preserved so round-trips stay byte-identical + /// when only metadata or unrelated cells change. + pub source_was_string: bool, +} + +/// Preserved notebook-level state used for rebuilding the JSON after edits. +#[derive(Clone)] +pub struct NotebookContext { + /// Cells in document order. + pub cells: Vec, + /// Top-level notebook fields other than `cells`, in original key order. + pub top_fields: Map, + /// Indent string detected from the original JSON (typically `" "`). + pub indent: String, + /// Whether the original file ended with a trailing `\n`. + pub trailing_newline: bool, + /// Normalized kernel language used for code cells (e.g. `python`). + pub kernel_language: String, +} + +/// Result of a notebook parse: the virtual source ready for chunking plus +/// the context needed to rebuild the JSON later. +pub struct NotebookParse { + pub virtual_source: String, + pub context: NotebookContext, +} + +// ──────────────────────────────────────────────────────────────────── +// JSON parsing +// ──────────────────────────────────────────────────────────────────── + +/// Parse the raw ipynb JSON into a [`NotebookContext`] and the derived +/// virtual source text. Returns a descriptive error if the JSON is invalid +/// or not a notebook document. +pub fn parse_notebook(source: &str) -> Result { + let normalized = strip_bom(source); + let value: Value = serde_json::from_str(normalized) + .map_err(|err| format!("Invalid Jupyter notebook JSON: {err}"))?; + let Value::Object(obj) = value else { + return Err("Invalid Jupyter notebook: top-level value must be an object.".to_string()); + }; + + let cells_val = obj + .get("cells") + .ok_or_else(|| "Invalid Jupyter notebook: missing `cells` array.".to_string())?; + let cells_arr = cells_val + .as_array() + .ok_or_else(|| "Invalid Jupyter notebook: `cells` is not an array.".to_string())?; + + let mut cells = Vec::with_capacity(cells_arr.len()); + for (i, raw) in cells_arr.iter().enumerate() { + cells.push(parse_cell(raw, i)?); + } + + // Top-level fields minus `cells`, preserving insertion order. + let mut top_fields = Map::new(); + for (k, v) in &obj { + if k != "cells" { + top_fields.insert(k.clone(), v.clone()); + } + } + + let kernel_language = + extract_kernel_language(&top_fields).unwrap_or_else(|| "python".to_string()); + let indent = detect_json_indent(normalized); + let trailing_newline = normalized.ends_with('\n'); + + let ctx = NotebookContext { cells, top_fields, indent, trailing_newline, kernel_language }; + let virtual_source = build_virtual_source(&ctx); + Ok(NotebookParse { virtual_source, context: ctx }) +} + +fn parse_cell(raw: &Value, index: usize) -> Result { + let obj = raw + .as_object() + .ok_or_else(|| format!("Invalid Jupyter notebook: cell {} is not an object.", index + 1))?; + + let cell_type = obj + .get("cell_type") + .and_then(Value::as_str) + .unwrap_or("code") + .to_string(); + + let (source, source_was_string) = match obj.get("source") { + Some(Value::String(s)) => (s.clone(), true), + Some(Value::Array(arr)) => { + let mut joined = String::new(); + for item in arr { + match item { + Value::String(s) => joined.push_str(s), + _ => { + return Err(format!( + "Invalid Jupyter notebook: cell {} source array contains non-string element.", + index + 1 + )); + }, + } + } + (joined, false) + }, + Some(Value::Null) | None => (String::new(), false), + Some(_) => { + return Err(format!( + "Invalid Jupyter notebook: cell {} has a non-string `source` field.", + index + 1 + )); + }, + }; + + let metadata = obj + .get("metadata") + .cloned() + .unwrap_or_else(|| Value::Object(Map::new())); + let outputs = obj.get("outputs").cloned(); + let execution_count = obj.get("execution_count").cloned(); + + // Additional fields (id, attachments, …) that we pass through untouched. + let mut other = Map::new(); + for (k, v) in obj { + if !matches!(k.as_str(), "cell_type" | "source" | "metadata" | "outputs" | "execution_count") + { + other.insert(k.clone(), v.clone()); + } + } + + Ok(NotebookCell { + cell_type, + source, + metadata, + outputs, + execution_count, + other, + source_was_string, + }) +} + +fn strip_bom(source: &str) -> &str { + source.strip_prefix('\u{feff}').unwrap_or(source) +} + +fn extract_kernel_language(top_fields: &Map) -> Option { + if let Some(meta) = top_fields.get("metadata").and_then(Value::as_object) { + if let Some(lang) = meta + .get("kernelspec") + .and_then(Value::as_object) + .and_then(|k| k.get("language")) + .and_then(Value::as_str) + && !lang.is_empty() + { + return Some(lang.to_ascii_lowercase()); + } + if let Some(lang) = meta + .get("language_info") + .and_then(Value::as_object) + .and_then(|k| k.get("name")) + .and_then(Value::as_str) + && !lang.is_empty() + { + return Some(lang.to_ascii_lowercase()); + } + } + None +} + +fn detect_json_indent(source: &str) -> String { + // Look for the first `\n` followed by whitespace inside the top-level object + // (i.e. after the opening `{`). This is a heuristic; Jupyter canonically + // uses a single space per level. + let Some(brace) = source.find('{') else { + return " ".to_string(); + }; + let rest = &source[brace + 1..]; + let Some(nl) = rest.find('\n') else { + return " ".to_string(); + }; + let after_nl = &rest[nl + 1..]; + let mut end = 0usize; + for ch in after_nl.chars() { + if ch == ' ' || ch == '\t' { + end += ch.len_utf8(); + } else { + break; + } + } + if end == 0 { + " ".to_string() + } else { + after_nl[..end].to_string() + } +} + +// ──────────────────────────────────────────────────────────────────── +// Virtual source assembly +// ──────────────────────────────────────────────────────────────────── + +/// Build the virtual source text from the current cell list. +/// +/// Each cell is preceded by a marker line and its body. Cells do not include +/// trailing blank separators — we rely purely on the marker line to delimit +/// adjacent cells so the reconstructed sources stay byte-identical to the +/// originals after whole-cell edits. +pub fn build_virtual_source(ctx: &NotebookContext) -> String { + let mut out = String::new(); + for (i, cell) in ctx.cells.iter().enumerate() { + out.push_str(&format_marker(i + 1, &cell.cell_type)); + out.push('\n'); + out.push_str(&cell.source); + if !cell.source.is_empty() && !cell.source.ends_with('\n') { + out.push('\n'); + } + } + out +} + +// ──────────────────────────────────────────────────────────────────── +// Chunk tree construction +// ──────────────────────────────────────────────────────────────────── + +/// Locate every cell marker in `source`, returning tuples of: +/// (`cell_number`, `marker_line_byte_start`, `content_byte_start`, +/// `content_byte_end`, `marker_line_number_1based`, +/// `content_start_line_1based`, +/// `content_end_line_1based_inclusive_or_zero_if_empty`) +struct CellRegion { + cell_num: usize, + cell_type: String, + marker_start: usize, // byte offset of the `#` starting the marker line + content_start: usize, // byte offset of the first byte of the cell body + content_end: usize, // byte offset one past the last byte of the cell body + marker_line: u32, // 1-based line number of the marker + content_line: u32, // 1-based line number of the first body line (or marker_line + 1) + content_end_line: u32, /* 1-based line number of the last body line (== content_line - 1 + * for empty bodies) */ +} + +/// Scan a virtual source text for cell markers and return the list of +/// regions. Assumes markers occur at the very start of their line. +fn scan_cells(virtual_source: &str) -> Vec { + let line_starts = line_start_offsets(virtual_source); + let mut regions: Vec = Vec::new(); + + for (line_idx, &line_start) in line_starts.iter().enumerate() { + let line_end = if line_idx + 1 < line_starts.len() { + // Exclude the trailing newline + line_starts[line_idx + 1] - 1 + } else { + virtual_source.len() + }; + let line = &virtual_source[line_start..line_end]; + if let Some((cell_num, cell_type)) = parse_marker_line(line) { + // Close the previous region if any. + if let Some(prev) = regions.last_mut() { + prev.content_end = line_start; + // Trim trailing newline from content_end if present (i.e. the body ended with + // \n). Actually we keep the newline: cell bodies end with \n except + // possibly the last. content_end_line = line of the last body byte. + if prev.content_end > prev.content_start { + let body_last_char_line = line_idx; // line_idx is 0-based, so this is the previous line + prev.content_end_line = body_last_char_line as u32; + } else { + prev.content_end_line = prev.content_line.saturating_sub(1); + } + } + let content_start = if line_idx + 1 < line_starts.len() { + line_starts[line_idx + 1] + } else { + virtual_source.len() + }; + regions.push(CellRegion { + cell_num, + cell_type: cell_type.to_string(), + marker_start: line_start, + content_start, + content_end: virtual_source.len(), // provisional, closed by the next marker + marker_line: (line_idx as u32) + 1, + content_line: (line_idx as u32) + 2, + content_end_line: 0, + }); + } + } + + // Close the last region. + if let Some(last) = regions.last_mut() { + last.content_end = virtual_source.len(); + if last.content_end > last.content_start { + // Count lines inside the body. + let body = &virtual_source[last.content_start..last.content_end]; + let body_lines = body.matches('\n').count(); + // If body doesn't end with '\n', the final partial line still counts. + let has_trailing_nl = body.ends_with('\n'); + let content_lines = if has_trailing_nl { + body_lines + } else { + body_lines + 1 + }; + if content_lines > 0 { + last.content_end_line = last.content_line + content_lines as u32 - 1; + } else { + last.content_end_line = last.content_line.saturating_sub(1); + } + } else { + last.content_end_line = last.content_line.saturating_sub(1); + } + } + + regions +} + +/// Build a chunk tree from a virtual source text. +/// +/// Re-scans the virtual source for cell markers, parses each cell body with +/// its language, and wraps the results in `cell_` parent chunks. This is +/// the entry point used by both the initial JSON-based parse (via +/// [`parse_notebook`] → `build_virtual_source` → this function) and the +/// post-edit rebuilds that operate directly on the mutated virtual source. +pub fn build_notebook_tree_from_virtual( + virtual_source: &str, + kernel_language: &str, +) -> Result { + let total_lines = total_line_count(virtual_source); + let root_checksum = chunk_checksum(virtual_source.as_bytes()); + let regions = scan_cells(virtual_source); + + // Accumulated chunk nodes. Index 0 is reserved for the synthetic root. + let mut chunks: Vec = Vec::with_capacity(1 + regions.len() * 2); + chunks.push(ChunkNode { + path: String::new(), + name: "root".to_string(), + leaf: false, + parent_path: None, + children: Vec::new(), + signature: None, + start_line: u32::from(total_lines != 0), + end_line: total_lines as u32, + line_count: total_lines as u32, + start_byte: 0, + end_byte: virtual_source.len() as u32, + checksum_start_byte: 0, + body_start_byte: None, + body_end_byte: None, + checksum: root_checksum.clone(), + error: false, + indent: 0, + indent_char: String::new(), + group: false, + }); + + let mut root_children: Vec = Vec::with_capacity(regions.len()); + + for region in ®ions { + let cell_path = format!("cell_{}", region.cell_num); + let cell_language_str = match region.cell_type.as_str() { + "code" => kernel_language.to_string(), + "markdown" => "markdown".to_string(), + _ => String::new(), + }; + + let body = &virtual_source[region.content_start..region.content_end]; + let body_has_content = !body.is_empty(); + let cell_checksum = chunk_checksum(body.as_bytes()); + + // Build sub-chunks by parsing the cell body in isolation. Offsets in + // the returned tree are relative to `body`; we translate them into + // virtual-source coordinates by adding `region.content_start` bytes + // and `region.content_line - 1` lines. + let mut cell_children_paths: Vec = Vec::new(); + if body_has_content { + let sub_tree = build_chunk_tree(body, cell_language_str.as_str()) + .map_err(|err| format!("Failed to parse cell_{} body: {err}", region.cell_num))?; + // sub_tree.chunks[0] is the sub-root; skip it. + for sub_chunk in sub_tree.chunks.into_iter().skip(1) { + let translated_path = format!("{}.{}", cell_path, sub_chunk.path); + let translated_parent = match sub_chunk.parent_path.as_deref() { + Some("") | None => Some(cell_path.clone()), + Some(other) => Some(format!("{cell_path}.{other}")), + }; + let translated_children: Vec = sub_chunk + .children + .iter() + .map(|c| format!("{cell_path}.{c}")) + .collect(); + let shifted_start_byte = sub_chunk + .start_byte + .saturating_add(region.content_start as u32); + let shifted_end_byte = sub_chunk + .end_byte + .saturating_add(region.content_start as u32); + let line_shift = region.content_line.saturating_sub(1); + chunks.push(ChunkNode { + path: translated_path.clone(), + name: sub_chunk.name, + leaf: sub_chunk.leaf, + parent_path: translated_parent, + children: translated_children, + signature: sub_chunk.signature, + start_line: sub_chunk.start_line.saturating_add(line_shift), + end_line: sub_chunk.end_line.saturating_add(line_shift), + line_count: sub_chunk.line_count, + start_byte: shifted_start_byte, + end_byte: shifted_end_byte, + checksum_start_byte: sub_chunk + .checksum_start_byte + .saturating_add(region.content_start as u32), + body_start_byte: sub_chunk + .body_start_byte + .map(|b| b.saturating_add(region.content_start as u32)), + body_end_byte: sub_chunk + .body_end_byte + .map(|b| b.saturating_add(region.content_start as u32)), + checksum: sub_chunk.checksum, + error: sub_chunk.error, + indent: sub_chunk.indent, + indent_char: sub_chunk.indent_char, + group: false, + }); + } + // Collect direct children (those whose parent_path equals cell_path). + for sub_path in sub_tree.root_children { + cell_children_paths.push(format!("{cell_path}.{sub_path}")); + } + } + + // Cell parent chunk: spans from the marker line through the body. + let cell_line_count = { + let body_lines = if body_has_content { + if body.ends_with('\n') { + body.matches('\n').count() + } else { + body.matches('\n').count() + 1 + } + } else { + 0 + }; + 1 + body_lines as u32 // +1 for the marker line + }; + let cell_end_line = region.marker_line + cell_line_count.saturating_sub(1); + let cell_leaf = cell_children_paths.is_empty(); + chunks.push(ChunkNode { + path: cell_path.clone(), + name: cell_path.clone(), + leaf: cell_leaf, + parent_path: Some(String::new()), + children: cell_children_paths, + signature: Some(format!("cell_{} ({})", region.cell_num, region.cell_type)), + start_line: region.marker_line, + end_line: cell_end_line, + line_count: cell_line_count, + start_byte: region.marker_start as u32, + end_byte: region.content_end as u32, + checksum_start_byte: region.content_start as u32, + body_start_byte: Some(region.content_start as u32), + body_end_byte: Some(region.content_end as u32), + checksum: cell_checksum, + error: false, + indent: 0, + indent_char: String::new(), + group: false, + }); + root_children.push(cell_path); + } + + // Populate root children now that every cell is known. + if let Some(root) = chunks.get_mut(0) { + root.children.clone_from(&root_children); + } + + // Sort chunks so the cell parent always comes before its sub-chunks, + // matching the invariant that other paths rely on (render, edit + // scheduling, line-to-chunk lookup). Keep the root at index 0. + // The insertion order above places sub-chunks before the cell parent, so + // we need to reorder: for each cell region, move the cell parent ahead of + // its sub-chunks. + // + // Simpler: rebuild the chunks list by iterating cells, emitting the cell + // parent followed by its sub-chunks in path order. + let mut reordered: Vec = Vec::with_capacity(chunks.len()); + reordered.push(chunks.remove(0)); // root + + let mut remaining: Vec = chunks; + for cell_path in &root_children { + // Extract the cell parent first. + if let Some(pos) = remaining.iter().position(|c| &c.path == cell_path) { + reordered.push(remaining.remove(pos)); + } + // Then any descendants of this cell. + let prefix = format!("{cell_path}."); + let mut i = 0; + while i < remaining.len() { + if remaining[i].path.starts_with(&prefix) { + reordered.push(remaining.remove(i)); + } else { + i += 1; + } + } + } + // Anything left over (shouldn't happen, but be defensive). + reordered.extend(remaining); + + Ok(ChunkTree { + language: "ipynb".to_string(), + checksum: root_checksum, + line_count: total_lines as u32, + parse_errors: 0, + fallback: false, + root_path: String::new(), + root_children, + chunks: reordered, + }) +} + +fn total_line_count(source: &str) -> usize { + if source.is_empty() { + 0 + } else { + source.bytes().filter(|b| *b == b'\n').count() + 1 + } +} + +// ──────────────────────────────────────────────────────────────────── +// Virtual → JSON round-trip +// ──────────────────────────────────────────────────────────────────── + +/// Update a [`NotebookContext`] from a (possibly edited) virtual source, +/// then serialize it back to JSON. Cells that no longer appear in the +/// virtual source are dropped; cells whose markers survive get their +/// `source` field replaced with the current body. +/// +/// Returns the serialized JSON text ready to be written to disk. +pub fn notebook_to_json( + virtual_source: &str, + base_ctx: &NotebookContext, +) -> Result { + let mut ctx = base_ctx.clone(); + let regions = scan_cells(virtual_source); + + // Rebuild the cells array in the order markers appear in the virtual + // source. Look up each marker's original cell by 1-based cell_num so + // edits that reorder cells via sibling insertion continue to track the + // right metadata. + let mut new_cells: Vec = Vec::with_capacity(regions.len()); + for region in ®ions { + let body_slice = &virtual_source[region.content_start..region.content_end]; + // Trim the single trailing newline that the virtual source format + // adds so edits that replace an entire cell body don't grow by one + // line every round-trip. + let body = trim_virtual_body(body_slice); + + let original = ctx.cells.get(region.cell_num.saturating_sub(1)).cloned(); + let cell = match original { + Some(mut cell) => { + cell.source = body.to_string(); + cell.cell_type.clone_from(®ion.cell_type); + cell + }, + None => NotebookCell { + cell_type: region.cell_type.clone(), + source: body.to_string(), + metadata: Value::Object(Map::new()), + outputs: match region.cell_type.as_str() { + "code" => Some(Value::Array(Vec::new())), + _ => None, + }, + execution_count: match region.cell_type.as_str() { + "code" => Some(Value::Null), + _ => None, + }, + other: Map::new(), + source_was_string: false, + }, + }; + new_cells.push(cell); + } + ctx.cells = new_cells; + + let json = serialize_notebook(&ctx)?; + Ok(json) +} + +/// Strip a single trailing newline from `body`, if present. The virtual +/// source always terminates each cell body with `\n` to make the markers +/// start on a fresh line; we remove that byte so the cell's stored source +/// matches the semantic content. +fn trim_virtual_body(body: &str) -> &str { + body.strip_suffix('\n').unwrap_or(body) +} + +fn serialize_notebook(ctx: &NotebookContext) -> Result { + // Build the cells array first. + let mut cells_arr: Vec = Vec::with_capacity(ctx.cells.len()); + for cell in &ctx.cells { + cells_arr.push(cell_to_value(cell)); + } + + // Rebuild the top-level object preserving the original key order with + // `cells` injected at the position it originally occupied. If the input + // had no `cells` key (we wouldn't be here), we append. + let mut top = Map::new(); + let mut cells_inserted = false; + for (k, v) in &ctx.top_fields { + top.insert(k.clone(), v.clone()); + if k == "metadata" && !cells_inserted { + // Jupyter's canonical order is cells, metadata, nbformat, + // nbformat_minor. We preserve whatever we found. + } + } + // If the original document had `cells` somewhere, we want to re-insert + // it at roughly the same slot. Jupyter always writes `cells` first, so + // build a fresh Map in canonical order: cells then the preserved + // top_fields. + let mut final_top = Map::new(); + final_top.insert("cells".to_string(), Value::Array(cells_arr)); + cells_inserted = true; + for (k, v) in top { + if k != "cells" { + final_top.insert(k, v); + } + } + let _ = cells_inserted; + + let indent_bytes = ctx.indent.as_bytes().to_vec(); + let formatter = serde_json::ser::PrettyFormatter::with_indent(&indent_bytes); + let mut buf: Vec = Vec::with_capacity(1024); + { + let mut ser = serde_json::Serializer::with_formatter(&mut buf, formatter); + Value::Object(final_top) + .serialize(&mut ser) + .map_err(|err| format!("Failed to serialize notebook JSON: {err}"))?; + } + let mut text = String::from_utf8(buf) + .map_err(|err| format!("Serialized notebook is not valid UTF-8: {err}"))?; + if ctx.trailing_newline && !text.ends_with('\n') { + text.push('\n'); + } + Ok(text) +} + +fn cell_to_value(cell: &NotebookCell) -> Value { + let mut obj = Map::new(); + obj.insert("cell_type".to_string(), Value::String(cell.cell_type.clone())); + // Preserve `id` and similar fields that idiomatically appear before + // `metadata` in nbformat 4+. + for (k, v) in &cell.other { + if !matches!(k.as_str(), "metadata" | "outputs" | "execution_count" | "source") { + obj.insert(k.clone(), v.clone()); + } + } + obj.insert("metadata".to_string(), cell.metadata.clone()); + if cell.cell_type == "code" { + obj.insert( + "execution_count".to_string(), + cell.execution_count.clone().unwrap_or(Value::Null), + ); + obj.insert( + "outputs".to_string(), + cell + .outputs + .clone() + .unwrap_or_else(|| Value::Array(Vec::new())), + ); + } else { + if let Some(outputs) = &cell.outputs { + obj.insert("outputs".to_string(), outputs.clone()); + } + if let Some(ec) = &cell.execution_count { + obj.insert("execution_count".to_string(), ec.clone()); + } + } + obj.insert("source".to_string(), source_to_value(&cell.source, cell.source_was_string)); + Value::Object(obj) +} + +/// Convert a flat source string to the Jupyter `source` field representation. +/// +/// If the cell originally used a string (or a new cell was inserted), we keep +/// it as a string. Otherwise we split into the canonical `Vec` with +/// each element preserving its trailing `\n`. +fn source_to_value(source: &str, was_string: bool) -> Value { + if was_string { + return Value::String(source.to_string()); + } + if source.is_empty() { + return Value::Array(Vec::new()); + } + let mut parts: Vec = Vec::new(); + for line in source.split_inclusive('\n') { + parts.push(Value::String(line.to_string())); + } + Value::Array(parts) +} + +// ──────────────────────────────────────────────────────────────────── +// Shared helpers +// ──────────────────────────────────────────────────────────────────── + +/// Atomically-shareable notebook context; used by [`ChunkStateInner`] to +/// carry the notebook metadata through edit cycles. +pub type SharedNotebookContext = Arc; + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::chunk::state::ChunkStateInner; + + fn sample_notebook() -> String { + let value = json!({ + "cells": [ + { + "cell_type": "code", + "source": ["def foo():\n", " return 1\n"], + "metadata": {}, + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "source": ["# Hello\n", "World\n"], + "metadata": {} + }, + { + "cell_type": "code", + "source": ["class Bar:\n", " def baz(self):\n", " pass\n"], + "metadata": {}, + "outputs": [], + "execution_count": 3 + } + ], + "metadata": { + "kernelspec": { + "language": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 + }); + serde_json::to_string_pretty(&value).expect("static json") + } + + #[test] + fn parses_notebook_into_cells() { + let nb = parse_notebook(&sample_notebook()).expect("valid notebook"); + assert_eq!(nb.context.cells.len(), 3); + assert_eq!(nb.context.cells[0].cell_type, "code"); + assert_eq!(nb.context.cells[1].cell_type, "markdown"); + assert_eq!(nb.context.cells[2].cell_type, "code"); + assert_eq!(nb.context.kernel_language, "python"); + } + + #[test] + fn virtual_source_contains_all_cells() { + let nb = parse_notebook(&sample_notebook()).expect("valid notebook"); + let vs = &nb.virtual_source; + assert!(vs.contains("def foo():"), "cell 1 body missing"); + assert!(vs.contains("# Hello"), "cell 2 body missing"); + assert!(vs.contains("class Bar:"), "cell 3 body missing"); + assert!(vs.contains("# %%% oh-my-pi cell_1 [code]"), "cell_1 marker missing"); + assert!(vs.contains("# %%% oh-my-pi cell_2 [markdown]"), "cell_2 marker missing"); + assert!(vs.contains("# %%% oh-my-pi cell_3 [code]"), "cell_3 marker missing"); + } + + #[test] + fn builds_cell_level_chunks() { + let nb = parse_notebook(&sample_notebook()).expect("valid notebook"); + let tree = + build_notebook_tree_from_virtual(&nb.virtual_source, "python").expect("tree should build"); + assert_eq!( + tree.root_children, + vec!["cell_1", "cell_2", "cell_3"], + "root children should be the three cells" + ); + let cell1 = tree + .chunks + .iter() + .find(|c| c.path == "cell_1") + .expect("cell_1 chunk"); + assert!(!cell1.leaf, "code cell with a function should not be a leaf"); + assert!( + cell1 + .children + .iter() + .any(|p| p.starts_with("cell_1.fn_foo")), + "cell_1 should contain fn_foo, got {:?}", + cell1.children + ); + } + + #[test] + fn sub_chunk_paths_are_prefixed_with_cell() { + let nb = parse_notebook(&sample_notebook()).expect("valid notebook"); + let tree = + build_notebook_tree_from_virtual(&nb.virtual_source, "python").expect("tree should build"); + let cell3 = tree + .chunks + .iter() + .find(|c| c.path == "cell_3") + .expect("cell_3 chunk"); + assert!( + cell3 + .children + .iter() + .any(|p| p.starts_with("cell_3.class_Bar")), + "cell_3 should contain class_Bar, got {:?}", + cell3.children + ); + let bar_method = tree + .chunks + .iter() + .find(|c| c.path == "cell_3.class_Bar.fn_baz"); + assert!(bar_method.is_some(), "cell_3.class_Bar.fn_baz should exist"); + } + + #[test] + fn chunk_state_parse_ipynb_carries_notebook_context() { + let json = sample_notebook(); + let state = + ChunkStateInner::parse(json, "ipynb".to_string()).expect("ChunkState should parse ipynb"); + assert_eq!(state.language(), "ipynb"); + // Source is the virtual source, not the JSON + assert!(state.source().contains("# %%% oh-my-pi cell_1")); + // The notebook context is preserved for JSON round-trip + let ctx = state + .notebook + .as_ref() + .expect("notebook context should be set"); + let json_out = notebook_to_json(state.source(), ctx).expect("should serialize back to JSON"); + let reparsed: serde_json::Value = + serde_json::from_str(&json_out).expect("output should be valid JSON"); + let cells = reparsed["cells"].as_array().expect("cells array"); + assert_eq!(cells.len(), 3, "should still have 3 cells"); + assert_eq!(cells[0]["cell_type"], "code"); + assert_eq!(cells[2]["execution_count"], 3); + } + + #[test] + fn empty_notebook_produces_empty_tree() { + let json = r#"{"cells": [], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}"#; + let state = ChunkStateInner::parse(json.to_string(), "ipynb".to_string()) + .expect("empty notebook should parse"); + assert!(state.tree().root_children.is_empty()); + } +} diff --git a/crates/pi-natives/src/chunk/ast_just.rs b/crates/pi-natives/src/chunk/ast_just.rs new file mode 100644 index 000000000..6e9c9b8c9 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_just.rs @@ -0,0 +1,90 @@ +//! Language-specific chunk classifier for Just. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct JustClassifier; + +fn first_named_child(node: Node<'_>) -> Option> { + named_children(node).into_iter().next() +} + +fn first_named_child_of_kind<'t>(node: Node<'t>, kind: &str) -> Option> { + named_children(node) + .into_iter() + .find(|child| child.kind() == kind) +} + +fn child_text<'a>(source: &'a str, node: Node<'_>) -> &'a str { + node_text(source, node.start_byte(), node.end_byte()) +} + +/// `set shell := ...` uses a dedicated `shell` token instead of a named +/// identifier, so parse the assignment head text instead of relying on fields. +fn extract_setting_name(node: Node<'_>, source: &str) -> Option { + let header = child_text(source, node).lines().next()?.trim(); + let rest = header.strip_prefix("set ")?; + let name = rest.split_once(":=")?.0.trim(); + sanitize_identifier(name) +} + +fn extract_alias_name(node: Node<'_>, source: &str) -> Option { + first_named_child(node).and_then(|child| sanitize_identifier(child_text(source, child))) +} + +fn extract_recipe_name(node: Node<'_>, source: &str) -> Option { + let header = first_named_child_of_kind(node, "recipe_header")?; + first_named_child(header).and_then(|child| sanitize_identifier(child_text(source, child))) +} + +fn classify_just_root_node<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "setting" => { + let name = extract_setting_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("setting_{name}"), source, None) + }, + "alias" => { + let name = extract_alias_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + make_named_chunk(node, format!("alias_{name}"), source, None) + }, + "recipe" => { + let name = extract_recipe_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + make_container_chunk( + node, + format!("recipe_{name}"), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["recipe_body"]), + ) + }, + _ => return None, + }) +} + +fn classify_just_body_node<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + // Just recipe bodies are line-oriented; tree-sitter exposes shell lines as + // `recipe_line` leaves rather than a nested shell AST. + "recipe_line" => group_candidate(node, "cmd", source), + "shebang" => make_named_chunk(node, "shebang".to_string(), source, None), + _ => return None, + }) +} + +impl LangClassifier for JustClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_just_root_node(node, source) + } + + fn classify_class<'t>(&self, _node: Node<'t>, _source: &str) -> Option> { + None + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_just_body_node(node, source) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + kind == "source_file" + } +} diff --git a/crates/pi-natives/src/chunk/ast_ocaml.rs b/crates/pi-natives/src/chunk/ast_ocaml.rs new file mode 100644 index 000000000..7b03d2d78 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_ocaml.rs @@ -0,0 +1,189 @@ +//! OCaml-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct OcamlClassifier; + +impl LangClassifier for OcamlClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_ocaml_item(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "method_definition" => Some(make_named_chunk( + node, + format!("fn_{}", ocaml_named_text(node, source, &["method_name"])?), + source, + ocaml_method_recurse(node), + )), + "method_specification" => Some(make_named_chunk( + node, + format!("fn_{}", ocaml_named_text(node, source, &["method_name"])?), + source, + None, + )), + "instance_variable_definition" => { + Some(match ocaml_named_text(node, source, &["instance_variable_name"]) { + Some(name) => make_named_chunk(node, format!("field_{name}"), source, None), + None => group_candidate(node, "fields", source), + }) + }, + _ => classify_ocaml_item(node, source), + } + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "function_expression" | "match_expression" => Some(make_named_chunk( + node, + "match".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )), + "match_case" => Some(make_named_chunk( + node, + "case".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )), + "let_expression" => Some(make_named_chunk( + node, + "let".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )), + _ => classify_ocaml_item(node, source), + } + } +} + +fn classify_ocaml_item<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "open_module" => group_candidate(node, "imports", source), + "module_definition" => make_container_chunk( + node, + format!("mod_{}", ocaml_named_text(node, source, &["module_name"])?), + source, + ocaml_module_recurse(node), + ), + "module_type_definition" => make_container_chunk( + node, + format!("modtype_{}", ocaml_named_text(node, source, &["module_type_name"])?), + source, + ocaml_module_type_recurse(node), + ), + "class_definition" => make_container_chunk( + node, + format!("class_{}", ocaml_named_text(node, source, &["class_name"])?), + source, + ocaml_class_recurse(node), + ), + "class_type_definition" => make_container_chunk( + node, + format!("classtype_{}", ocaml_named_text(node, source, &["class_type_name"])?), + source, + ocaml_class_type_recurse(node), + ), + "type_definition" => make_named_chunk( + node, + format!("type_{}", ocaml_named_text(node, source, &["type_constructor"])?), + source, + None, + ), + "exception_definition" => make_named_chunk( + node, + format!("exception_{}", ocaml_named_text(node, source, &["constructor_name"])?), + source, + None, + ), + "value_definition" => classify_ocaml_value_definition(node, source)?, + "value_specification" => make_named_chunk( + node, + format!("val_{}", ocaml_named_text(node, source, &["value_name"])?), + source, + None, + ), + _ => return None, + }) +} + +fn classify_ocaml_value_definition<'t>( + node: Node<'t>, + source: &str, +) -> Option> { + let name = ocaml_named_text(node, source, &["value_name"])?; + let recurse = ocaml_value_recurse(node); + if ocaml_value_definition_is_function(node) { + Some(make_named_chunk(node, format!("fn_{name}"), source, recurse)) + } else { + Some(make_named_chunk(node, format!("val_{name}"), source, recurse)) + } +} + +fn ocaml_named_text(node: Node<'_>, source: &str, kinds: &[&str]) -> Option { + find_named_text(node, source, kinds).and_then(sanitize_identifier) +} + +fn find_named_text<'a>(node: Node<'_>, source: &'a str, kinds: &[&str]) -> Option<&'a str> { + if kinds.iter().any(|kind| node.kind() == *kind) { + return Some(node_text(source, node.start_byte(), node.end_byte())); + } + for child in named_children(node) { + if let Some(text) = find_named_text(child, source, kinds) { + return Some(text); + } + } + None +} + +fn ocaml_module_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::ClassBody, &[], &["module_binding"]).and_then(|binding| { + recurse_into(binding.node, ChunkContext::ClassBody, &["body"], &["structure", "signature"]) + }) +} + +fn ocaml_module_type_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::ClassBody, &["body"], &["signature"]) +} + +fn ocaml_class_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::ClassBody, &[], &["class_binding"]).and_then(|binding| { + recurse_into(binding.node, ChunkContext::ClassBody, &["body"], &["object_expression"]) + }) +} + +fn ocaml_class_type_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::ClassBody, &[], &["class_type_binding"]).and_then(|binding| { + recurse_into(binding.node, ChunkContext::ClassBody, &["body"], &["class_body_type"]) + }) +} + +fn ocaml_method_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &["body"], &[ + "function_expression", + "match_expression", + "let_expression", + ]) +} + +fn ocaml_value_recurse(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &[], &["let_binding"]).and_then(|binding| { + named_children(binding.node) + .into_iter() + .find(|child| { + matches!(child.kind(), "function_expression" | "match_expression" | "let_expression") + }) + .map(|child| RecurseSpec { node: child, context: ChunkContext::FunctionBody }) + }) +} + +fn ocaml_value_definition_is_function(node: Node<'_>) -> bool { + recurse_into(node, ChunkContext::FunctionBody, &[], &["let_binding"]).is_some_and(|binding| { + named_children(binding.node) + .into_iter() + .any(|child| matches!(child.kind(), "parameter" | "function_expression")) + }) +} diff --git a/crates/pi-natives/src/chunk/ast_perl.rs b/crates/pi-natives/src/chunk/ast_perl.rs new file mode 100644 index 000000000..8c3adb916 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_perl.rs @@ -0,0 +1,97 @@ +//! Language-specific chunk classifier for Perl. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct PerlClassifier; + +impl LangClassifier for PerlClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_perl_node(node, source) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_perl_node(node, source) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + kind == "statement_list" + } +} + +fn classify_perl_node<'t>(node: Node<'t>, source: &str) -> Option> { + let body_recurse = || recurse_into(node, ChunkContext::FunctionBody, &["body"], &["block"]); + + Some(match node.kind() { + "package_statement" => { + make_named_chunk(node, format!("mod_{}", perl_name(node, source)?), source, None) + }, + "use_statement" => group_candidate(node, "imports", source), + "subroutine_declaration_statement" => { + make_named_chunk(node, format!("fn_{}", perl_name(node, source)?), source, body_recurse()) + }, + "conditional_statement" => make_candidate( + node, + "if".to_string(), + NameStyle::Named, + None, + body_recurse(), + false, + source, + ), + "for_statement" | "loop_statement" => make_candidate( + node, + "loop".to_string(), + NameStyle::Named, + None, + body_recurse(), + false, + source, + ), + "expression_statement" => classify_perl_statement(node, source), + _ => return None, + }) +} + +fn classify_perl_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + if perl_declares_variable(node) { + group_candidate(node, "decls", source) + } else { + group_candidate(node, "stmts", source) + } +} + +fn perl_declares_variable(node: Node<'_>) -> bool { + if node.kind() == "variable_declaration" { + return true; + } + + if node.kind() == "assignment_expression" + && named_children(node) + .into_iter() + .any(|child| child.kind() == "variable_declaration") + { + return true; + } + + named_children(node).into_iter().any(perl_declares_variable) +} + +fn perl_name(node: Node<'_>, source: &str) -> Option { + find_named_text(node, source, &["bareword", "package", "varname"]).and_then(sanitize_identifier) +} + +fn find_named_text<'a>(node: Node<'_>, source: &'a str, kinds: &[&str]) -> Option<&'a str> { + if kinds.iter().any(|kind| node.kind() == *kind) { + return Some(node_text(source, node.start_byte(), node.end_byte())); + } + + for child in named_children(node) { + if let Some(text) = find_named_text(child, source, kinds) { + return Some(text); + } + } + + None +} diff --git a/crates/pi-natives/src/chunk/ast_powershell.rs b/crates/pi-natives/src/chunk/ast_powershell.rs new file mode 100644 index 000000000..4ca68f68b --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_powershell.rs @@ -0,0 +1,207 @@ +//! PowerShell-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct PowershellClassifier; + +impl LangClassifier for PowershellClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "param_block" => group_candidate(node, "params", source), + "statement_list" => make_container_chunk( + node, + "body".to_string(), + source, + Some(recurse_self(node, ChunkContext::Root)), + ), + "class_statement" => make_container_chunk( + node, + format!("class_{}", powershell_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ), + "function_statement" => make_named_chunk( + node, + format!("fn_{}", powershell_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "pipeline" => classify_powershell_pipeline(node, source), + "switch_statement" | "if_statement" | "foreach_statement" => { + return self.classify_function(node, source); + }, + "flow_control_statement" => group_candidate(node, "stmts", source), + _ => return None, + }) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "class_property_definition" => match powershell_name(node, source) { + Some(name) => make_named_chunk(node, format!("field_{name}"), source, None), + None => group_candidate(node, "fields", source), + }, + "class_method_definition" => classify_class_method(node, source)?, + _ => return None, + }) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "class_method_parameter_list" | "param_block" => group_candidate(node, "params", source), + "script_block" => make_container_chunk( + node, + block_name_for_parent(node).to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "script_block_body" | "statement_block" => make_container_chunk( + node, + "block".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["statement_list"]), + ), + "pipeline" => classify_powershell_pipeline(node, source), + "if_statement" => make_named_chunk( + node, + "if".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["statement_block"]), + ), + "foreach_statement" => make_named_chunk( + node, + "loop".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["statement_block"]), + ), + "switch_statement" => make_named_chunk( + node, + "switch".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["switch_body"]), + ), + "switch_clauses" => make_container_chunk( + node, + "cases".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "switch_clause" => make_named_chunk( + node, + "case".to_string(), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["statement_block"]), + ), + "flow_control_statement" => group_candidate(node, "stmts", source), + _ => return None, + }) + } + + fn is_trivia(&self, kind: &str) -> bool { + matches!(kind, "function_name" | "simple_name" | "type_literal" | "switch_condition") + } +} + +fn classify_class_method<'t>(node: Node<'t>, source: &str) -> Option> { + let name = powershell_name(node, source)?; + let class_name = powershell_name(node.parent()?, source)?; + let chunk_name = if name == "new" || name == class_name { + "constructor".to_string() + } else { + format!("fn_{name}") + }; + + Some(make_named_chunk( + node, + chunk_name, + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )) +} + +fn classify_powershell_pipeline<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + if let Some(command_name) = powershell_command_name(node, source) + && matches!(command_name.as_str(), "using" | "using-module" | "Import-Module") + { + return group_candidate(node, "imports", source); + } + + if let Some((name, script_block)) = assigned_script_block(node, source) { + return make_container_chunk_from( + node, + node, + format!("block_{name}"), + source, + Some(recurse_self(script_block, ChunkContext::FunctionBody)), + ); + } + + if child_by_kind(node, &["assignment_expression"]).is_some() { + group_candidate(node, "decls", source) + } else { + group_candidate(node, "stmts", source) + } +} + +fn assigned_script_block<'t>(node: Node<'t>, source: &str) -> Option<(String, Node<'t>)> { + let assignment = child_by_kind(node, &["assignment_expression"])?; + let lhs = child_by_kind(assignment, &["left_assignment_expression"])?; + let name = sanitize_identifier( + node_text(source, lhs.start_byte(), lhs.end_byte()).trim_start_matches('$'), + )?; + let script_block = named_children(assignment) + .into_iter() + .filter(|child| child.kind() != "left_assignment_expression") + .find_map(find_script_block)?; + Some((name, script_block)) +} + +fn find_script_block(node: Node<'_>) -> Option> { + if node.kind() == "script_block" { + return Some(node); + } + for child in named_children(node) { + if let Some(script_block) = find_script_block(child) { + return Some(script_block); + } + } + None +} + +fn block_name_for_parent(node: Node<'_>) -> &'static str { + match node.parent().map(|parent| parent.kind()) { + Some("function_statement" | "class_method_definition") => "body", + _ => "block", + } +} + +fn powershell_name(node: Node<'_>, source: &str) -> Option { + find_named_text(node, source, &[ + "function_name", + "simple_name", + "member_name", + "type_identifier", + "variable", + ]) + .and_then(|text| sanitize_identifier(text.trim_start_matches('$'))) +} + +fn powershell_command_name(node: Node<'_>, source: &str) -> Option { + find_named_text(node, source, &["command_name"]).and_then(sanitize_identifier) +} + +fn find_named_text<'a>(node: Node<'_>, source: &'a str, kinds: &[&str]) -> Option<&'a str> { + if kinds.iter().any(|kind| node.kind() == *kind) { + return Some(node_text(source, node.start_byte(), node.end_byte())); + } + + for child in named_children(node) { + if let Some(text) = find_named_text(child, source, kinds) { + return Some(text); + } + } + + None +} diff --git a/crates/pi-natives/src/chunk/ast_proto.rs b/crates/pi-natives/src/chunk/ast_proto.rs new file mode 100644 index 000000000..0d1ba8510 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_proto.rs @@ -0,0 +1,111 @@ +//! Chunk classifier for Protocol Buffers. +//! +//! Mirror the grammar's declaration structure directly: the root owns headers, +//! imports, options, messages, enums, and services; message bodies own fields, +//! oneofs, and nested messages/enums; services own rpc declarations and service +//! options; rpc blocks may contain rpc-scoped options. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct ProtoClassifier; + +impl LangClassifier for ProtoClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_proto_root(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_proto_class(node, source) + } +} + +fn classify_proto_root<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "syntax" | "package" => group_candidate(node, "headers", source), + "import" => group_candidate(node, "imports", source), + "option" => group_candidate(node, "options", source), + "message" => make_container_chunk( + node, + format!("msg_{}", proto_name(node, source)?), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["message_body"]), + ), + "enum" => make_container_chunk( + node, + format!("enum_{}", proto_name(node, source)?), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["enum_body"]), + ), + "service" => make_container_chunk( + node, + format!("service_{}", proto_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ), + _ => return None, + }) +} + +fn classify_proto_class<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "option" => group_candidate(node, "options", source), + "field" if is_proto_message_field(node) => { + make_named_chunk(node, format!("field_{}", proto_name(node, source)?), source, None) + }, + "oneof" => make_container_chunk( + node, + format!("oneof_{}", proto_name(node, source)?), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ), + "oneof_field" => { + make_named_chunk(node, format!("field_{}", proto_name(node, source)?), source, None) + }, + "message" => make_container_chunk( + node, + format!("msg_{}", proto_name(node, source)?), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["message_body"]), + ), + "enum" => make_container_chunk( + node, + format!("enum_{}", proto_name(node, source)?), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["enum_body"]), + ), + "enum_field" => { + make_named_chunk(node, format!("variant_{}", proto_name(node, source)?), source, None) + }, + "rpc" => make_container_chunk( + node, + format!("rpc_{}", proto_name(node, source)?), + source, + proto_rpc_recurse(node), + ), + _ => return None, + }) +} + +fn is_proto_message_field(node: Node<'_>) -> bool { + node + .parent() + .is_some_and(|parent| parent.kind() == "message_body") +} + +fn proto_rpc_recurse(node: Node<'_>) -> Option> { + let has_nested_option = named_children(node) + .into_iter() + .any(|child| child.kind() == "option"); + if has_nested_option { + Some(recurse_self(node, ChunkContext::ClassBody)) + } else { + None + } +} + +fn proto_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["message_name", "enum_name", "service_name", "rpc_name", "identifier"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) +} diff --git a/crates/pi-natives/src/chunk/ast_r.rs b/crates/pi-natives/src/chunk/ast_r.rs new file mode 100644 index 000000000..f6b832050 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_r.rs @@ -0,0 +1,144 @@ +//! R-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct RClassifier; + +impl LangClassifier for RClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + // ── Imports ── + "call" if is_import_call(node, source) => group_candidate(node, "imports", source), + "call" => group_candidate(node, "stmts", source), + + // ── Function / value assignments ── + "binary_operator" => classify_assignment(node, source, ChunkScope::Root)?, + + // ── Control flow at script scope ── + "if_statement" => control_candidate(node, "if", source, recurse_if(node)), + "for_statement" | "while_statement" | "repeat_statement" => { + control_candidate(node, "loop", source, recurse_loop(node)) + }, + + // ── Bare expressions ── + "identifier" | "subset" | "subset2" | "extract_operator" => { + group_candidate(node, "stmts", source) + }, + + _ => return None, + }) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + // ── Local assignments ── + "binary_operator" => classify_assignment(node, source, ChunkScope::Function)?, + + // ── Control flow ── + "if_statement" => control_candidate(node, "if", source, recurse_if(node)), + "for_statement" | "while_statement" | "repeat_statement" => { + control_candidate(node, "loop", source, recurse_loop(node)) + }, + + // ── Calls / bare expressions ── + "call" | "identifier" | "subset" | "subset2" | "extract_operator" | "break" | "next" + | "return" => group_candidate(node, "stmts", source), + + _ => return None, + }) + } +} + +#[derive(Clone, Copy)] +enum ChunkScope { + Root, + Function, +} + +fn classify_assignment<'t>( + node: Node<'t>, + source: &str, + scope: ChunkScope, +) -> Option> { + let (lhs, rhs) = assignment_sides(node, source)?; + + if rhs.kind() == "function_definition" { + let name = simple_lhs_name(lhs, source).unwrap_or_else(|| "anonymous".to_string()); + return Some(make_named_chunk_from( + node, + rhs, + format!("fn_{name}"), + source, + recurse_body(rhs, ChunkContext::FunctionBody), + )); + } + + match (scope, simple_lhs_name(lhs, source)) { + (ChunkScope::Root, Some(name)) => { + Some(make_named_chunk(node, format!("var_{name}"), source, None)) + }, + (ChunkScope::Function, Some(name)) if spans_multiple_lines(node) => { + Some(make_named_chunk(node, format!("var_{name}"), source, None)) + }, + _ => Some(group_candidate( + node, + match scope { + ChunkScope::Root => "decls", + ChunkScope::Function => "stmts", + }, + source, + )), + } +} + +fn assignment_sides<'t>(node: Node<'t>, source: &str) -> Option<(Node<'t>, Node<'t>)> { + if node.kind() != "binary_operator" { + return None; + } + + let operator = node.child_by_field_name("operator")?; + let operator_text = node_text(source, operator.start_byte(), operator.end_byte()); + if !matches!(operator_text, "<-" | "<<-" | "=") { + return None; + } + + Some((node.child_by_field_name("lhs")?, node.child_by_field_name("rhs")?)) +} + +fn simple_lhs_name(lhs: Node<'_>, source: &str) -> Option { + (lhs.kind() == "identifier") + .then(|| extract_identifier(lhs, source)) + .flatten() +} + +fn is_import_call(node: Node<'_>, source: &str) -> bool { + matches!( + extract_identifier(node, source).as_deref(), + Some("library" | "require" | "requireNamespace" | "source") + ) +} + +fn recurse_if(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &["consequence", "alternative"], &[ + "braced_expression", + ]) +} + +fn recurse_loop(node: Node<'_>) -> Option> { + recurse_into(node, ChunkContext::FunctionBody, &["body"], &["braced_expression"]) +} + +fn control_candidate<'t>( + node: Node<'t>, + name: &str, + source: &str, + recurse: Option>, +) -> RawChunkCandidate<'t> { + make_candidate(node, name.to_string(), NameStyle::Named, None, recurse, true, source) +} + +fn spans_multiple_lines(node: Node<'_>) -> bool { + node.start_position().row != node.end_position().row +} diff --git a/crates/pi-natives/src/chunk/ast_sql.rs b/crates/pi-natives/src/chunk/ast_sql.rs new file mode 100644 index 000000000..56e3c71f7 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_sql.rs @@ -0,0 +1,232 @@ +//! SQL-specific chunk classifier. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct SqlClassifier; + +impl LangClassifier for SqlClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_sql_root(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_sql_class(node, source) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_sql_function(node, source) + } + + fn is_trivia(&self, kind: &str) -> bool { + matches!(kind, "empty_statement" | "dollar_quote" | "keyword_from") + } +} + +fn classify_sql_root<'t>(node: Node<'t>, source: &str) -> Option> { + if node.kind() == "statement" { + return classify_sql_statement_root(node, source); + } + + classify_sql_root_node(node, node, source).or_else(|| classify_sql_query_node(node, source)) +} + +fn classify_sql_statement_root<'t>(node: Node<'t>, source: &str) -> Option> { + let children = named_children(node); + if children.len() == 1 { + return classify_sql_root_node(node, children[0], source) + .or_else(|| classify_sql_query_node(children[0], source)); + } + + if children.iter().any(|child| is_sql_query_kind(child.kind())) { + return Some(make_container_chunk( + node, + "query".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )); + } + + None +} + +fn classify_sql_root_node<'t>( + range_node: Node<'t>, + node: Node<'t>, + source: &str, +) -> Option> { + Some(match node.kind() { + "create_schema" => make_named_chunk_from( + range_node, + node, + format!( + "schema_{}", + extract_sql_identifier(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + ), + "create_table" => make_container_chunk_from( + range_node, + node, + format!( + "table_{}", + extract_sql_object_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::ClassBody, &[], &["column_definitions"]), + ), + "create_view" => make_container_chunk_from( + range_node, + node, + format!( + "view_{}", + extract_sql_object_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["create_query"]), + ), + "create_materialized_view" => make_container_chunk_from( + range_node, + node, + format!( + "matview_{}", + extract_sql_object_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["create_query"]), + ), + "create_function" => make_container_chunk_from( + range_node, + node, + format!( + "fn_{}", + extract_sql_object_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_sql_function_query(node), + ), + "create_trigger" => make_named_chunk_from( + range_node, + node, + format!( + "trigger_{}", + extract_sql_object_name(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + ), + "create_index" => make_named_chunk_from( + range_node, + node, + format!( + "index_{}", + extract_sql_identifier(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + ), + _ => return None, + }) +} + +fn classify_sql_class<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "column_definition" => make_named_chunk( + node, + format!( + "field_{}", + extract_sql_identifier(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + None, + ), + _ => return None, + }) +} + +fn classify_sql_function<'t>(node: Node<'t>, source: &str) -> Option> { + if node.kind() == "statement" { + return Some(make_container_chunk( + node, + "query".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + )); + } + + classify_sql_query_node(node, source) +} + +fn classify_sql_query_node<'t>(node: Node<'t>, source: &str) -> Option> { + Some(match node.kind() { + "insert" => group_candidate(node, "stmts", source), + "keyword_with" => group_candidate(node, "with", source), + "cte" => make_container_chunk( + node, + format!( + "cte_{}", + extract_sql_identifier(node, source).unwrap_or_else(|| "anonymous".to_string()) + ), + source, + recurse_into(node, ChunkContext::FunctionBody, &[], &["statement"]), + ), + "select" => positional_candidate(node, "select", source), + "from" => make_container_chunk( + node, + "from".to_string(), + source, + Some(recurse_self(node, ChunkContext::FunctionBody)), + ), + "relation" => group_candidate(node, "relations", source), + "join" => positional_candidate(node, "join", source), + "where" => positional_candidate(node, "where", source), + "group_by" => positional_candidate(node, "group_by", source), + "order_by" => positional_candidate(node, "order_by", source), + _ => return None, + }) +} + +fn recurse_sql_function_query(node: Node<'_>) -> Option> { + let body = child_by_kind(node, &["function_body"])?; + recurse_into(body, ChunkContext::FunctionBody, &[], &["statement"]) +} + +fn is_sql_query_kind(kind: &str) -> bool { + matches!( + kind, + "insert" + | "keyword_with" + | "cte" + | "select" + | "from" + | "where" + | "group_by" + | "order_by" + | "join" + ) +} + +fn extract_sql_identifier(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["identifier"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) +} + +fn extract_sql_object_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["object_reference"]).and_then(|name| last_identifier(name, source)) +} + +fn last_identifier(node: Node<'_>, source: &str) -> Option { + if node.kind() == "identifier" { + return sanitize_identifier(node_text(source, node.start_byte(), node.end_byte())); + } + + for child in named_children(node).into_iter().rev() { + if let Some(identifier) = last_identifier(child, source) { + return Some(identifier); + } + } + + None +} diff --git a/crates/pi-natives/src/chunk/ast_svelte.rs b/crates/pi-natives/src/chunk/ast_svelte.rs new file mode 100644 index 000000000..0997d733e --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_svelte.rs @@ -0,0 +1,218 @@ +//! Language-specific chunk classifier for Svelte. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct SvelteClassifier; + +impl LangClassifier for SvelteClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_svelte_node(node, source, true) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_svelte_node(node, source, false) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_svelte_node(node, source, false) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + kind == "document" + } +} + +fn classify_svelte_node<'t>( + node: Node<'t>, + source: &str, + include_plain_elements: bool, +) -> Option> { + match node.kind() { + "script_element" => Some(classify_script_element(node, source)), + "style_element" => Some(classify_style_element(node, source)), + "snippet_statement" => Some(classify_snippet_statement(node, source)), + "if_statement" => Some(classify_if_statement(node, source)), + "else_if_statement" => Some(classify_else_if_statement(node, source)), + "else_statement" => Some(make_block_chunk(node, "else", source)), + "each_statement" => Some(classify_each_statement(node, source)), + "await_statement" => Some(classify_await_statement(node, source)), + "then_statement" => Some(classify_then_statement(node, source)), + "catch_statement" => Some(classify_catch_statement(node, source)), + "render_expr" => Some(classify_render_expr(node, source)), + "html_interpolation" => Some(group_candidate(node, "html", source)), + "interpolation" => Some(group_candidate(node, "interpolation", source)), + "expression" => Some(group_candidate(node, "expr", source)), + "element" if include_plain_elements || element_has_structure(node) => { + classify_element(node, source) + }, + _ => None, + } +} + +fn classify_script_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "module", source) + || attribute_value(node, "context", source).as_deref() == Some("module") + { + "script_module".to_string() + } else { + "script".to_string() + }; + + // The grammar exposes script contents as a single `raw_text` child, so the + // element boundary is the most truthful chunk. + make_named_chunk(node, name, source, None) +} + +fn classify_style_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "scoped", source) { + "style_scoped".to_string() + } else { + "style".to_string() + }; + make_named_chunk(node, name, source, None) +} + +fn classify_snippet_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = child_by_kind(node, &["snippet_start_expr"]) + .and_then(|start| child_by_kind(start, &["snippet_name"])) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) + .map_or_else(|| "snippet".to_string(), |name| format!("snippet_{name}")); + make_block_chunk(node, &name, source) +} + +fn classify_if_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "if_start_expr", &["raw_text_expr"], "if"); + make_block_chunk(node, &name, source) +} + +fn classify_else_if_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "else_if_expr", &["raw_text_expr"], "else_if"); + make_block_chunk(node, &name, source) +} + +fn classify_each_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "each_start_expr", &["raw_text_each"], "each"); + make_block_chunk(node, &name, source) +} + +fn classify_await_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "await_start_expr", &["raw_text_expr"], "await"); + make_block_chunk(node, &name, source) +} + +fn classify_then_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "then_expr", &["raw_text_expr"], "then"); + make_block_chunk(node, &name, source) +} + +fn classify_catch_statement<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = block_expr_name(node, source, "catch_expr", &["raw_text_expr"], "catch"); + make_block_chunk(node, &name, source) +} + +fn classify_render_expr<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = child_by_kind(node, &["snippet_name"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) + .map_or_else(|| "render".to_string(), |name| format!("render_{name}")); + make_named_chunk(node, name, source, None) +} + +fn classify_element<'t>(node: Node<'t>, source: &str) -> Option> { + let tag_name = extract_markup_tag_name(node, source)?; + Some(make_container_chunk( + node, + format!("tag_{tag_name}"), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + )) +} + +fn make_block_chunk<'t>(node: Node<'t>, name: &str, source: &str) -> RawChunkCandidate<'t> { + make_container_chunk( + node, + name.to_string(), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ) +} + +fn block_expr_name( + node: Node<'_>, + source: &str, + header_kind: &str, + expr_kinds: &[&str], + prefix: &str, +) -> String { + child_by_kind(node, &[header_kind]) + .and_then(|header| child_by_kind(header, expr_kinds)) + .and_then(|expr| sanitize_identifier(node_text(source, expr.start_byte(), expr.end_byte()))) + .map_or_else(|| prefix.to_string(), |expr| format!("{prefix}_{expr}")) +} + +fn element_has_structure(node: Node<'_>) -> bool { + named_children(node).into_iter().any(|child| { + matches!( + child.kind(), + "snippet_statement" + | "if_statement" + | "else_if_statement" + | "else_statement" + | "each_statement" + | "await_statement" + | "then_statement" + | "catch_statement" + | "render_expr" + | "html_interpolation" + | "interpolation" + | "expression" + | "element" + ) + }) +} + +fn extract_markup_tag_name(node: Node<'_>, source: &str) -> Option { + start_like(node) + .and_then(|start| child_by_kind(start, &["tag_name"])) + .and_then(|tag| sanitize_identifier(node_text(source, tag.start_byte(), tag.end_byte()))) +} + +fn has_attribute(node: Node<'_>, name: &str, source: &str) -> bool { + start_like(node) + .into_iter() + .flat_map(named_children) + .filter(|child| child.kind() == "attribute") + .filter_map(|attr| extract_attribute_name(attr, source)) + .any(|attr_name| attr_name == name) +} + +fn attribute_value(node: Node<'_>, name: &str, source: &str) -> Option { + let start = start_like(node)?; + for child in named_children(start) { + if child.kind() != "attribute" { + continue; + } + if extract_attribute_name(child, source).as_deref() != Some(name) { + continue; + } + if let Some(value) = child_by_kind(child, &["attribute_value", "quoted_attribute_value"]) { + return sanitize_identifier(&unquote_text(node_text( + source, + value.start_byte(), + value.end_byte(), + ))); + } + return Some(name.to_string()); + } + None +} + +fn extract_attribute_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["attribute_name"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) +} + +fn start_like(node: Node<'_>) -> Option> { + child_by_kind(node, &["start_tag", "self_closing_tag"]) +} diff --git a/crates/pi-natives/src/chunk/ast_tlaplus.rs b/crates/pi-natives/src/chunk/ast_tlaplus.rs index dc4ba4e5d..94af892a9 100644 --- a/crates/pi-natives/src/chunk/ast_tlaplus.rs +++ b/crates/pi-natives/src/chunk/ast_tlaplus.rs @@ -265,10 +265,14 @@ fn translation_chunk( line_count: range.end_line.saturating_sub(range.start_line) + 1, start_byte, end_byte, + checksum_start_byte: start_byte, + body_start_byte: None, + body_end_byte: None, checksum, error: false, indent: 0, indent_char: String::new(), + group: false, } } diff --git a/crates/pi-natives/src/chunk/ast_vue.rs b/crates/pi-natives/src/chunk/ast_vue.rs new file mode 100644 index 000000000..e12fa70f6 --- /dev/null +++ b/crates/pi-natives/src/chunk/ast_vue.rs @@ -0,0 +1,223 @@ +//! Language-specific chunk classifier for Vue single-file components. + +use tree_sitter::Node; + +use super::{classify::LangClassifier, common::*}; + +pub struct VueClassifier; + +impl LangClassifier for VueClassifier { + fn classify_root<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_root_node(node, source) + } + + fn classify_class<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_nested_node(node, source) + } + + fn classify_function<'t>(&self, node: Node<'t>, source: &str) -> Option> { + classify_nested_node(node, source) + } + + fn is_root_wrapper(&self, kind: &str) -> bool { + kind == "document" + } +} + +fn classify_root_node<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "template_element" => Some(classify_template_element(node, source)), + "script_element" => Some(classify_script_element(node, source)), + "style_element" => Some(classify_style_element(node, source)), + // Vue custom blocks (for example ) currently parse as plain `element` + // nodes at the document root, so infer custom-block semantics from position. + "element" => Some(classify_custom_block(node, source)), + _ => None, + } +} + +fn classify_nested_node<'t>(node: Node<'t>, source: &str) -> Option> { + match node.kind() { + "template_element" => Some(classify_template_element(node, source)), + "element" => classify_element(node, source), + "start_tag" => classify_start_tag(node, source), + "directive_attribute" => Some(classify_directive_attribute(node, source)), + "attribute" => Some(classify_attribute(node, source)), + "interpolation" => Some(make_named_chunk(node, "expr".to_string(), source, None)), + "text" => Some(group_candidate(node, "text", source)), + "raw_text" => Some(group_candidate(node, "text", source)), + _ => None, + } +} + +fn classify_template_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = extract_slot_name(node, source) + .map_or_else(|| "template".to_string(), |slot_name| format!("slot_{slot_name}")); + make_container_chunk(node, name, source, Some(recurse_self(node, ChunkContext::ClassBody))) +} + +fn classify_script_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "setup", source) { + "script_setup".to_string() + } else if attribute_value(node, "context", source).as_deref() == Some("module") { + "script_module".to_string() + } else { + "script".to_string() + }; + // tree-sitter-vue exposes script bodies as `raw_text`, not injected JS/TS. + make_named_chunk(node, name, source, None) +} + +fn classify_style_element<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = if has_attribute(node, "scoped", source) { + "style_scoped".to_string() + } else { + "style".to_string() + }; + // Styles are likewise exposed as `raw_text`, so preserve only the SFC block. + make_named_chunk(node, name, source, None) +} + +fn classify_custom_block<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let tag_name = extract_markup_tag_name(node, source).unwrap_or_else(|| "anonymous".to_string()); + make_container_chunk( + node, + format!("custom_{tag_name}"), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + ) +} + +fn classify_element<'t>(node: Node<'t>, source: &str) -> Option> { + let tag_name = extract_markup_tag_name(node, source)?; + Some(make_container_chunk( + node, + format!("tag_{tag_name}"), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + )) +} + +fn classify_start_tag<'t>(node: Node<'t>, source: &str) -> Option> { + if !named_children(node) + .into_iter() + .any(|child| matches!(child.kind(), "attribute" | "directive_attribute")) + { + return None; + } + let tag_name = child_by_kind(node, &["tag_name"]) + .and_then(|tag| sanitize_identifier(node_text(source, tag.start_byte(), tag.end_byte()))) + .unwrap_or_else(|| "anonymous".to_string()); + Some(make_container_chunk( + node, + format!("attrs_{tag_name}"), + source, + Some(recurse_self(node, ChunkContext::ClassBody)), + )) +} + +fn classify_attribute<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let name = child_by_kind(node, &["attribute_name"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) + .unwrap_or_else(|| "attr".to_string()); + make_named_chunk(node, format!("attr_{name}"), source, None) +} + +fn classify_directive_attribute<'t>(node: Node<'t>, source: &str) -> RawChunkCandidate<'t> { + let raw = node_text(source, node.start_byte(), node.end_byte()).trim(); + let directive_name = + extract_directive_name(node, source).unwrap_or_else(|| "directive".to_string()); + let modifier_suffix = extract_directive_modifiers(node, source) + .filter(|mods| !mods.is_empty()) + .map(|mods| format!("_{mods}")) + .unwrap_or_default(); + let chunk_name = if raw.starts_with('@') { + format!("on_{directive_name}{modifier_suffix}") + } else if raw.starts_with(':') { + format!("bind_{directive_name}{modifier_suffix}") + } else if raw.starts_with('#') { + format!("slot_{directive_name}{modifier_suffix}") + } else { + format!("dir_{directive_name}{modifier_suffix}") + }; + make_named_chunk(node, chunk_name, source, None) +} + +fn extract_markup_tag_name(node: Node<'_>, source: &str) -> Option { + start_like(node) + .and_then(|start| child_by_kind(start, &["tag_name"])) + .and_then(|tag| sanitize_identifier(node_text(source, tag.start_byte(), tag.end_byte()))) +} + +fn extract_slot_name(node: Node<'_>, source: &str) -> Option { + let start = start_like(node)?; + named_children(start) + .into_iter() + .find(|child| { + node_text(source, child.start_byte(), child.end_byte()) + .trim() + .starts_with('#') + }) + .and_then(|child| extract_directive_name(child, source)) +} + +fn extract_directive_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["directive_name", "directive_value"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) +} + +fn extract_directive_modifiers(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["directive_modifiers"]) + .and_then(|mods| sanitize_identifier(node_text(source, mods.start_byte(), mods.end_byte()))) +} + +fn has_attribute(node: Node<'_>, name: &str, source: &str) -> bool { + start_like(node) + .into_iter() + .flat_map(named_children) + .filter(|child| matches!(child.kind(), "attribute" | "directive_attribute")) + .filter_map(|attr| extract_attribute_name(attr, source)) + .any(|attr_name| attr_name == name) +} + +fn attribute_value(node: Node<'_>, name: &str, source: &str) -> Option { + let start = start_like(node)?; + for child in named_children(start) { + if !matches!(child.kind(), "attribute" | "directive_attribute") { + continue; + } + if extract_attribute_name(child, source).as_deref() != Some(name) { + continue; + } + if let Some(value) = + child_by_kind(child, &["attribute_value", "quoted_attribute_value", "directive_value"]) + { + return sanitize_identifier(&unquote_text(node_text( + source, + value.start_byte(), + value.end_byte(), + ))); + } + return Some(name.to_string()); + } + None +} + +fn extract_attribute_name(node: Node<'_>, source: &str) -> Option { + child_by_kind(node, &["attribute_name", "directive_name"]) + .and_then(|name| sanitize_identifier(node_text(source, name.start_byte(), name.end_byte()))) + .or_else(|| { + if node_text(source, node.start_byte(), node.end_byte()) + .trim() + .starts_with('#') + { + extract_directive_name(node, source) + } else { + None + } + }) +} + +fn start_like(node: Node<'_>) -> Option> { + child_by_kind(node, &["start_tag", "self_closing_tag"]) +} diff --git a/crates/pi-natives/src/chunk/classify.rs b/crates/pi-natives/src/chunk/classify.rs index 20bfb8244..2e77090c5 100644 --- a/crates/pi-natives/src/chunk/classify.rs +++ b/crates/pi-natives/src/chunk/classify.rs @@ -45,7 +45,7 @@ pub trait LangClassifier { } /// Post-process the chunk tree after initial construction. - /// Used for structural transformations like Go receiver reparenting. + /// Used for structural transformations fn post_process( &self, _chunks: &mut Vec, @@ -86,6 +86,7 @@ pub trait LangClassifier { /// Resolve a [`LangClassifier`] for the given language. pub fn classifier_for(lang: &str) -> &'static dyn LangClassifier { match lang { + "astro" => &super::ast_astro::AstroClassifier, // JS / TS family "javascript" | "js" | "jsx" | "typescript" | "ts" | "tsx" => { &super::ast_js_ts::JsTsClassifier @@ -102,26 +103,54 @@ pub fn classifier_for(lang: &str) -> &'static dyn LangClassifier { "csharp" | "java" => &super::ast_csharp_java::CSharpJavaClassifier, // Clojure "clojure" => &super::ast_clojure::ClojureClassifier, - // Elixir - "elixir" => &super::ast_elixir::ElixirClassifier, - // Ruby / Lua - "ruby" | "lua" => &super::ast_ruby_lua::RubyLuaClassifier, - // Haskell / Scala - "haskell" | "scala" => &super::ast_haskell_scala::HaskellScalaClassifier, + // CMake + "cmake" => &super::ast_cmake::CMakeClassifier, // CSS "css" => &super::ast_css::CssClassifier, - // HTML / XML - "html" | "xml" => &super::ast_html_xml::HtmlXmlClassifier, // Data formats "json" | "toml" | "yaml" => &super::ast_data_formats::DataFormatsClassifier, - // Nix / HCL - "nix" | "hcl" => &super::ast_nix_hcl::NixHclClassifier, + // Dockerfile + "dockerfile" => &super::ast_dockerfile::DockerfileClassifier, + // Elixir + "elixir" => &super::ast_elixir::ElixirClassifier, + // Erlang + "erlang" => &super::ast_erlang::ErlangClassifier, + // GraphQL + "graphql" => &super::ast_graphql::GraphqlClassifier, + // Haskell / Scala + "haskell" | "scala" => &super::ast_haskell_scala::HaskellScalaClassifier, + // HTML / XML + "html" | "xml" => &super::ast_html_xml::HtmlXmlClassifier, + // INI + "ini" => &super::ast_ini::IniClassifier, + // Just + "just" => &super::ast_just::JustClassifier, // Markdown / Handlebars "markdown" | "handlebars" => &super::ast_markup::MarkupClassifier, + // Nix / HCL + "nix" | "hcl" => &super::ast_nix_hcl::NixHclClassifier, + // OCaml + "ocaml" => &super::ast_ocaml::OcamlClassifier, + // Perl + "perl" => &super::ast_perl::PerlClassifier, + // PowerShell + "powershell" => &super::ast_powershell::PowershellClassifier, + // Protobuf + "protobuf" | "proto" => &super::ast_proto::ProtoClassifier, + // R + "r" => &super::ast_r::RClassifier, + // Ruby / Lua + "ruby" | "lua" => &super::ast_ruby_lua::RubyLuaClassifier, + // SQL + "sql" => &super::ast_sql::SqlClassifier, + // Svelte + "svelte" => &super::ast_svelte::SvelteClassifier, // TLA+ / PlusCal "tlaplus" | "pluscal" | "pcal" | "tla" | "tla+" => &super::ast_tlaplus::TlaplusClassifier, // Bash / Make / Diff "bash" | "make" | "diff" => &super::ast_bash_make_diff::ShellBuildClassifier, + // Vue + "vue" => &super::ast_vue::VueClassifier, // Everything else (Kotlin, Swift, PHP, Solidity, etc.) _ => &super::ast_misc::MiscClassifier, } diff --git a/crates/pi-natives/src/highlight.rs b/crates/pi-natives/src/highlight.rs index e2ba692da..6bd9642d1 100644 --- a/crates/pi-natives/src/highlight.rs +++ b/crates/pi-natives/src/highlight.rs @@ -168,9 +168,8 @@ const LANG_ALIASES: &[(&[&str], &str)] = &[ (&["cs", "csharp"], "C#"), (&["php"], "PHP"), (&["sh", "bash", "zsh", "shell"], "Bash"), - (&["fish"], "Shell-Unix-Generic"), (&["ps1", "powershell"], "PowerShell"), - (&["html", "htm"], "HTML"), + (&["html", "htm", "astro", "vue", "svelte"], "HTML"), (&["css"], "CSS"), (&["scss"], "SCSS"), (&["sass"], "Sass"), @@ -182,7 +181,7 @@ const LANG_ALIASES: &[(&[&str], &str)] = &[ (&["md", "markdown"], "Markdown"), (&["sql"], "SQL"), (&["lua"], "Lua"), - (&["perl", "pl"], "Perl"), + (&["perl", "pl", "pm"], "Perl"), (&["r"], "R"), (&["scala"], "Scala"), (&["clj", "clojure"], "Clojure"), @@ -194,9 +193,9 @@ const LANG_ALIASES: &[(&[&str], &str)] = &[ (&["graphql", "gql"], "GraphQL"), (&["proto", "protobuf"], "Protocol Buffers"), (&["tf", "hcl", "terraform"], "Terraform"), - (&["dockerfile", "docker"], "Dockerfile"), - (&["makefile", "make"], "Makefile"), - (&["cmake"], "CMake"), + (&["dockerfile", "docker", "containerfile"], "Dockerfile"), + (&["makefile", "make", "just", "justfile"], "Makefile"), + (&["cmake", "cmakelists"], "CMake"), (&["ini", "cfg", "conf", "config", "properties"], "INI"), (&["diff", "patch"], "Diff"), (&["gitignore", "gitattributes", "gitmodules"], "Git Ignore"), diff --git a/crates/pi-natives/src/language/mod.rs b/crates/pi-natives/src/language/mod.rs index 0414b327a..f566376c1 100644 --- a/crates/pi-natives/src/language/mod.rs +++ b/crates/pi-natives/src/language/mod.rs @@ -107,18 +107,30 @@ macro_rules! impl_lang_expando { impl_lang_expando!(C, language_c, '𐀀'); impl_lang_expando!(Cpp, language_cpp, '𐀀'); impl_lang_expando!(CSharp, language_c_sharp, 'µ'); +impl_lang_expando!(Cmake, language_cmake, 'µ'); impl_lang_expando!(Css, language_css, '_'); +impl_lang_expando!(Dockerfile, language_dockerfile, 'µ'); impl_lang_expando!(Elixir, language_elixir, 'µ'); +impl_lang_expando!(Erlang, language_erlang, 'µ'); impl_lang_expando!(Go, language_go, 'µ'); +impl_lang!(Graphql, language_graphql); impl_lang!(Handlebars, language_handlebars); impl_lang_expando!(Haskell, language_haskell, 'µ'); impl_lang_expando!(Hcl, language_hcl, 'µ'); +impl_lang_expando!(Ini, language_ini, 'µ'); +impl_lang_expando!(Just, language_just, 'µ'); impl_lang_expando!(Kotlin, language_kotlin, 'µ'); impl_lang_expando!(Nix, language_nix, '_'); +impl_lang_expando!(Ocaml, language_ocaml, 'µ'); +impl_lang_expando!(Perl, language_perl, 'µ'); impl_lang_expando!(Php, language_php, 'µ'); +impl_lang_expando!(Powershell, language_powershell, 'µ'); +impl_lang_expando!(Proto, language_proto, 'µ'); impl_lang_expando!(Python, language_python, 'µ'); +impl_lang_expando!(R, language_r, 'µ'); impl_lang_expando!(Ruby, language_ruby, 'µ'); impl_lang_expando!(Rust, language_rust, 'µ'); +impl_lang_expando!(Sql, language_sql, 'µ'); impl_lang_expando!(Swift, language_swift, 'µ'); // New expando languages @@ -133,6 +145,7 @@ impl_lang_expando!(Tlaplus, language_tlaplus, 'µ'); // ── Stub languages ($ accepted in grammar) ────────────────────────────── +impl_lang!(Astro, language_astro); impl_lang!(Bash, language_bash); impl_lang!(Clojure, language_clojure); impl_lang!(Java, language_java); @@ -141,8 +154,10 @@ impl_lang!(Json, language_json); impl_lang!(Lua, language_lua); impl_lang!(Scala, language_scala); impl_lang!(Solidity, language_solidity); +impl_lang!(Svelte, language_svelte); impl_lang!(Tsx, language_tsx); impl_lang!(TypeScript, language_typescript); +impl_lang!(Vue, language_vue); impl_lang!(Yaml, language_yaml); // New stub languages @@ -252,22 +267,29 @@ fn node_to_range(node: &Node) -> TSRange { /// All supported languages for ast-grep structural search/replace. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum SupportLang { + Astro, Bash, C, + Cmake, Cpp, CSharp, Clojure, Css, Diff, + Dockerfile, Elixir, + Erlang, Go, + Graphql, Handlebars, Haskell, Hcl, Html, + Ini, Java, JavaScript, Json, + Just, Julia, Kotlin, Lua, @@ -275,21 +297,29 @@ pub enum SupportLang { Markdown, Nix, ObjC, + Ocaml, Odin, + Perl, Php, + Powershell, + Proto, Python, + R, Regex, Ruby, Rust, Scala, Solidity, + Sql, Starlark, + Svelte, Swift, Toml, Tlaplus, Tsx, TypeScript, Verilog, + Vue, Xml, Yaml, Zig, @@ -305,10 +335,11 @@ impl SupportLang { pub const fn all_langs() -> &'static [Self] { use SupportLang::*; &[ - Bash, C, Cpp, CSharp, Clojure, Css, Diff, Elixir, Go, Handlebars, Haskell, Hcl, Html, - Java, JavaScript, Json, Julia, Kotlin, Lua, Make, Markdown, Nix, ObjC, Odin, Php, Python, - Regex, Ruby, Rust, Scala, Solidity, Starlark, Swift, Toml, Tlaplus, Tsx, TypeScript, - Verilog, Xml, Yaml, Zig, + Astro, Bash, C, Cmake, Cpp, CSharp, Clojure, Css, Diff, Dockerfile, Elixir, Erlang, Go, + Graphql, Handlebars, Haskell, Hcl, Html, Ini, Java, JavaScript, Json, Just, Julia, Kotlin, + Lua, Make, Markdown, Nix, ObjC, Ocaml, Odin, Perl, Php, Powershell, Proto, Python, R, + Regex, Ruby, Rust, Scala, Solidity, Sql, Starlark, Svelte, Swift, Toml, Tlaplus, Tsx, + TypeScript, Verilog, Vue, Xml, Yaml, Zig, ] } @@ -316,22 +347,29 @@ impl SupportLang { /// file-type inference results, and error messages. pub const fn canonical_name(self) -> &'static str { match self { + Self::Astro => "astro", Self::Bash => "bash", Self::C => "c", + Self::Cmake => "cmake", Self::Cpp => "cpp", Self::CSharp => "csharp", Self::Clojure => "clojure", Self::Css => "css", Self::Diff => "diff", + Self::Dockerfile => "dockerfile", Self::Elixir => "elixir", + Self::Erlang => "erlang", Self::Go => "go", + Self::Graphql => "graphql", Self::Handlebars => "handlebars", Self::Haskell => "haskell", Self::Hcl => "hcl", Self::Html => "html", + Self::Ini => "ini", Self::Java => "java", Self::JavaScript => "javascript", Self::Json => "json", + Self::Just => "just", Self::Julia => "julia", Self::Kotlin => "kotlin", Self::Lua => "lua", @@ -339,21 +377,29 @@ impl SupportLang { Self::Markdown => "markdown", Self::Nix => "nix", Self::ObjC => "objc", + Self::Ocaml => "ocaml", Self::Odin => "odin", + Self::Perl => "perl", Self::Php => "php", + Self::Powershell => "powershell", + Self::Proto => "protobuf", Self::Python => "python", + Self::R => "r", Self::Regex => "regex", Self::Ruby => "ruby", Self::Rust => "rust", Self::Scala => "scala", Self::Solidity => "solidity", + Self::Sql => "sql", Self::Starlark => "starlark", + Self::Svelte => "svelte", Self::Swift => "swift", Self::Toml => "toml", Self::Tlaplus => "tlaplus", Self::Tsx => "tsx", Self::TypeScript => "typescript", Self::Verilog => "verilog", + Self::Vue => "vue", Self::Xml => "xml", Self::Yaml => "yaml", Self::Zig => "zig", @@ -379,25 +425,32 @@ impl fmt::Display for SupportLang { // ── Dispatch macro ────────────────────────────────────────────────────── macro_rules! execute_lang_method { - ($me:path, $method:ident, $($pname:tt),*) => { + ($me:expr, $method:ident, $($pname:tt),*) => { use SupportLang as S; - match $me { + match *$me { + S::Astro => Astro.$method($($pname,)*), S::Bash => Bash.$method($($pname,)*), S::C => C.$method($($pname,)*), + S::Cmake => Cmake.$method($($pname,)*), S::Cpp => Cpp.$method($($pname,)*), S::CSharp => CSharp.$method($($pname,)*), S::Clojure => Clojure.$method($($pname,)*), S::Css => Css.$method($($pname,)*), S::Diff => Diff.$method($($pname,)*), + S::Dockerfile => Dockerfile.$method($($pname,)*), S::Elixir => Elixir.$method($($pname,)*), + S::Erlang => Erlang.$method($($pname,)*), S::Go => Go.$method($($pname,)*), + S::Graphql => Graphql.$method($($pname,)*), S::Handlebars => Handlebars.$method($($pname,)*), S::Haskell => Haskell.$method($($pname,)*), S::Hcl => Hcl.$method($($pname,)*), S::Html => Html.$method($($pname,)*), + S::Ini => Ini.$method($($pname,)*), S::Java => Java.$method($($pname,)*), S::JavaScript => JavaScript.$method($($pname,)*), S::Json => Json.$method($($pname,)*), + S::Just => Just.$method($($pname,)*), S::Julia => Julia.$method($($pname,)*), S::Kotlin => Kotlin.$method($($pname,)*), S::Lua => Lua.$method($($pname,)*), @@ -405,21 +458,29 @@ macro_rules! execute_lang_method { S::Markdown => Markdown.$method($($pname,)*), S::Nix => Nix.$method($($pname,)*), S::ObjC => ObjC.$method($($pname,)*), + S::Ocaml => Ocaml.$method($($pname,)*), S::Odin => Odin.$method($($pname,)*), + S::Perl => Perl.$method($($pname,)*), S::Php => Php.$method($($pname,)*), + S::Powershell => Powershell.$method($($pname,)*), + S::Proto => Proto.$method($($pname,)*), S::Python => Python.$method($($pname,)*), + S::R => R.$method($($pname,)*), S::Regex => Regex.$method($($pname,)*), S::Ruby => Ruby.$method($($pname,)*), S::Rust => Rust.$method($($pname,)*), S::Scala => Scala.$method($($pname,)*), S::Solidity => Solidity.$method($($pname,)*), + S::Sql => Sql.$method($($pname,)*), S::Starlark => Starlark.$method($($pname,)*), + S::Svelte => Svelte.$method($($pname,)*), S::Swift => Swift.$method($($pname,)*), S::Toml => Toml.$method($($pname,)*), S::Tlaplus => Tlaplus.$method($($pname,)*), S::Tsx => Tsx.$method($($pname,)*), S::TypeScript => TypeScript.$method($($pname,)*), S::Verilog => Verilog.$method($($pname,)*), + S::Vue => Vue.$method($($pname,)*), S::Xml => Xml.$method($($pname,)*), S::Yaml => Yaml.$method($($pname,)*), S::Zig => Zig.$method($($pname,)*), @@ -479,24 +540,31 @@ impl LanguageExt for SupportLang { const fn extensions(lang: SupportLang) -> &'static [&'static str] { use SupportLang::*; match lang { + Astro => &["astro"], Bash => { &["bash", "bats", "cgi", "command", "env", "fcgi", "ksh", "sh", "tmux", "tool", "zsh"] }, C => &["c", "h"], + Cmake => &["cmake"], Cpp => &["cc", "hpp", "cpp", "c++", "hh", "cxx", "cu", "ino"], CSharp => &["cs"], Clojure => &["clj", "cljs", "cljc", "edn"], Css => &["css", "scss"], Diff => &["diff", "patch"], + Dockerfile => &["dockerfile"], Elixir => &["ex", "exs"], + Erlang => &["erl", "hrl"], Go => &["go"], + Graphql => &["graphql", "gql"], Handlebars => &["hbs", "hsb", "handlebars"], Haskell => &["hs"], Hcl => &["hcl", "tf", "tfvars"], Html => &["html", "htm", "xhtml"], + Ini => &["ini", "cfg", "conf", "properties"], Java => &["java"], JavaScript => &["cjs", "js", "mjs", "jsx"], Json => &["json"], + Just => &[], Julia => &["jl"], Kotlin => &["kt", "ktm", "kts"], Lua => &["lua"], @@ -504,21 +572,29 @@ const fn extensions(lang: SupportLang) -> &'static [&'static str] { Markdown => &["md", "markdown", "mdx"], Nix => &["nix"], ObjC => &["m"], + Ocaml => &["ml"], Odin => &["odin"], + Perl => &["pl", "pm"], Php => &["php"], + Powershell => &["ps1", "psm1"], + Proto => &["proto"], Python => &["py", "py3", "pyi", "bzl"], - Regex => &[], // regex has no file extension + R => &["r"], + Regex => &[], Ruby => &["rb", "rbw", "gemspec"], Rust => &["rs"], Scala => &["scala", "sc", "sbt"], Solidity => &["sol"], + Sql => &["sql"], Starlark => &["star", "bzl"], + Svelte => &["svelte"], Swift => &["swift"], Toml => &["toml"], Tlaplus => &["tla"], Tsx => &["tsx"], TypeScript => &["ts", "cts", "mts"], Verilog => &["v", "sv", "svh", "vh"], + Vue => &["vue"], Xml => &["xml", "xsl", "xslt", "svg", "plist"], Yaml => &["yaml", "yml"], Zig => &["zig"], @@ -527,15 +603,27 @@ const fn extensions(lang: SupportLang) -> &'static [&'static str] { /// Guess language from file extension. fn from_extension(path: &Path) -> Option { - let ext = path.extension()?.to_str()?; - // Special cases: Makefile has no extension - if ext.is_empty() { - let name = path.file_name()?.to_str()?; - return match name { - "Makefile" | "makefile" | "GNUmakefile" => Some(SupportLang::Make), - _ => None, - }; + let name = path.file_name()?.to_str()?; + if name == "Makefile" || name == "makefile" || name == "GNUmakefile" { + return Some(SupportLang::Make); } + if name == "Justfile" || name == "justfile" { + return Some(SupportLang::Just); + } + if name == "CMakeLists.txt" { + return Some(SupportLang::Cmake); + } + if name == "Dockerfile" + || name == "dockerfile" + || name.starts_with("Dockerfile.") + || name.starts_with("dockerfile.") + || name == "Containerfile" + || name == "containerfile" + { + return Some(SupportLang::Dockerfile); + } + + let ext = path.extension()?.to_str()?; SupportLang::all_langs() .iter() .copied() @@ -543,6 +631,7 @@ fn from_extension(path: &Path) -> Option { } static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { +"astro" => SupportLang::Astro, "bash" => SupportLang::Bash, "sh" => SupportLang::Bash, "zsh" => SupportLang::Bash, @@ -550,6 +639,7 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "bats" => SupportLang::Bash, "c" => SupportLang::C, "h" => SupportLang::C, +"cmake" => SupportLang::Cmake, "cpp" => SupportLang::Cpp, "c++" => SupportLang::Cpp, "cc" => SupportLang::Cpp, @@ -570,11 +660,19 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "edn" => SupportLang::Clojure, "diff" => SupportLang::Diff, "patch" => SupportLang::Diff, +"docker" => SupportLang::Dockerfile, +"dockerfile" => SupportLang::Dockerfile, +"containerfile" => SupportLang::Dockerfile, "elixir" => SupportLang::Elixir, "ex" => SupportLang::Elixir, "exs" => SupportLang::Elixir, +"erlang" => SupportLang::Erlang, +"erl" => SupportLang::Erlang, +"hrl" => SupportLang::Erlang, "go" => SupportLang::Go, "golang" => SupportLang::Go, +"graphql" => SupportLang::Graphql, +"gql" => SupportLang::Graphql, "handlebars" => SupportLang::Handlebars, "hbs" => SupportLang::Handlebars, "hsb" => SupportLang::Handlebars, @@ -588,6 +686,11 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "html" => SupportLang::Html, "htm" => SupportLang::Html, "xhtml" => SupportLang::Html, +"ini" => SupportLang::Ini, +"cfg" => SupportLang::Ini, +"conf" => SupportLang::Ini, +"config" => SupportLang::Ini, +"properties" => SupportLang::Ini, "java" => SupportLang::Java, "javascript" => SupportLang::JavaScript, "js" => SupportLang::JavaScript, @@ -595,6 +698,8 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "mjs" => SupportLang::JavaScript, "cjs" => SupportLang::JavaScript, "json" => SupportLang::Json, +"just" => SupportLang::Just, +"justfile" => SupportLang::Just, "julia" => SupportLang::Julia, "jl" => SupportLang::Julia, "kotlin" => SupportLang::Kotlin, @@ -616,12 +721,23 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "objective-c" => SupportLang::ObjC, "m" => SupportLang::ObjC, "mm" => SupportLang::ObjC, +"ocaml" => SupportLang::Ocaml, +"ml" => SupportLang::Ocaml, "odin" => SupportLang::Odin, +"perl" => SupportLang::Perl, +"pl" => SupportLang::Perl, +"pm" => SupportLang::Perl, "php" => SupportLang::Php, +"powershell" => SupportLang::Powershell, +"ps1" => SupportLang::Powershell, +"psm1" => SupportLang::Powershell, +"protobuf" => SupportLang::Proto, +"proto" => SupportLang::Proto, "python" => SupportLang::Python, "py" => SupportLang::Python, "py3" => SupportLang::Python, "pyi" => SupportLang::Python, +"r" => SupportLang::R, "regex" => SupportLang::Regex, "re" => SupportLang::Regex, "ruby" => SupportLang::Ruby, @@ -635,11 +751,13 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "sbt" => SupportLang::Scala, "solidity" => SupportLang::Solidity, "sol" => SupportLang::Solidity, +"sql" => SupportLang::Sql, "starlark" => SupportLang::Starlark, "star" => SupportLang::Starlark, "bzl" => SupportLang::Starlark, "bazel" => SupportLang::Starlark, "skylark" => SupportLang::Starlark, +"svelte" => SupportLang::Svelte, "swift" => SupportLang::Swift, "toml" => SupportLang::Toml, "tla" => SupportLang::Tlaplus, @@ -658,6 +776,7 @@ static LANG_ALIASES: phf::Map<&'static str, SupportLang> = phf_map! { "svh" => SupportLang::Verilog, "vh" => SupportLang::Verilog, "v" => SupportLang::Verilog, +"vue" => SupportLang::Vue, "xml" => SupportLang::Xml, "xsl" => SupportLang::Xml, "xslt" => SupportLang::Xml, diff --git a/crates/pi-natives/src/language/parsers.rs b/crates/pi-natives/src/language/parsers.rs index d31e823f7..750127d2b 100644 --- a/crates/pi-natives/src/language/parsers.rs +++ b/crates/pi-natives/src/language/parsers.rs @@ -2,6 +2,9 @@ use ast_grep_core::tree_sitter::TSLanguage; +pub fn language_astro() -> TSLanguage { + tree_sitter_astro::LANGUAGE.into() +} pub fn language_bash() -> TSLanguage { tree_sitter_bash::LANGUAGE.into() } @@ -11,6 +14,9 @@ pub fn language_c() -> TSLanguage { pub fn language_clojure() -> TSLanguage { tree_sitter_clojure::LANGUAGE.into() } +pub fn language_cmake() -> TSLanguage { + tree_sitter_cmake::LANGUAGE.into() +} pub fn language_cpp() -> TSLanguage { tree_sitter_cpp::LANGUAGE.into() } @@ -23,12 +29,21 @@ pub fn language_css() -> TSLanguage { pub fn language_diff() -> TSLanguage { tree_sitter_diff::LANGUAGE.into() } +pub fn language_dockerfile() -> TSLanguage { + tree_sitter_dockerfile::language() +} pub fn language_elixir() -> TSLanguage { tree_sitter_elixir::LANGUAGE.into() } +pub fn language_erlang() -> TSLanguage { + tree_sitter_erlang::LANGUAGE.into() +} pub fn language_go() -> TSLanguage { tree_sitter_go::LANGUAGE.into() } +pub fn language_graphql() -> TSLanguage { + tree_sitter_graphql::LANGUAGE.into() +} pub fn language_handlebars() -> TSLanguage { tree_sitter_glimmer::language() } @@ -41,6 +56,9 @@ pub fn language_hcl() -> TSLanguage { pub fn language_html() -> TSLanguage { tree_sitter_html::LANGUAGE.into() } +pub fn language_ini() -> TSLanguage { + tree_sitter_ini::LANGUAGE.into() +} pub fn language_java() -> TSLanguage { tree_sitter_java::LANGUAGE.into() } @@ -50,6 +68,9 @@ pub fn language_javascript() -> TSLanguage { pub fn language_json() -> TSLanguage { tree_sitter_json::LANGUAGE.into() } +pub fn language_just() -> TSLanguage { + tree_sitter_just::LANGUAGE.into() +} pub fn language_julia() -> TSLanguage { tree_sitter_julia::LANGUAGE.into() } @@ -71,15 +92,30 @@ pub fn language_nix() -> TSLanguage { pub fn language_objc() -> TSLanguage { tree_sitter_objc::LANGUAGE.into() } +pub fn language_ocaml() -> TSLanguage { + tree_sitter_ocaml::LANGUAGE_OCAML.into() +} pub fn language_odin() -> TSLanguage { tree_sitter_odin::LANGUAGE.into() } +pub fn language_perl() -> TSLanguage { + tree_sitter_perl::LANGUAGE.into() +} pub fn language_php() -> TSLanguage { tree_sitter_php::LANGUAGE_PHP_ONLY.into() } +pub fn language_powershell() -> TSLanguage { + tree_sitter_powershell::LANGUAGE.into() +} +pub fn language_proto() -> TSLanguage { + tree_sitter_proto::LANGUAGE.into() +} pub fn language_python() -> TSLanguage { tree_sitter_python::LANGUAGE.into() } +pub fn language_r() -> TSLanguage { + tree_sitter_r::LANGUAGE.into() +} pub fn language_regex() -> TSLanguage { tree_sitter_regex::LANGUAGE.into() } @@ -95,9 +131,15 @@ pub fn language_scala() -> TSLanguage { pub fn language_solidity() -> TSLanguage { tree_sitter_solidity::LANGUAGE.into() } +pub fn language_sql() -> TSLanguage { + tree_sitter_sql::LANGUAGE.into() +} pub fn language_starlark() -> TSLanguage { tree_sitter_starlark::LANGUAGE.into() } +pub fn language_svelte() -> TSLanguage { + tree_sitter_svelte::LANGUAGE.into() +} pub fn language_swift() -> TSLanguage { tree_sitter_swift::LANGUAGE.into() } @@ -116,6 +158,9 @@ pub fn language_tlaplus() -> TSLanguage { pub fn language_verilog() -> TSLanguage { tree_sitter_verilog::LANGUAGE.into() } +pub fn language_vue() -> TSLanguage { + tree_sitter_vue::LANGUAGE.into() +} pub fn language_xml() -> TSLanguage { tree_sitter_xml::LANGUAGE_XML.into() }