devgen_splitter/splitter/
entity_splitter.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
use super::{
    CodeChunk,
    CodeEntity,
    EntityType,
    SplitOptions,
    Splitter,
};
use anyhow::Result;
use std::ops::Range;
use tree_sitter::Node;

impl Splitter {
    pub fn build_header_for_entities(entities: &Vec<CodeEntity>) -> Option<String> {
        if entities.is_empty() {
            return None;
        }

        // Check if all entities are methods and belong to the same class
        let all_methods = entities.iter().all(|e| e.entity_type == EntityType::Method);
        let same_class = entities
            .iter()
            .all(|e| e.parent_name == entities[0].parent_name);

        if all_methods && same_class {
            // Get the class name from the first entity's parent_name
            if let Some(class_name) = &entities[0].parent_name {
                return Some(format!(
                    "The following methods belong to class '{}'",
                    class_name
                ));
            }
        }

        None
    }

    pub fn build_header(entity: &CodeEntity) -> Option<String> {
        match entity.entity_type {
            EntityType::Class => Some(format!(
                "The incomplete part comes from the '{}' class definition",
                entity.name
            )),
            EntityType::Interface => Some(format!(
                "The incomplete part comes from the '{}' interface definition",
                entity.name
            )),
            EntityType::Function => Some(format!(
                "The incomplete part comes from the '{}' function implementation",
                entity.name
            )),
            EntityType::Method => Some(format!(
                "The incomplete part comes from the '{}' method implementation in the '{}' class",
                entity.name,
                entity.parent_name.as_deref().unwrap_or("")
            )),
            EntityType::Enum => Some(format!(
                "The incomplete part comes from the '{}' enum definition",
                entity.name
            )),
        }
    }

    pub fn split_entity<'a>(
        last_chunk_end_line: usize,
        entity: &CodeEntity,
        nodes: &Vec<Node<'a>>,
        options: &SplitOptions,
    ) -> Result<(Vec<CodeChunk>, usize)> {
        let mut chunks = Vec::new();
        let mut current_chunk_end_line = last_chunk_end_line;
        let header = if options.enable_header {
            Self::build_header(entity)
        } else {
            None
        };
        for node in nodes {
            current_chunk_end_line = Self::chunk_entity(
                node,
                current_chunk_end_line,
                &mut chunks,
                options,
                header.clone(),
            )?;
        }
        Ok((chunks, current_chunk_end_line))
    }

    fn chunk_entity(
        node: &Node,
        last_chunk_end_line: usize,
        chunks: &mut Vec<CodeChunk>,
        options: &SplitOptions,
        header: Option<String>,
    ) -> Result<usize> {
        let mut local_last_chunk_end_line = last_chunk_end_line;
        for i in 0..node.child_count() {
            let child = node.child(i).expect("Failed to get child node");
            let child_end_line = child.end_position().row;
            let child_start_line = child.start_position().row;
            let left_chunk_line_count = child_start_line - local_last_chunk_end_line;
            if left_chunk_line_count > options.chunk_line_limit {
                let chunk = CodeChunk {
                    line_range: Range {
                        start: local_last_chunk_end_line,
                        end: child_start_line,
                    },
                    entities: Vec::new(),
                    header: header.clone(),
                };
                chunks.push(chunk);
                local_last_chunk_end_line = child_start_line;
            }
            let node_line_count = child_end_line - child_start_line;
            if node_line_count > options.chunk_line_limit {
                if child.child_count() > 0 {
                    let update_last_chunk_end_line = Self::chunk_entity(
                        &child,
                        local_last_chunk_end_line,
                        chunks,
                        options,
                        header.clone(),
                    )?;
                    local_last_chunk_end_line = update_last_chunk_end_line;
                } else {
                    let chunk = CodeChunk {
                        line_range: Range {
                            start: local_last_chunk_end_line,
                            end: child_end_line,
                        },
                        entities: Vec::new(),
                        header: header.clone(),
                    };
                    chunks.push(chunk);
                    local_last_chunk_end_line = child_end_line;
                }
            } else if node_line_count + left_chunk_line_count >= options.chunk_line_limit {
                let chunk = CodeChunk {
                    line_range: Range {
                        start: local_last_chunk_end_line,
                        end: child_end_line,
                    },
                    entities: Vec::new(),
                    header: header.clone(),
                };
                chunks.push(chunk);
                local_last_chunk_end_line = child_end_line;
            }
        }
        Ok(local_last_chunk_end_line)
    }
}