Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
2026-04-28 08:43:33 -05:00
commit 0dbd3d567a
4982 changed files with 1431549 additions and 0 deletions
@@ -0,0 +1,108 @@
+use crate::index::full_source_code_embedding::chunker::{coalesce_fragments, Fragment};
+use itertools::Itertools;
+use line_span::{LineSpan, LineSpans};
+use std::path::Path;
+
+/// Chunks the given file into [`Fragment`]s. Each chunk is at most `num_lines_per_chunk` lines long, and contains at most `max_bytes_per_chunk` bytes.
+pub(super) fn chunk_code<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    num_lines_per_chunk: usize,
+) -> Vec<Fragment<'a>> {
+    let lines = code.line_spans().enumerate().collect_vec();
+    let chunks = lines.chunks(num_lines_per_chunk);
+
+    chunks
+        .into_iter()
+        .flat_map(|chunk| {
+            let (start_line, start_range) = chunk[0];
+            let (end_line, end_range) =
+                chunk.last().expect("Chunks must have at least one element");
+
+            if (end_range.end() - start_range.start()) > max_bytes_per_chunk {
+                let chunked_fragments = chunk.iter().flat_map(|(line, line_span)| {
+                    chunk_line_by_bytes(code, path, max_bytes_per_chunk, *line, line_span)
+                });
+
+                return coalesce_fragments(chunked_fragments, code, max_bytes_per_chunk);
+            }
+
+            vec![Fragment {
+                content: &code[start_range.start()..end_range.end()],
+                start_line,
+                end_line: *end_line,
+                file_path: path,
+                start_byte_index: start_range.start().into(),
+                end_byte_index: end_range.end().into(),
+            }]
+        })
+        .collect()
+}
+
+/// Chunks the line represented by `line_span` into multiple fragments if it exceeds `max_bytes_per_chunk`.
+fn chunk_line_by_bytes<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    line_number: usize,
+    line_span: &LineSpan<'a>,
+) -> Vec<Fragment<'a>> {
+    let line_start = line_span.start();
+    let line_end = line_span.end();
+    let line_content = &code[line_start..line_end];
+    let line_length = line_end - line_start;
+
+    // If the line is smaller than max_bytes_per_chunk, return it as a single fragment
+    if line_length <= max_bytes_per_chunk {
+        return vec![Fragment {
+            content: line_content,
+            start_line: line_number,
+            end_line: line_number,
+            file_path: path,
+            start_byte_index: line_start.into(),
+            end_byte_index: line_end.into(),
+        }];
+    }
+
+    // Otherwise, split the line into multiple fragments
+    let mut fragments = Vec::new();
+    let mut current_start = line_start;
+
+    while current_start < line_end {
+        let remaining_bytes = line_end - current_start;
+        let chunk_size = std::cmp::min(remaining_bytes, max_bytes_per_chunk);
+        let mut chunk_end = current_start + chunk_size;
+
+        // Ensure chunk_end is on a UTF-8 character boundary
+        while chunk_end > current_start && !code.is_char_boundary(chunk_end) {
+            chunk_end -= 1;
+        }
+
+        // If we couldn't find a valid boundary within reasonable distance,
+        // move forward to the next character boundary instead
+        if chunk_end <= current_start {
+            chunk_end = current_start + chunk_size;
+            while chunk_end < line_end && !code.is_char_boundary(chunk_end) {
+                chunk_end += 1;
+            }
+        }
+
+        fragments.push(Fragment {
+            content: &code[current_start..chunk_end],
+            start_line: line_number,
+            end_line: line_number,
+            file_path: path,
+            start_byte_index: current_start.into(),
+            end_byte_index: chunk_end.into(),
+        });
+
+        current_start = chunk_end;
+    }
+
+    fragments
+}
+
+#[cfg(test)]
+#[path = "naive_tests.rs"]
+mod tests;
@@ -0,0 +1,354 @@
+use super::*;
+use std::path::Path;
+
+#[test]
+fn test_chunker() {
+    let code = "This is some text content\nthat should be chunked\nusing the naive chunker\nbecause the language isn't recognized.";
+    let path = Path::new("test_file.xyz");
+
+    let max_lines = 1;
+    let fragments = chunk_code(code, path, 10000, max_lines);
+
+    assert!(!fragments.is_empty(), "Expected at least one fragment");
+
+    assert_eq!(fragments.len(), code.lines().count());
+    for (idx, line) in code.lines().enumerate() {
+        assert_eq!(fragments[idx].content, line);
+        assert_eq!(fragments[idx].start_line, idx);
+        assert_eq!(fragments[idx].end_line, idx);
+    }
+}
+
+#[test]
+fn test_chunker_large_chunk() {
+    let code = "This is some text content\nthat should be chunked\nusing the naive chunker\nbecause the language isn't recognized.";
+    let path = Path::new("test_file.xyz");
+
+    let fragments = chunk_code(code, path, 10000, 100);
+
+    // We should have only one fragment
+    assert_eq!(fragments.len(), 1);
+
+    assert_eq!(fragments[0].content, code);
+    assert_eq!(fragments[0].start_line, 0);
+    assert_eq!(fragments[0].end_line, code.lines().count() - 1);
+}
+
+#[test]
+fn test_chunker_max_bytes() {
+    // Create a string with known byte size - each line is exactly 20 bytes including newline
+    let code = "line1\nline2\nline3\nline4abcdefghijklmnopqrstuvwxyz";
+    let path = Path::new("test_file.xyz");
+
+    // Set max_bytes_per_chunk to 25 bytes to force multiple chunks for the last line (which is 30 bytes).
+    let max_bytes_per_chunk = 25;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that no chunk exceeds the max_bytes_per_chunk limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.trim().len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk
+        );
+    }
+
+    // The first fragment contains all of the lines except the last one.
+    assert_eq!(fragments[0].content, "line1\nline2\nline3");
+
+    // The last two fragments contains the contents of the line line.
+    assert_eq!(fragments[1].content, "line4abcdefghijklmnopqrst");
+    assert_eq!(fragments[2].content, "uvwxyz");
+
+    // Verify that the chunks together contain all the original content
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    // Ignore any newlines when doing comparisons--the chunker may drop newlines at fragment boundaries
+    // and that's not necessary for testing the correctness of the naive chunker.
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_emoji_chunking() {
+    // Test with emojis (4-byte UTF-8 characters) to ensure byte boundaries are respected
+    let code = "Hello 🦀 Rust\nWorld 🌍 Test\n🚀 Rocket 🎯 Target";
+    let path = Path::new("test_emoji.txt");
+
+    // Set a small max_bytes_per_chunk to force splitting through emoji characters
+    let max_bytes_per_chunk = 15; // This will force splits in the middle of emoji sequences
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that no chunk exceeds the max_bytes_per_chunk limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes. Content: '{}'",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk,
+            fragment.content
+        );
+    }
+
+    // Verify that all fragments contain valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.is_ascii() || std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content
+        );
+    }
+
+    // Verify that reassembled content matches original (ignoring newlines)
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_accented_characters() {
+    // Test with accented characters (2-byte UTF-8)
+    let code = "Café résumé naïve\nÉlève découvrir\nMañana piñata";
+    let path = Path::new("test_accents.txt");
+
+    // Set max_bytes_per_chunk to force splitting through accented characters
+    let max_bytes_per_chunk = 10;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that all fragments contain valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify that reassembled content matches original (ignoring newlines)
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content.replace('\n', ""),
+        code.replace('\n', ""),
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_mixed_characters() {
+    // Test with a mix of ASCII, 2-byte, 3-byte, and 4-byte UTF-8 characters
+    let code = "ASCII text 中文 🦀 résumé ℘ math symbols";
+    let path = Path::new("test_mixed.txt");
+
+    // Set a small chunk size to force many splits
+    let max_bytes_per_chunk = 8;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we have multiple chunks
+    assert!(
+        fragments.len() > 1,
+        "Expected multiple chunks due to size limit"
+    );
+
+    // Verify that all fragments contain valid UTF-8 and don't exceed size limit
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            fragment.content.len() <= max_bytes_per_chunk,
+            "Fragment {} has size {} bytes, which exceeds limit of {} bytes",
+            i,
+            fragment.content.len(),
+            max_bytes_per_chunk
+        );
+
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify that reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_boundary_edge_cases() {
+    // Test edge case where chunk boundary falls exactly on a multi-byte character
+    let code = "ab🦀cd"; // 'ab' (2 bytes) + '🦀' (4 bytes) + 'cd' (2 bytes) = 8 bytes total
+    let path = Path::new("test_edge.txt");
+
+    // Set chunk size to 3 bytes, which would split in the middle of the emoji without our fix
+    let max_bytes_per_chunk = 3;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Should have multiple fragments
+    assert!(fragments.len() >= 2, "Expected at least 2 fragments");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}
+
+#[test]
+fn test_utf8_single_multibyte_character() {
+    // Test with a single multi-byte character that's larger than chunk size
+    let code = "🦀"; // 4-byte emoji
+    let path = Path::new("test_single.txt");
+
+    // Set chunk size smaller than the character
+    let max_bytes_per_chunk = 2;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Should have exactly one fragment (can't split a single character)
+    assert_eq!(fragments.len(), 1, "Should have exactly one fragment");
+
+    // The fragment should contain the complete character
+    assert_eq!(fragments[0].content, code);
+
+    // Verify it's valid UTF-8
+    assert!(
+        std::str::from_utf8(fragments[0].content.as_bytes()).is_ok(),
+        "Fragment contains invalid UTF-8"
+    );
+}
+
+#[test]
+fn test_utf8_line_endings_with_multibyte() {
+    // Test multi-byte characters at line boundaries
+    let code = "Hello🌍\nWorld🦀\nTest🎯";
+    let path = Path::new("test_lines.txt");
+
+    let max_bytes_per_chunk = 10;
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1); // 1 line per chunk
+
+    // Should have 3 fragments (one per line)
+    assert_eq!(fragments.len(), 3, "Should have 3 fragments for 3 lines");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify line numbers are correct
+    assert_eq!(fragments[0].start_line, 0);
+    assert_eq!(fragments[0].end_line, 0);
+    assert_eq!(fragments[1].start_line, 1);
+    assert_eq!(fragments[1].end_line, 1);
+    assert_eq!(fragments[2].start_line, 2);
+    assert_eq!(fragments[2].end_line, 2);
+}
+
+#[test]
+fn test_panic_regression_byte_boundary() {
+    // This is a regression test for the "byte index is not a char boundary" panic.
+    // Before the fix, this would panic when trying to slice at byte index 3,
+    // which is in the middle of the 4-byte emoji '🦀'.
+    let code = "Hi🦀Test";
+    let path = Path::new("test_panic.txt");
+
+    // This chunk size would cause the original code to panic
+    let max_bytes_per_chunk = 3;
+
+    // This should not panic
+    let fragments = chunk_code(code, path, max_bytes_per_chunk, 1000);
+
+    // Verify we get valid fragments
+    assert!(!fragments.is_empty(), "Should have at least one fragment");
+
+    // Verify all fragments are valid UTF-8
+    for (i, fragment) in fragments.iter().enumerate() {
+        assert!(
+            std::str::from_utf8(fragment.content.as_bytes()).is_ok(),
+            "Fragment {} contains invalid UTF-8: {:?}",
+            i,
+            fragment.content.as_bytes()
+        );
+    }
+
+    // Verify reassembled content matches original
+    let reassembled_content: String = fragments
+        .iter()
+        .map(|f| f.content)
+        .collect::<Vec<_>>()
+        .join("");
+
+    assert_eq!(
+        reassembled_content, code,
+        "Reassembled content does not match original"
+    );
+}
@@ -0,0 +1,170 @@
+use std::path::Path;
+
+use arborium::tree_sitter::{Language, Node, Parser, TreeCursor};
+use itertools::Itertools;
+
+use super::{coalesce_fragments, Fragment};
+
+/// Maximum depth for recursive tree traversal to prevent infinite recursion
+/// or excessive depth in malformed/deeply nested code.
+const MAX_TRAVERSAL_DEPTH: usize = 200;
+
+/// Chunks code into an ordered list of fragments, where each fragment is at most
+/// `max_bytes_per_chunk` bytes.
+pub(super) fn chunk_code<'a>(
+    code: &'a str,
+    path: &'a Path,
+    max_bytes_per_chunk: usize,
+    language: &Language,
+) -> anyhow::Result<Vec<Fragment<'a>>> {
+    // Wrap this in a block to ensure the treesitter Parser / Tree are dropped
+    // after creating the fragments.
+    let fragments = {
+        let mut parser = Parser::new();
+        parser.set_language(language)?;
+
+        let tree = parser
+            .parse(code, None /* old_tree */)
+            .ok_or_else(|| anyhow::anyhow!("Failed to parse code"))?;
+
+        let mut cursor = tree.walk();
+
+        let nodes = split_node(
+            tree.root_node(),
+            code,
+            max_bytes_per_chunk,
+            path,
+            &mut cursor,
+            0, // initial depth
+        )?;
+
+        coalesce_fragments(nodes.into_iter(), code, max_bytes_per_chunk)
+    };
+
+    // Release extra unused memory from malloc to the system.  For some
+    // reason, the memory obtained by the allocator is often not released
+    // back to the OS after we're done with it, resulting in high memory
+    // usage (from the perspective of the OS, though not from the perspective
+    // of the allocator).
+    //
+    // See: https://github.com/tree-sitter/tree-sitter/issues/3129
+    #[cfg(all(target_os = "linux", target_env = "gnu", not(feature = "jemalloc")))]
+    unsafe {
+        nix::libc::malloc_trim(0);
+    }
+
+    Ok(fragments)
+}
+
+/// Splits a [`Node`] into a series of [`Fragment`]s that are at most `max_bytes_per_chunk` bytes.
+fn split_node<'a, 'b>(
+    node: Node<'b>,
+    code: &'a str,
+    max_bytes_per_chunk: usize,
+    path: &'a Path,
+    cursor: &mut TreeCursor<'b>,
+    depth: usize,
+) -> anyhow::Result<Vec<Fragment<'a>>> {
+    // Check if we've exceeded the maximum traversal depth
+    if depth > MAX_TRAVERSAL_DEPTH {
+        return Err(anyhow::anyhow!(
+            "Maximum traversal depth {} exceeded, falling back to naive chunking",
+            MAX_TRAVERSAL_DEPTH
+        ));
+    }
+
+    let mut current_fragment = Fragment::from_node_start(node, path);
+    let mut fragments = vec![];
+
+    // Collect into a vec to avoid a double mutable borrow with `cursor` when we make
+    // the recursive call below.
+    for child in node.children(cursor).collect_vec() {
+        let child_size = child.end_byte().saturating_sub(child.start_byte());
+
+        // The child is larger than the max chunk size, so we need to split it recursively.
+        if child_size > max_bytes_per_chunk {
+            let mut new_fragment = Fragment::from_node_end(child, path);
+            std::mem::swap(&mut current_fragment, &mut new_fragment);
+            fragments.push(new_fragment);
+
+            fragments.append(&mut split_node(
+                child,
+                code,
+                max_bytes_per_chunk,
+                path,
+                cursor,
+                depth + 1,
+            )?);
+        } else if child_size + current_fragment.size() > max_bytes_per_chunk {
+            // The child would make the current fragment too large, so we finalize the current
+            // fragment and create a new one.
+            fragments.push(current_fragment);
+            current_fragment = Fragment::from_node_start(child, path);
+            current_fragment.append(&Fragment::from_node_end(child, path), code);
+        } else {
+            // The child fits within the current fragment.
+            current_fragment.end_line = child.end_position().row;
+            current_fragment.end_byte_index = child.end_byte().into();
+            current_fragment.content =
+                &code[current_fragment.start_byte_index.as_usize()..child.end_byte()];
+        }
+    }
+
+    fragments.push(current_fragment);
+
+    Ok(fragments)
+}
+
+impl<'a> Fragment<'a> {
+    /// Creates an empty fragment.
+    fn empty() -> Fragment<'a> {
+        Fragment {
+            content: "",
+            start_line: 0,
+            end_line: 0,
+            start_byte_index: 0.into(),
+            end_byte_index: 0.into(),
+            file_path: Path::new(""),
+        }
+    }
+
+    /// Creates a fragment comprised solely of the start of the given node.
+    fn from_node_start(node: Node<'_>, path: &'a Path) -> Self {
+        Fragment {
+            content: "",
+            start_line: node.start_position().row,
+            end_line: node.start_position().row,
+            start_byte_index: node.start_byte().into(),
+            end_byte_index: node.start_byte().into(),
+            file_path: path,
+        }
+    }
+
+    /// Creates a fragment comprised solely of the end of the given node.
+    fn from_node_end(node: Node<'_>, path: &'a Path) -> Self {
+        Fragment {
+            content: "",
+            start_line: node.end_position().row,
+            end_line: node.end_position().row,
+            start_byte_index: node.end_byte().into(),
+            end_byte_index: node.end_byte().into(),
+            file_path: path,
+        }
+    }
+
+    /// Creates a fragment comprised solely of the end of the given fragment.
+    fn from_fragment_end(fragment: &Fragment<'a>) -> Self {
+        Fragment {
+            content: "",
+            start_line: fragment.end_line,
+            end_line: fragment.end_line,
+            start_byte_index: fragment.end_byte_index,
+            end_byte_index: fragment.end_byte_index,
+            file_path: fragment.file_path,
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "semantic_tests.rs"]
+mod tests;
@@ -0,0 +1,90 @@
+use std::path::Path;
+
+use languages::language_by_filename;
+
+use super::*;
+
+#[test]
+fn test_basic_rust_chunking() {
+    let source_code = r#"
+#[derive(Debug)]
+struct Rectangle {
+    width: u32,
+    height: u32,
+}
+
+impl Rectangle {
+    fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}
+
+fn main() {
+    let rect1 = Rectangle {
+        width: 30,
+        height: 50,
+    };
+
+    println!(
+        "The area of the rectangle is {} square pixels.",
+        rect1.area()
+    );
+}
+"#;
+
+    let max_chunk_size = 128;
+
+    let chunks = chunk_code(
+        source_code,
+        Path::new("test.rs"),
+        max_chunk_size,
+        &language_by_filename(Path::new("test.rs"))
+            .expect("Rust language must exist")
+            .grammar,
+    )
+    .unwrap();
+
+    assert_eq!(chunks.len(), 4);
+
+    // None of the chunks should exceed the chunk size.
+    for chunk in &chunks {
+        assert!(
+            chunk.content.len() <= max_chunk_size,
+            "Chunk should not exceed max size of {max_chunk_size} but was: {}",
+            chunk.content.len()
+        );
+    }
+
+    assert_eq!(
+        chunks[0].content.trim(),
+        r#"#[derive(Debug)]
+struct Rectangle {
+    width: u32,
+    height: u32,
+}"#
+    );
+    assert_eq!(
+        chunks[1].content.trim(),
+        r#"impl Rectangle {
+    fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}"#
+    );
+    assert_eq!(
+        chunks[2].content.trim(),
+        r#"fn main() {
+    let rect1 = Rectangle {
+        width: 30,
+        height: 50,
+    };"#
+    );
+    assert_eq!(
+        chunks[3].content.trim(),
+        r#"println!(
+        "The area of the rectangle is {} square pixels.",
+        rect1.area()
+    );
+}"#
+    );
+}