#[cfg(test)]
mod gap_survey;
pub mod hash;
pub mod language;
pub mod metadata;
pub mod similarity;
mod similarity_corpus_tests;
pub mod tip;
use anyhow::{Result, anyhow};
use std::fmt;
#[derive(Debug, Default)]
pub struct Code {
pub contents: String,
pub metadata: Metadata,
pub ast: Option<tree_sitter::Tree>,
}
impl Clone for Code {
fn clone(&self) -> Self {
Code {
contents: self.contents.clone(),
metadata: Metadata {
ast_metadata: None,
..self.metadata.clone()
},
ast: self.ast.clone(),
}
}
}
impl Code {
pub fn parse(&mut self, parser: &mut tree_sitter::Parser) {
let language = match self.metadata.language.as_ref() {
Some(lang) => lang,
None => return,
};
let ts_language = match crate::code::language::to_treesitter(language) {
Some(ts_lang) => ts_lang,
None => return,
};
if parser.set_language(&ts_language).is_err() {
return;
}
self.ast = parser.parse(&self.contents, None);
}
pub fn ensure_parsed(&mut self) -> Result<()> {
let language = match self.metadata.language.as_ref() {
Some(lang) => lang,
None => return Err(anyhow!("Language must be set to parse code")),
};
let needs_parsing = self.ast.is_none();
let needs_metadata = self.metadata.ast_metadata.is_none();
if !needs_parsing && !needs_metadata {
return Ok(());
}
if needs_parsing {
let ts_language = match crate::code::language::to_treesitter(language) {
Some(ts_lang) => ts_lang,
None => {
return Err(anyhow!(
"Language {} is not supported by tree-sitter",
language
));
}
};
let mut parser = tree_sitter::Parser::new();
if parser.set_language(&ts_language).is_err() {
return Err(anyhow!("Failed to set tree-sitter language"));
}
self.ast = parser.parse(&self.contents, None);
}
if needs_metadata && self.ast.is_some() {
self.metadata.ast_metadata = Some(crate::code::metadata::compute_ast_metadata(self)?);
}
Ok(())
}
pub fn from_string(contents: &str, language: &Language) -> Self {
let mut code = Code {
contents: contents.to_string(),
metadata: Metadata {
path: None,
tip: Some(Type::Code("Code".to_string())),
language: Some(*language),
ast_metadata: None,
},
..Default::default()
};
let mut parser = tree_sitter::Parser::new();
code.parse(&mut parser);
if code.ast.is_some() {
match crate::code::metadata::compute_ast_metadata(&code) {
Ok(ast_metadata) => code.metadata.ast_metadata = Some(ast_metadata),
Err(e) => eprintln!("Failed to compute AST metadata: {:?}", e),
}
}
code
}
pub fn from_file(path: &std::path::Path) -> Result<Self> {
use std::fs;
let contents = fs::read_to_string(path)
.map_err(|e| anyhow!("Failed to read file {}: {}", path.display(), e))?;
let language =
language::language_for_path_and_content(path, &contents).unwrap_or(Language::Unknown);
let mut code = Code::from_string(&contents, &language);
code.metadata.path = Some(path.to_path_buf());
Ok(code)
}
}
pub fn is_binary_file(path: &std::path::Path) -> Result<bool> {
let bytes = std::fs::read(path)
.map_err(|e| anyhow!("Failed to read file {}: {}", path.display(), e))?;
Ok(std::str::from_utf8(&bytes).is_err())
}
#[derive(Debug, Clone, Default)]
pub struct Metadata {
pub path: Option<std::path::PathBuf>,
pub tip: Option<Type>,
pub language: Option<Language>,
pub ast_metadata: Option<ASTMetadata>,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ASTNodeMetadata {
pub kind: String,
pub text: String,
pub owned_text_hash: u64,
pub children: Vec<usize>,
pub start_byte: usize,
pub preorder_index: usize,
pub is_named: bool,
pub kind_cost_class: KindCostClass,
}
impl ASTNodeMetadata {
pub fn new(
kind: String,
text: String,
children: Vec<usize>,
start_byte: usize,
preorder_index: usize,
) -> Self {
let kind_cost_class = KindCostClass {
identifier_like: crate::diff::nodes::is_identifier_kind(&kind),
literal_like: crate::diff::nodes::is_literal_kind(&kind),
operator_families: crate::diff::nodes::operator_family_mask(&kind),
};
ASTNodeMetadata {
kind,
text,
owned_text_hash: 0,
children,
start_byte,
preorder_index,
is_named: true,
kind_cost_class,
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct KindCostClass {
pub identifier_like: bool,
pub literal_like: bool,
pub operator_families: crate::diff::nodes::FamilyMask,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ASTMetadata {
pub node_to_full_hash: rustc_hash::FxHashMap<usize, u64>,
pub full_hash_to_node: rustc_hash::FxHashMap<u64, Vec<usize>>,
pub node_to_structural_hash: rustc_hash::FxHashMap<usize, u64>,
pub structural_hash_to_node: rustc_hash::FxHashMap<u64, Vec<usize>>,
pub node_to_kind_and_value_hash: rustc_hash::FxHashMap<usize, u64>,
pub kind_and_value_hash_to_node: rustc_hash::FxHashMap<u64, Vec<usize>>,
pub node_to_kind_only_hash: rustc_hash::FxHashMap<usize, u64>,
pub kind_only_hash_to_node: rustc_hash::FxHashMap<u64, Vec<usize>>,
pub node_to_similarity_sketch:
rustc_hash::FxHashMap<usize, crate::code::similarity::SimilaritySketch>,
pub node_to_subtree_size: rustc_hash::FxHashMap<usize, usize>,
pub node_to_widest_subtree_node: rustc_hash::FxHashMap<usize, (usize, usize)>,
pub node_to_depth: rustc_hash::FxHashMap<usize, usize>,
pub node_to_parent: rustc_hash::FxHashMap<usize, usize>,
pub reference_nodes_ordered: Vec<usize>,
pub node_info: rustc_hash::FxHashMap<usize, ASTNodeMetadata>,
pub language: Language,
}
impl ASTMetadata {
pub fn is_leaf(&self, id: usize) -> bool {
self.node_info
.get(&id)
.is_some_and(|info| info.children.is_empty())
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, strum::EnumIter)]
pub enum Language {
#[default]
Unknown,
Bazel,
C,
CPP,
CSS,
CSharp,
Dart,
Go,
HTML,
JSON,
Java,
JavaScript,
Kotlin,
LUA,
Lisp,
MarkDown,
PHP,
ProtoBuf,
Python,
R,
Ruby,
Rust,
SQL,
Scala,
ShellScript,
Swift,
TSX,
TypeScript,
Vimscript,
XML,
YAML,
}
impl std::fmt::Display for Language {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
std::fmt::Debug::fmt(self, f)
}
}
#[derive(Debug, Clone, Default, PartialEq)]
pub enum Type {
#[default]
Unknown,
Code(String),
Configuration(String),
Data(String),
Documentation(String),
}
impl std::fmt::Display for Type {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
std::fmt::Debug::fmt(self, f)
}
}
#[cfg(test)]
mod tests {
#[test]
fn is_binary_file_agrees_with_from_file_on_valid_utf8() {
let mut file = tempfile::NamedTempFile::new().expect("create temp file");
std::io::Write::write_all(&mut file, "fn main() {}\n".as_bytes()).expect("write");
assert!(!is_binary_file(file.path()).expect("classify"));
assert!(Code::from_file(file.path()).is_ok());
}
#[test]
fn is_binary_file_agrees_with_from_file_on_invalid_utf8() {
let mut file = tempfile::NamedTempFile::new().expect("create temp file");
std::io::Write::write_all(&mut file, b"%PDF-1.7\n\xff\xfe\x00binary").expect("write");
assert!(is_binary_file(file.path()).expect("classify"));
assert!(Code::from_file(file.path()).is_err());
}
#[cfg(unix)]
#[test]
fn is_binary_file_says_dev_null_is_not_binary() {
assert!(!is_binary_file(std::path::Path::new("/dev/null")).expect("classify"));
}
#[test]
fn is_binary_file_says_valid_utf8_containing_a_nul_byte_is_not_binary() {
let mut file = tempfile::NamedTempFile::new().expect("create temp file");
std::io::Write::write_all(&mut file, b"let s = \"a\x00b\";\n").expect("write");
assert!(!is_binary_file(file.path()).expect("classify"));
assert!(Code::from_file(file.path()).is_ok());
}
use super::*;
use crate::test::helper;
#[test]
fn code_from_empty_string() {
let code = Code::from_string("", &Language::Rust);
assert_eq!(code.contents, "");
assert_eq!(code.metadata.language, Some(Language::Rust));
assert!(code.ast.is_some());
}
#[test]
fn code_from_string_skips_ast_metadata_when_the_language_has_no_grammar() {
let code = Code::from_string("", &Language::Unknown);
assert!(code.ast.is_none());
assert!(code.metadata.ast_metadata.is_none());
}
#[test]
fn cloning_code_drops_ast_metadata_so_its_ids_cannot_outlive_the_tree() -> Result<()> {
let code = Code::from_string("fn main() { let x = 1; }", &Language::Rust);
assert!(code.metadata.ast_metadata.is_some());
let mut clone = code.clone();
assert!(clone.ast.is_some());
assert!(clone.metadata.ast_metadata.is_none());
clone.ensure_parsed()?;
let root = clone.ast.as_ref().unwrap().root_node().id();
let metadata = clone.metadata.ast_metadata.as_ref().unwrap();
assert!(metadata.node_info.contains_key(&root));
Ok(())
}
#[test]
fn code_from_file() -> Result<()> {
let paths = helper::handmade_test_code_as_paths()?;
let hello_world_path = paths
.get("hello-world.rs")
.expect("hello-world.rs should exist in test data");
let code = Code::from_file(hello_world_path)?;
assert_eq!(code.metadata.language, Some(Language::Rust));
assert!(code.metadata.path.is_some());
assert!(code.contents.contains("fn main()"));
assert!(code.contents.contains("Hello, World"));
assert!(code.ast.is_some());
Ok(())
}
#[test]
fn parse_code() -> Result<()> {
let mut codes = helper::handmade_test_code()?;
let hello_world = codes
.get_mut("hello-world.rs")
.expect("hello-world.rs should exist in test data");
let language = hello_world
.metadata
.language
.as_ref()
.expect("Language should be set");
let ts_language = crate::code::language::to_treesitter(language).expect("Unable to convert OmniDiff language to TreeSitter language in tests. Something is wrong with the test infrastructure.");
let mut parser = tree_sitter::Parser::new();
parser.set_language(&ts_language)?;
hello_world.parse(&mut parser);
assert!(hello_world.ast.is_some());
Ok(())
}
#[test]
fn ast_metadata_computed_in_from_string() {
let code = Code::from_string("fn main() { println!(\"Hello, World\"); }", &Language::Rust);
assert!(code.metadata.ast_metadata.is_some());
let ast_metadata = code.metadata.ast_metadata.as_ref().unwrap();
assert!(!ast_metadata.node_to_full_hash.is_empty());
assert!(!ast_metadata.full_hash_to_node.is_empty());
assert!(!ast_metadata.node_to_structural_hash.is_empty());
assert!(!ast_metadata.structural_hash_to_node.is_empty());
}
#[test]
fn ast_metadata_computed_in_from_file() -> Result<()> {
let paths = helper::handmade_test_code_as_paths()?;
let hello_world_path = paths
.get("hello-world.rs")
.expect("hello-world.rs should exist in test data");
let code = Code::from_file(hello_world_path)?;
assert!(code.metadata.ast_metadata.is_some());
let ast_metadata = code.metadata.ast_metadata.as_ref().unwrap();
assert!(!ast_metadata.node_to_full_hash.is_empty());
assert!(!ast_metadata.full_hash_to_node.is_empty());
assert!(!ast_metadata.node_to_structural_hash.is_empty());
assert!(!ast_metadata.structural_hash_to_node.is_empty());
assert!(!ast_metadata.reference_nodes_ordered.is_empty());
Ok(())
}
#[test]
fn ast_metadata_consistency() -> Result<()> {
let paths = helper::handmade_test_code_as_paths()?;
let hello_world_path = paths
.get("hello-world.rs")
.expect("hello-world.rs should exist in test data");
let content = std::fs::read_to_string(hello_world_path)?;
let code_from_file = Code::from_file(hello_world_path)?;
let code_from_string = Code::from_string(&content, &Language::Rust);
assert!(code_from_file.metadata.ast_metadata.is_some());
assert!(code_from_string.metadata.ast_metadata.is_some());
let metadata_from_file = code_from_file.metadata.ast_metadata.as_ref().unwrap();
let metadata_from_string = code_from_string.metadata.ast_metadata.as_ref().unwrap();
assert_eq!(
metadata_from_file.node_to_full_hash.len(),
metadata_from_string.node_to_full_hash.len()
);
assert_eq!(
metadata_from_file.full_hash_to_node.len(),
metadata_from_string.full_hash_to_node.len()
);
assert_eq!(
metadata_from_file.node_to_structural_hash.len(),
metadata_from_string.node_to_structural_hash.len()
);
assert_eq!(
metadata_from_file.structural_hash_to_node.len(),
metadata_from_string.structural_hash_to_node.len()
);
Ok(())
}
#[test]
fn ensure_parsed_already_parsed_and_metadata_set() -> Result<()> {
let mut code = Code::from_string("fn main() { println!(\"Hello\"); }", &Language::Rust);
assert!(code.ast.is_some());
assert!(code.metadata.ast_metadata.is_some());
let original_metadata = code.metadata.ast_metadata.clone();
code.ensure_parsed()?;
assert!(code.ast.is_some());
assert_eq!(code.metadata.ast_metadata, original_metadata);
Ok(())
}
#[test]
fn ensure_parsed_parsed_but_no_metadata() -> Result<()> {
let mut code = Code {
contents: "fn main() { println!(\"Hello\"); }".to_string(),
metadata: Metadata {
path: None,
tip: Some(Type::Code("Code".to_string())),
language: Some(Language::Rust),
ast_metadata: None,
},
..Default::default()
};
let mut parser = tree_sitter::Parser::new();
code.parse(&mut parser);
assert!(code.ast.is_some());
assert!(code.metadata.ast_metadata.is_none());
code.ensure_parsed()?;
assert!(code.ast.is_some());
assert!(code.metadata.ast_metadata.is_some());
let metadata = code.metadata.ast_metadata.as_ref().unwrap();
assert!(!metadata.node_to_full_hash.is_empty());
assert!(!metadata.node_to_structural_hash.is_empty());
Ok(())
}
#[test]
fn ensure_parsed_not_parsed() -> Result<()> {
let mut code = Code {
contents: "fn main() { println!(\"Hello\"); }".to_string(),
metadata: Metadata {
path: None,
tip: Some(Type::Code("Code".to_string())),
language: Some(Language::Rust),
ast_metadata: None,
},
..Default::default()
};
assert!(code.ast.is_none());
assert!(code.metadata.ast_metadata.is_none());
code.ensure_parsed()?;
assert!(code.ast.is_some());
assert!(code.metadata.ast_metadata.is_some());
let metadata = code.metadata.ast_metadata.as_ref().unwrap();
assert!(!metadata.node_to_full_hash.is_empty());
assert!(!metadata.node_to_structural_hash.is_empty());
Ok(())
}
#[test]
fn ensure_parsed_no_language() {
let mut code = Code {
contents: "fn main() { println!(\"Hello\"); }".to_string(),
metadata: Metadata {
path: None,
tip: Some(Type::Code("Code".to_string())),
language: None,
ast_metadata: None,
},
..Default::default()
};
let result = code.ensure_parsed();
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Language must be set")
);
assert!(code.ast.is_none());
assert!(code.metadata.ast_metadata.is_none());
}
#[test]
fn ensure_parsed_unsupported_language() {
let mut code = Code {
contents: "fn main() { println!(\"Hello\"); }".to_string(),
metadata: Metadata {
path: None,
tip: Some(Type::Code("Code".to_string())),
language: Some(Language::Bazel), ast_metadata: None,
},
..Default::default()
};
let result = code.ensure_parsed();
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("not supported by tree-sitter")
);
assert!(code.ast.is_none());
assert!(code.metadata.ast_metadata.is_none());
}
}