code_graph/
lib.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::{fs::read_dir, path::Path};
4
5use eframe::egui::{CollapsingHeader, Ui};
6use egui::{emath, Color32, Pos2, Rect, Stroke, Vec2};
7use lang::{CQuery, JavaQuery, JsQuery, RustQuery, SymbolQuery};
8use lazy_static::lazy_static;
9use tree_sitter::Node;
10use tree_sitter::Parser;
11use uuid::Uuid;
12
13pub mod lang;
14
15#[derive(Clone, PartialEq)]
16pub enum TreeEvent {
17    Clicked(String),
18    None,
19}
20#[derive(Debug, Clone, PartialEq)]
21pub enum TreeType {
22    File,
23    Directory,
24}
25
26#[derive(Clone, Default, Debug)]
27pub struct Tree {
28    pub label: String,
29    full_path: String,
30    select_path: String,
31    children: Vec<Tree>,
32    tree_type: Option<TreeType>,
33    clicked: bool,
34}
35
36impl Tree {
37    pub fn new(name: &str, full_path: &str, tree_type: TreeType) -> Self {
38        Self {
39            label: name.to_owned(),
40            full_path: full_path.to_owned(),
41            children: vec![],
42            tree_type: Some(tree_type),
43            clicked: false,
44            select_path: "".to_owned(),
45        }
46    }
47
48    pub fn ui(&mut self, ui: &mut Ui) -> TreeEvent {
49        let root_name = self.label.clone();
50        self.ui_impl(ui, 0, root_name.as_str())
51    }
52}
53
54impl Tree {
55    fn ui_impl(&mut self, ui: &mut Ui, depth: usize, name: &str) -> TreeEvent {
56        let tree_type = self.tree_type.clone().unwrap_or(TreeType::File);
57        if self.children.len() > 0 || tree_type == TreeType::Directory {
58            return CollapsingHeader::new(name)
59                .default_open(depth < 1)
60                .show(ui, |ui| self.children_ui(ui, depth))
61                .body_returned
62                .unwrap_or(TreeEvent::None);
63        } else {
64            let full_path = self.full_path.clone();
65            if ui
66                .selectable_value(&mut self.select_path, full_path, name)
67                .clicked()
68            {
69                return TreeEvent::Clicked(self.full_path.to_string());
70            }
71            return TreeEvent::None;
72        }
73    }
74
75    pub fn clicked(&self) -> bool {
76        if self.clicked {
77            return true;
78        } else if self.children.len() > 0 {
79            for child in &self.children {
80                if child.clicked() {
81                    return true;
82                }
83            }
84        }
85        return false;
86    }
87
88    fn children_ui(&mut self, ui: &mut Ui, depth: usize) -> TreeEvent {
89        for ele in &mut self.children {
90            let name = ele.label.clone();
91            let event = ele.ui_impl(ui, depth + 1, &name);
92            if let TreeEvent::Clicked(_) = event {
93                return event;
94            }
95        }
96        TreeEvent::None
97    }
98}
99pub fn recursion_dir(root_path: &Path, pathes: &mut Vec<PathBuf>, mut root_tree: Tree) -> Tree {
100    if root_path.is_dir() {
101        for entry in read_dir(root_path).expect("Error read Dir") {
102            let dir_entry = entry.expect("Error");
103            let path_buf = dir_entry.path();
104            let is_dir = path_buf.is_dir();
105            let tree_type = if is_dir {
106                TreeType::Directory
107            } else {
108                TreeType::File
109            };
110            let mut tree = Tree::new(
111                path_buf.file_name().unwrap().to_str().unwrap(),
112                path_buf.as_os_str().to_str().unwrap(),
113                tree_type,
114            );
115            if path_buf.is_dir() {
116                tree = recursion_dir(path_buf.as_path(), pathes, tree);
117            } else if path_buf.is_file() {
118                pathes.push(path_buf);
119            }
120            root_tree.children.push(tree);
121        }
122    }
123    return root_tree;
124}
125#[derive(Debug, Clone, Hash, PartialEq, Eq)]
126pub enum CodeBlockType {
127    FUNCTION,
128    METHOD,
129    STRUCT,
130    IMPL,
131    CLASS,
132    CONST,
133    NORMAL,
134    CALL,
135}
136#[derive(Debug, Clone)]
137pub struct CodeNode {
138    id: String,
139    // 标签
140    pub label: String,
141    // 代码内容
142    pub block: String,
143    // 文件定位 LineNumber
144    pub file_location: usize,
145    // 文件路径
146    pub file_path: String,
147    // 等级
148    level: usize,
149    // block
150    block_type: CodeBlockType,
151    // position
152    position: Pos2,
153}
154
155impl Default for CodeNode {
156    fn default() -> Self {
157        Self {
158            block_type: CodeBlockType::NORMAL,
159            id: "".to_owned(),
160            label: "".to_owned(),
161            block: "".to_owned(),
162            file_location: 0,
163            level: 0,
164            file_path: "".to_owned(),
165            position: Pos2::ZERO,
166        }
167    }
168}
169
170impl CodeNode {
171    pub fn new(
172        id: &str,
173        label: &str,
174        block: &str,
175        file_location: usize,
176        block_type: CodeBlockType,
177        level: usize,
178    ) -> Self {
179        Self {
180            id: id.to_owned(),
181            label: label.to_owned(),
182            block: block.to_owned(),
183            file_location: file_location.to_owned(),
184            file_path: "".to_owned(),
185            block_type,
186            position: Pos2::new(0.0, 0.0),
187            level,
188        }
189    }
190}
191#[derive(Clone, Copy)]
192pub struct CodeNodeIndex(usize);
193
194pub struct Edge {
195    from: usize,
196    to: usize,
197}
198
199pub struct Graph {
200    nodes: Vec<CodeNode>,
201    edges: Vec<Edge>,
202    focus_node: Option<CodeNodeIndex>,
203}
204
205lazy_static! {
206    static ref BLOCK_TYPE_DARK_COLORS: HashMap<CodeBlockType, egui::Color32> = {
207        let mut m = HashMap::new();
208        m.insert(CodeBlockType::NORMAL, egui::Color32::DARK_GRAY);
209        m.insert(CodeBlockType::FUNCTION, egui::Color32::DARK_BLUE);
210        m.insert(CodeBlockType::STRUCT, egui::Color32::from_rgb(204, 112, 0));
211        m.insert(CodeBlockType::CONST, egui::Color32::from_rgb(204, 112, 0));
212        m.insert(CodeBlockType::CLASS, egui::Color32::DARK_GREEN);
213        m
214    };
215    static ref BLOCK_TYPE_LIGHT_COLORS: HashMap<CodeBlockType, egui::Color32> = {
216        let mut m = HashMap::new();
217        m.insert(CodeBlockType::NORMAL, egui::Color32::LIGHT_GRAY);
218        m.insert(CodeBlockType::FUNCTION, egui::Color32::LIGHT_BLUE);
219        m.insert(CodeBlockType::STRUCT, egui::Color32::LIGHT_YELLOW);
220        m.insert(CodeBlockType::CONST, egui::Color32::LIGHT_YELLOW);
221        m.insert(CodeBlockType::CLASS, egui::Color32::LIGHT_GREEN);
222        m
223    };
224}
225
226impl Graph {
227    pub fn new() -> Self {
228        Self {
229            nodes: vec![],
230            edges: vec![],
231            focus_node: None,
232        }
233    }
234
235    pub fn get_focus_idx(&mut self) -> Option<CodeNodeIndex> {
236        return self.focus_node;
237    }
238
239    pub fn add_node(&mut self, node: CodeNode) -> CodeNodeIndex {
240        let index = self.nodes.len();
241        self.nodes.push(node);
242        return CodeNodeIndex(index);
243    }
244
245    pub fn add_edge(&mut self, from: CodeNodeIndex, to: CodeNodeIndex) {
246        self.edges.push(Edge {
247            from: from.0,
248            to: to.0,
249        })
250    }
251
252    pub fn clear(&mut self) {
253        self.nodes.clear();
254        self.edges.clear();
255        self.focus_node = None;
256    }
257    /**
258     * 对节点进行布局
259     */
260    pub fn layout(&mut self, ui: &mut Ui) {
261        let (_, painter) = ui.allocate_painter(ui.available_size(), egui::Sense::click());
262        let mut sum_height = 0.0;
263        // 直线布局
264        for (index, node) in self.nodes.iter_mut().enumerate() {
265            let text_size = painter
266                .layout_no_wrap(
267                    node.label.clone(),
268                    egui::FontId::default(),
269                    egui::Color32::WHITE,
270                )
271                .size();
272            node.position = Pos2::new(
273                ui.available_size().x / 2.0 + node.level as f32 * 20.0,
274                index as f32 * 16.0 + sum_height + 32.0,
275            );
276            sum_height += text_size.y;
277        }
278    }
279
280    pub fn ui(&mut self, ui: &mut Ui) -> egui::Response {
281        let (response, painter) =
282            ui.allocate_painter(ui.available_size(), egui::Sense::click_and_drag());
283
284        let focus_stroke_color;
285        let stroke_color;
286        let text_color;
287        let grid_color;
288
289        if ui.ctx().style().visuals.dark_mode {
290            stroke_color = egui::Color32::LIGHT_GRAY;
291            text_color = egui::Color32::WHITE;
292            focus_stroke_color = egui::Color32::LIGHT_BLUE;
293            grid_color = Color32::from_gray(50);
294        } else {
295            focus_stroke_color = egui::Color32::BLUE;
296            stroke_color = egui::Color32::DARK_GRAY;
297            text_color = egui::Color32::DARK_GRAY;
298            grid_color = Color32::from_gray(220);
299        }
300
301        // 获取可用区域
302        let rect = ui.max_rect();
303
304        // 定义网格参数
305        let cell_size = 10.0; // 网格单元格大小
306        let stroke = Stroke::new(0.5, grid_color); // 线条宽度和颜色
307
308        // 绘制垂直线
309        let mut x = rect.left();
310        while x <= rect.right() {
311            let line = [Pos2::new(x, rect.top()), Pos2::new(x, rect.bottom())];
312            painter.line_segment(line, stroke);
313            x += cell_size;
314        }
315
316        // 绘制水平线
317        let mut y = rect.top();
318        while y <= rect.bottom() {
319            let line = [Pos2::new(rect.left(), y), Pos2::new(rect.right(), y)];
320            painter.line_segment(line, stroke);
321            y += cell_size;
322        }
323
324        let to_screen = emath::RectTransform::from_to(
325            Rect::from_min_size(Pos2::ZERO, response.rect.size()),
326            response.rect,
327        );
328        let mut node_size_list = vec![];
329
330        // 绘制节点
331        for (index, node) in self.nodes.iter_mut().enumerate() {
332            let node_pos = to_screen.transform_pos(node.position);
333            let text_size = painter
334                .layout_no_wrap(
335                    node.label.clone(),
336                    egui::FontId::default(),
337                    egui::Color32::WHITE,
338                )
339                .size();
340
341            node_size_list.push(text_size + Vec2::new(16.0, 8.0));
342
343            let rect = egui::Rect::from_min_size(
344                node_pos,
345                egui::vec2(text_size.x + 16.0, text_size.y + 8.0),
346            );
347            let fill_color = if ui.ctx().style().visuals.dark_mode {
348                BLOCK_TYPE_DARK_COLORS
349                    .get(&node.block_type)
350                    .copied()
351                    .unwrap_or(egui::Color32::DARK_GRAY)
352            } else {
353                BLOCK_TYPE_LIGHT_COLORS
354                    .get(&node.block_type)
355                    .copied()
356                    .unwrap_or(egui::Color32::LIGHT_GRAY)
357            };
358            painter.rect(rect, 5.0, fill_color, Stroke::new(1.0, stroke_color));
359
360            painter.text(
361                node_pos + Vec2::new(8.0, 4.0),
362                egui::Align2::LEFT_TOP,
363                &node.label,
364                egui::FontId::default(),
365                text_color,
366            );
367
368            let point_id = response.id.with(&node.id);
369
370            let node_response = ui.interact(rect, point_id, egui::Sense::click_and_drag());
371            if node_response.dragged() {
372                // 更新节点位置
373                node.position += node_response.drag_delta();
374            }
375            if node_response.clicked() {
376                self.focus_node = Some(CodeNodeIndex(index));
377            }
378            if let Some(f_node) = self.focus_node {
379                if f_node.0 == index {
380                    // ui.ctx().request_repaint();
381                    // let time = ui.input(|i| i.time);
382                    painter.rect(
383                        rect,
384                        5.0,
385                        egui::Color32::TRANSPARENT,
386                        Stroke::new(2.5, focus_stroke_color),
387                    );
388                }
389            }
390
391            if response.dragged() {
392                // 更新节点位置
393                node.position += response.drag_delta();
394            }
395        }
396
397        // 绘制边
398        for edge in &self.edges {
399            let from = to_screen.transform_pos(self.nodes[edge.from].position)
400                + Vec2::new(0.0, node_size_list[edge.from].y / 2.0);
401            let to = to_screen.transform_pos(self.nodes[edge.to].position)
402                + Vec2::new(0.0, node_size_list[edge.to].y / 2.0);
403            painter.line_segment(
404                [from, from + Vec2::new(-10.0, 0.0)],
405                (1.0, egui::Color32::GRAY),
406            );
407            painter.line_segment(
408                [from + Vec2::new(-10.0, 0.0), Pos2::new(from.x - 10.0, to.y)],
409                (1.0, egui::Color32::GRAY),
410            );
411            painter.line_segment(
412                [Pos2::new(from.x - 10.0, to.y), to],
413                (1.0, egui::Color32::GRAY),
414            );
415        }
416        response
417    }
418
419    pub fn get_node(&mut self, index: CodeNodeIndex) -> CodeNode {
420        let default_node = CodeNode::default();
421        let node = self.nodes.get(index.0).unwrap_or(&default_node);
422        return node.clone();
423    }
424
425    pub fn node_index(&mut self, node_id: &str) -> CodeNodeIndex {
426        for (index, node) in self.nodes.iter().enumerate() {
427            if node.id == node_id {
428                return CodeNodeIndex(index);
429            }
430        }
431        CodeNodeIndex(0)
432    }
433}
434
435pub fn valid_file_extention(extension: &str) -> bool {
436    return vec!["rs", "c", "h", "java", "js", "jsx"].contains(&extension);
437}
438
439pub fn get_symbol_query(extention: &str) -> Box<dyn SymbolQuery> {
440    match extention {
441        "rs" => Box::new(RustQuery),
442        "java" => Box::new(JavaQuery),
443        "c" | "h" => Box::new(CQuery),
444        "js" | "jsx" => Box::new(JsQuery),
445        _ => Box::new(RustQuery),
446    }
447}
448
449pub fn fetch_calls(path: &str, code: &str, symbol_query: Box<dyn SymbolQuery>) -> Vec<CodeNode> {
450    let mut parser = Parser::new();
451    parser
452        .set_language(&symbol_query.get_lang())
453        .expect("Error load Rust grammer");
454    let tree = parser.parse(code, None).unwrap();
455    let root_node = tree.root_node();
456    recursion_call(root_node, path, code, &symbol_query)
457}
458
459pub fn recursion_call(
460    node: Node,
461    path: &str,
462    code: &str,
463    symbol_query: &Box<dyn SymbolQuery>,
464) -> Vec<CodeNode> {
465    let mut nodes = vec![];
466    let code_node = symbol_query.get_call(code, &node);
467    if let Some(mut node) = code_node {
468        node.file_path = path.to_string();
469        nodes.push(node);
470    }
471
472    for child in node.children(&mut node.walk()) {
473        let sub_nodes = recursion_call(child, path, code, symbol_query);
474        if sub_nodes.len() > 0 {
475            for sub_node in sub_nodes {
476                nodes.push(sub_node);
477            }
478        }
479    }
480    return nodes;
481}
482/**
483* 打印大纲
484*/
485pub fn fetch_symbols(
486    path: &str,
487    code: &str,
488    symbol_query: Box<dyn SymbolQuery>,
489    graph: &mut Graph,
490) {
491    let mut parser = Parser::new();
492    parser
493        .set_language(&symbol_query.get_lang())
494        .expect("Error load Rust grammer");
495    let tree = parser.parse(code, None).unwrap();
496    let root_node = tree.root_node();
497    let root_code_node = CodeNode::new(
498        format!("{}", Uuid::new_v4()).as_str(),
499        path,
500        code,
501        0,
502        CodeBlockType::NORMAL,
503        0,
504    );
505    graph.add_node(root_code_node);
506    recursion_outline(
507        root_node,
508        CodeNodeIndex(0),
509        path,
510        code,
511        1,
512        &symbol_query,
513        graph,
514    );
515}
516
517pub fn recursion_outline(
518    node: Node,
519    parent_id: CodeNodeIndex,
520    path: &str,
521    code: &str,
522    level: usize,
523    symbol_query: &Box<dyn SymbolQuery>,
524    graph: &mut Graph,
525) {
526    let mut current_id = parent_id;
527    let code_node = symbol_query.get_definition(code, &node);
528    let mut level = level;
529    if let Some(mut node) = code_node {
530        node.file_path = path.to_string();
531        node.level = level;
532        let index = graph.add_node(node);
533        current_id = index;
534        graph.add_edge(parent_id, index);
535        level += 1;
536    }
537
538    for child in node.children(&mut node.walk()) {
539        recursion_outline(child, current_id, path, code, level, symbol_query, graph)
540    }
541}
code_graph/lib.rs

code_graph/
lib.rs
