1use std::cell::RefCell;
2use std::collections::HashMap;
3
4use tree_sitter::{Language, Parser, Tree};
5
6use crate::core::code_graph::CodeGraph;
7use crate::core::moniker::Moniker;
8
9use crate::lang::canonical_walker::CanonicalWalker;
10
11pub mod build;
12mod canonicalize;
13mod kinds;
14mod strategy;
15
16use canonicalize::compute_module_moniker;
17use strategy::{ImportEntry, Strategy, collect_callable_table, collect_type_table};
18
19#[derive(Clone, Debug, Default)]
20pub struct Presets {}
21
22pub fn parse(source: &str) -> Tree {
23 let mut parser = Parser::new();
24 let language: Language = tree_sitter_go::LANGUAGE.into();
25 parser
26 .set_language(&language)
27 .expect("failed to load tree-sitter Go grammar");
28 parser
29 .parse(source, None)
30 .expect("tree-sitter parse returned None on a non-cancelled call")
31}
32
33pub fn extract(
34 uri: &str,
35 source: &str,
36 anchor: &Moniker,
37 deep: bool,
38 _presets: &Presets,
39) -> CodeGraph {
40 let tree = parse(source);
41 let module = compute_module_moniker(anchor, uri);
42 let (def_cap, ref_cap) = CodeGraph::capacity_for_source(source.len());
43 let mut graph = CodeGraph::with_capacity(module.clone(), kinds::MODULE, def_cap, ref_cap);
44 let mut type_table: HashMap<&[u8], Moniker> = HashMap::new();
45 collect_type_table(
46 tree.root_node(),
47 source.as_bytes(),
48 &module,
49 &mut graph,
50 &mut type_table,
51 );
52 let mut callable_table: HashMap<(Moniker, Vec<u8>), Vec<u8>> = HashMap::new();
53 collect_callable_table(
54 tree.root_node(),
55 source.as_bytes(),
56 &module,
57 &type_table,
58 &mut callable_table,
59 );
60 let strat = Strategy {
61 module: module.clone(),
62 source_bytes: source.as_bytes(),
63 deep,
64 imports: RefCell::new(HashMap::<Vec<u8>, ImportEntry>::new()),
65 local_scope: RefCell::new(Vec::new()),
66 type_table,
67 callable_table,
68 };
69 let walker = CanonicalWalker::new(&strat, source.as_bytes());
70 walker.walk(tree.root_node(), &module, &mut graph);
71 graph
72}
73
74pub struct Lang;
75
76impl crate::lang::LangExtractor for Lang {
77 type Presets = Presets;
78 const LANG_TAG: &'static str = "go";
79 const ALLOWED_KINDS: &'static [&'static str] = &[
80 "type",
81 "struct",
82 "interface",
83 "func",
84 "method",
85 "var",
86 "const",
87 ];
88 const ALLOWED_VISIBILITIES: &'static [&'static str] = &["public", "module"];
89
90 fn extract(
91 uri: &str,
92 source: &str,
93 anchor: &Moniker,
94 deep: bool,
95 presets: &Self::Presets,
96 ) -> CodeGraph {
97 extract(uri, source, anchor, deep, presets)
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104 use crate::core::moniker::MonikerBuilder;
105 use crate::lang::assert_conformance;
106
107 fn make_anchor() -> Moniker {
108 MonikerBuilder::new().project(b"app").build()
109 }
110
111 fn extract_default(uri: &str, source: &str, anchor: &Moniker, deep: bool) -> CodeGraph {
112 let g = extract(uri, source, anchor, deep, &Presets::default());
113 assert_conformance::<super::Lang>(&g, anchor);
114 g
115 }
116
117 #[test]
118 fn parse_empty_returns_source_file() {
119 let tree = parse("");
120 assert_eq!(tree.root_node().kind(), "source_file");
121 }
122
123 #[test]
124 fn extract_module_uses_path_segments() {
125 let g = extract_default("acme/util/text.go", "package text\n", &make_anchor(), false);
126 let expected = MonikerBuilder::new()
127 .project(b"app")
128 .segment(b"lang", b"go")
129 .segment(b"package", b"acme")
130 .segment(b"package", b"util")
131 .segment(b"module", b"text")
132 .build();
133 assert_eq!(g.root(), &expected);
134 }
135
136 #[test]
137 fn extract_method_when_type_declared_after_method() {
138 let src = "package foo\nfunc (r *Foo) Bar() {}\ntype Foo struct{}\n";
139 let g = extract_default("foo.go", src, &make_anchor(), false);
140 let bar = MonikerBuilder::new()
141 .project(b"app")
142 .segment(b"lang", b"go")
143 .segment(b"module", b"foo")
144 .segment(b"struct", b"Foo")
145 .segment(b"method", b"Bar()")
146 .build();
147 assert!(
148 g.contains(&bar),
149 "method emitted before its type declaration must still be reparented; defs: {:?}",
150 g.def_monikers()
151 );
152 }
153
154 #[test]
155 fn extract_simple_call_to_unresolved_callee_uses_name_only() {
156 let src = "package foo\nfunc Run() { Helper(1, 2) }\n";
157 let g = extract_default("foo.go", src, &make_anchor(), false);
158 let r = g
159 .refs()
160 .find(|r| {
161 r.kind == b"calls"
162 && r.target.as_view().segments().last().unwrap().name == b"Helper"
163 })
164 .expect("calls Helper (name-only, no parens)");
165 assert_eq!(r.confidence, b"name_match".to_vec());
166 }
167
168 #[test]
169 fn extract_composite_literal_unresolved_type_marks_name_match() {
170 let src = "package foo\nfunc Run() { _ = Bar{} }\n";
171 let g = extract_default("foo.go", src, &make_anchor(), false);
172 let r = g
173 .refs()
174 .find(|r| r.kind == b"instantiates")
175 .expect("instantiates ref");
176 assert_eq!(r.confidence, b"name_match".to_vec());
177 }
178
179 #[test]
180 fn extract_shallow_skips_param_and_local_defs() {
181 let src = "package foo\nfunc Run(x int) { y := 1; _ = y }\n";
182 let g = extract_default("foo.go", src, &make_anchor(), false);
183 assert!(
184 g.defs().all(|d| d.kind != b"param" && d.kind != b"local"),
185 "shallow extraction must not emit param/local defs"
186 );
187 }
188
189 #[test]
190 fn extract_deep_emits_param_defs_under_function() {
191 let src = "package foo\nfunc Run(a int, b string) {}\n";
192 let g = extract_default("foo.go", src, &make_anchor(), true);
193 let pa = MonikerBuilder::new()
194 .project(b"app")
195 .segment(b"lang", b"go")
196 .segment(b"module", b"foo")
197 .segment(b"func", b"Run(a:int,b:string)")
198 .segment(b"param", b"a")
199 .build();
200 let pb = MonikerBuilder::new()
201 .project(b"app")
202 .segment(b"lang", b"go")
203 .segment(b"module", b"foo")
204 .segment(b"func", b"Run(a:int,b:string)")
205 .segment(b"param", b"b")
206 .build();
207 assert!(g.contains(&pa));
208 assert!(g.contains(&pb));
209 }
210
211 #[test]
212 fn extract_deep_emits_receiver_param_for_method() {
213 let src = "package foo\ntype Foo struct{}\nfunc (r *Foo) Bar(x int) {}\n";
214 let g = extract_default("foo.go", src, &make_anchor(), true);
215 let recv = MonikerBuilder::new()
216 .project(b"app")
217 .segment(b"lang", b"go")
218 .segment(b"module", b"foo")
219 .segment(b"struct", b"Foo")
220 .segment(b"method", b"Bar(x:int)")
221 .segment(b"param", b"r")
222 .build();
223 assert!(g.contains(&recv));
224 }
225
226 #[test]
227 fn extract_deep_skips_blank_param() {
228 let src = "package foo\nfunc Run(_ int, b string) {}\n";
229 let g = extract_default("foo.go", src, &make_anchor(), true);
230 let params: Vec<&[u8]> = g
231 .defs()
232 .filter(|d| d.kind == b"param")
233 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
234 .collect();
235 assert_eq!(params, vec![&b"b"[..]]);
236 }
237
238 #[test]
239 fn extract_deep_emits_local_def_for_short_var() {
240 let src = "package foo\nfunc Run() { x := 1; _ = x }\n";
241 let g = extract_default("foo.go", src, &make_anchor(), true);
242 let lx = MonikerBuilder::new()
243 .project(b"app")
244 .segment(b"lang", b"go")
245 .segment(b"module", b"foo")
246 .segment(b"func", b"Run()")
247 .segment(b"local", b"x")
248 .build();
249 assert!(g.contains(&lx));
250 }
251
252 #[test]
253 fn extract_deep_emits_local_defs_for_multi_assign() {
254 let src = "package foo\nfunc Run() { x, y := 1, 2; _, _ = x, y }\n";
255 let g = extract_default("foo.go", src, &make_anchor(), true);
256 let names: Vec<&[u8]> = g
257 .defs()
258 .filter(|d| d.kind == b"local")
259 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
260 .collect();
261 assert!(names.contains(&&b"x"[..]));
262 assert!(names.contains(&&b"y"[..]));
263 }
264
265 #[test]
266 fn extract_deep_emits_local_def_for_var_declaration() {
267 let src = "package foo\nfunc Run() { var z int = 5; _ = z }\n";
268 let g = extract_default("foo.go", src, &make_anchor(), true);
269 let lz = MonikerBuilder::new()
270 .project(b"app")
271 .segment(b"lang", b"go")
272 .segment(b"module", b"foo")
273 .segment(b"func", b"Run()")
274 .segment(b"local", b"z")
275 .build();
276 assert!(g.contains(&lz));
277 }
278
279 #[test]
280 fn extract_deep_emits_local_defs_for_range_vars() {
281 let src =
282 "package foo\nfunc Run(m map[string]int) { for k, v := range m { _, _ = k, v } }\n";
283 let g = extract_default("foo.go", src, &make_anchor(), true);
284 let names: Vec<&[u8]> = g
285 .defs()
286 .filter(|d| d.kind == b"local")
287 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
288 .collect();
289 assert!(names.contains(&&b"k"[..]));
290 assert!(names.contains(&&b"v"[..]));
291 }
292
293 #[test]
294 fn extract_top_level_var_does_not_pollute_locals() {
295 let src = "package foo\nvar GlobalCount int\nfunc Run() { GlobalCount = 1 }\n";
296 let g = extract_default("foo.go", src, &make_anchor(), false);
297 let local_names: Vec<&[u8]> = g
298 .defs()
299 .filter(|d| d.kind == b"local")
300 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
301 .collect();
302 assert!(
303 local_names.is_empty(),
304 "a package-level var must not be emitted as a local. found locals: {:?}",
305 local_names
306 );
307 let vars: Vec<&[u8]> = g
308 .defs()
309 .filter(|d| d.kind == b"var")
310 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
311 .collect();
312 assert_eq!(vars, vec![&b"GlobalCount"[..]]);
313 }
314
315 #[test]
316 fn extract_deep_skips_blank_in_short_var() {
317 let src = "package foo\nfunc Run() { _, y := 1, 2; _ = y }\n";
318 let g = extract_default("foo.go", src, &make_anchor(), true);
319 let names: Vec<&[u8]> = g
320 .defs()
321 .filter(|d| d.kind == b"local")
322 .map(|d| d.moniker.as_view().segments().last().unwrap().name)
323 .collect();
324 assert_eq!(names, vec![&b"y"[..]]);
325 }
326}