Skip to main content

codesynapse_core/ts_extract/
markdown.rs

1use super::make_file_node;
2use crate::error::Result;
3use crate::extract::{make_id, ImportNode, LanguageExtractor};
4use crate::types::{Edge, ExtractionFragment, Node};
5use std::collections::HashMap;
6use std::path::Path;
7
8pub struct MarkdownExtractor;
9
10impl MarkdownExtractor {
11    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
12        let (file_id, _, file_node) = make_file_node(path);
13        let mut fragment = ExtractionFragment {
14            nodes: vec![file_node],
15            edges: vec![],
16        };
17
18        let content = std::str::from_utf8(source).unwrap_or("");
19        let str_path = path.to_string_lossy().to_string();
20        let stem = file_id.clone();
21
22        let mut heading_stack: Vec<(usize, String)> = Vec::new();
23        let mut in_code_block = false;
24        let mut code_block_lang: Option<String> = None;
25        let mut code_block_start = 0usize;
26        let mut code_block_first_line: Option<String> = None;
27        let mut code_block_count = 0usize;
28
29        for (line_num_0, line_text) in content.lines().enumerate() {
30            let line_num = line_num_0 + 1;
31            let stripped = line_text.trim();
32
33            if let Some(fence_rest) = stripped.strip_prefix("```") {
34                if !in_code_block {
35                    in_code_block = true;
36                    let after = fence_rest.trim();
37                    code_block_lang = if !after.is_empty() {
38                        after.split_whitespace().next().map(str::to_string)
39                    } else {
40                        None
41                    };
42                    code_block_start = line_num;
43                    code_block_first_line = None;
44                } else {
45                    in_code_block = false;
46                    code_block_count += 1;
47                    let label = match (&code_block_lang, &code_block_first_line) {
48                        (Some(lang), Some(first)) => {
49                            let preview: String = first.chars().take(60).collect();
50                            format!("code:{} ({})", lang, preview)
51                        }
52                        (Some(lang), None) => format!("code:{}", lang),
53                        _ => format!("code:block{}", code_block_count),
54                    };
55                    let cb_nid = make_id(&[&stem, &format!("codeblock_{}", code_block_count)]);
56                    fragment.nodes.push(Node {
57                        id: cb_nid.clone(),
58                        label,
59                        file_type: "document".to_string(),
60                        source_file: str_path.clone(),
61                        source_location: Some(format!("L{}", code_block_start)),
62                        community: None,
63                        rationale: None,
64                        docstring: None,
65                        metadata: HashMap::new(),
66                    });
67                    let parent = heading_stack
68                        .last()
69                        .map(|(_, nid)| nid.clone())
70                        .unwrap_or_else(|| file_id.clone());
71                    fragment.edges.push(Edge {
72                        source: parent,
73                        target: cb_nid,
74                        relation: "contains".to_string(),
75                        confidence: "EXTRACTED".to_string(),
76                        source_file: Some(str_path.clone()),
77                        weight: 1.0,
78                        context: None,
79                    });
80                }
81                continue;
82            }
83
84            if in_code_block {
85                if code_block_first_line.is_none() && !stripped.is_empty() {
86                    code_block_first_line = Some(stripped.to_string());
87                }
88                continue;
89            }
90
91            // ATX headings: ^(#{1,6})\s+(.+)
92            let hash_count = line_text.chars().take_while(|&c| c == '#').count();
93            if (1..=6).contains(&hash_count) {
94                let after = &line_text[hash_count..];
95                if after.starts_with(' ') || after.starts_with('\t') {
96                    let title = after.trim();
97                    if !title.is_empty() {
98                        let h_nid_base = make_id(&[&stem, title]);
99                        let h_nid = if fragment.nodes.iter().any(|n| n.id == h_nid_base) {
100                            make_id(&[&stem, title, &line_num.to_string()])
101                        } else {
102                            h_nid_base
103                        };
104                        fragment.nodes.push(Node {
105                            id: h_nid.clone(),
106                            label: title.to_string(),
107                            file_type: "document".to_string(),
108                            source_file: str_path.clone(),
109                            source_location: Some(format!("L{}", line_num)),
110                            community: None,
111                            rationale: None,
112                            docstring: None,
113                            metadata: HashMap::new(),
114                        });
115                        while heading_stack
116                            .last()
117                            .map(|(l, _)| *l >= hash_count)
118                            .unwrap_or(false)
119                        {
120                            heading_stack.pop();
121                        }
122                        let parent = heading_stack
123                            .last()
124                            .map(|(_, nid)| nid.clone())
125                            .unwrap_or_else(|| file_id.clone());
126                        fragment.edges.push(Edge {
127                            source: parent,
128                            target: h_nid.clone(),
129                            relation: "contains".to_string(),
130                            confidence: "EXTRACTED".to_string(),
131                            source_file: Some(str_path.clone()),
132                            weight: 1.0,
133                            context: None,
134                        });
135                        heading_stack.push((hash_count, h_nid));
136                    }
137                }
138            }
139        }
140
141        Ok(fragment)
142    }
143}
144
145impl LanguageExtractor for MarkdownExtractor {
146    fn file_extensions(&self) -> Vec<&'static str> {
147        vec!["md", "mdx"]
148    }
149    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
150        Self::extract(source, path)
151    }
152    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
153        vec![]
154    }
155    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
156}