1use crate::completion::{
2 context::TreeResolver,
3 model::{
4 CommandLine, CompletionAnalysis, CompletionContext, CompletionNode, CompletionTree,
5 ContextScope, CursorState, MatchKind, ParsedLine, SuggestionOutput, TailItem,
6 },
7 parse::CommandLineParser,
8 suggest::SuggestionEngine,
9};
10use std::collections::BTreeSet;
11
12#[derive(Debug, Clone)]
13pub struct CompletionEngine {
14 parser: CommandLineParser,
15 suggester: SuggestionEngine,
16 tree: CompletionTree,
17 global_context_flags: BTreeSet<String>,
18}
19
20impl CompletionEngine {
21 pub fn new(tree: CompletionTree) -> Self {
22 let global_context_flags = collect_global_context_flags(&tree.root);
23 Self {
24 parser: CommandLineParser,
25 suggester: SuggestionEngine::new(tree.clone()),
26 tree,
27 global_context_flags,
28 }
29 }
30
31 pub fn complete(&self, line: &str, cursor: usize) -> (CursorState, Vec<SuggestionOutput>) {
32 let analysis = self.analyze(line, cursor);
33 let suggestions = self.suggestions_for_analysis(&analysis);
34 (analysis.cursor, suggestions)
35 }
36
37 pub fn suggestions_for_analysis(&self, analysis: &CompletionAnalysis) -> Vec<SuggestionOutput> {
38 self.suggester.generate(analysis)
39 }
40
41 pub fn analyze(&self, line: &str, cursor: usize) -> CompletionAnalysis {
42 let parsed = self.parser.analyze(line, cursor);
43
44 self.analyze_command_parts(parsed.parsed, parsed.cursor)
45 }
46
47 pub fn analyze_command(
48 &self,
49 full_cmd: CommandLine,
50 cursor_cmd: CommandLine,
51 cursor: CursorState,
52 ) -> CompletionAnalysis {
53 self.analyze_command_parts(
54 ParsedLine {
55 safe_cursor: 0,
56 full_tokens: Vec::new(),
57 cursor_tokens: Vec::new(),
58 full_cmd,
59 cursor_cmd,
60 },
61 cursor,
62 )
63 }
64
65 fn analyze_command_parts(
66 &self,
67 mut parsed: ParsedLine,
68 cursor: CursorState,
69 ) -> CompletionAnalysis {
70 let mut context =
71 self.resolve_completion_context(&parsed.cursor_cmd, cursor.token_stub.as_str());
72 self.merge_prefilled_values(&mut parsed.cursor_cmd, &context.matched_path);
73 context = self.resolve_completion_context(&parsed.cursor_cmd, cursor.token_stub.as_str());
74
75 if !parsed.cursor_cmd.has_pipe() {
79 self.merge_context_flags(
80 &mut parsed.cursor_cmd,
81 &parsed.full_cmd,
82 cursor.token_stub.as_str(),
83 );
84 }
85
86 CompletionAnalysis {
87 parsed,
88 cursor,
89 context,
90 }
91 }
92
93 pub fn tokenize(&self, line: &str) -> Vec<String> {
94 self.parser.tokenize(line)
95 }
96
97 pub fn matched_command_len_tokens(&self, tokens: &[String]) -> usize {
98 TreeResolver::new(&self.tree).matched_command_len_tokens(tokens)
99 }
100
101 pub fn classify_match(&self, analysis: &CompletionAnalysis, value: &str) -> MatchKind {
102 if analysis.parsed.cursor_cmd.has_pipe() {
103 return MatchKind::Pipe;
104 }
105 let nodes = TreeResolver::new(&self.tree).resolved_nodes(&analysis.context);
106
107 if value.starts_with("--") || nodes.flag_scope_node.flags.contains_key(value) {
108 return MatchKind::Flag;
109 }
110 if nodes.context_node.children.contains_key(value) {
111 return if analysis.context.matched_path.is_empty() {
112 MatchKind::Command
113 } else {
114 MatchKind::Subcommand
115 };
116 }
117 MatchKind::Value
118 }
119
120 fn merge_context_flags(
121 &self,
122 cursor_cmd: &mut CommandLine,
123 full_cmd: &CommandLine,
124 stub: &str,
125 ) {
126 let context = self.resolve_completion_context(cursor_cmd, stub);
127 let mut scoped_flags = BTreeSet::new();
128 let resolver = TreeResolver::new(&self.tree);
129 for i in (0..=context.matched_path.len()).rev() {
130 let (node, matched) = resolver.resolve_context(&context.matched_path[..i]);
131 if matched.len() == i {
132 scoped_flags.extend(node.flags.keys().cloned());
133 }
134 }
135 scoped_flags.extend(self.global_context_flags.iter().cloned());
136
137 for item in full_cmd.tail().iter().skip(cursor_cmd.tail_len()) {
138 let TailItem::Flag(flag) = item else {
139 continue;
140 };
141 if cursor_cmd.has_flag(&flag.name) {
142 continue;
143 }
144 if !scoped_flags.contains(&flag.name) {
145 continue;
146 }
147 cursor_cmd.merge_flag_values(flag.name.clone(), flag.values.clone());
148 }
149 }
150
151 fn merge_prefilled_values(&self, cursor_cmd: &mut CommandLine, matched_path: &[String]) {
152 let resolver = TreeResolver::new(&self.tree);
153 let mut prefilled_positionals = Vec::new();
154 for i in 0..=matched_path.len() {
155 let Some(node) = resolver.resolve_exact(&matched_path[..i]) else {
156 continue;
157 };
158 prefilled_positionals.extend(node.prefilled_positionals.iter().cloned());
159 for (flag, values) in &node.prefilled_flags {
160 if cursor_cmd.has_flag(flag) {
161 continue;
162 }
163 cursor_cmd.merge_flag_values(flag.clone(), values.clone());
164 }
165 }
166 cursor_cmd.prepend_positional_values(prefilled_positionals);
167 }
168
169 fn resolve_completion_context(&self, cmd: &CommandLine, stub: &str) -> CompletionContext {
170 let resolver = TreeResolver::new(&self.tree);
171 let (pre_node, _) = resolver.resolve_context(cmd.head());
172 let has_subcommands = !pre_node.children.is_empty();
173 let head_without_partial_subcommand =
177 if !stub.is_empty() && !stub.starts_with('-') && has_subcommands {
178 &cmd.head()[..cmd.head().len().saturating_sub(1)]
179 } else {
180 cmd.head()
181 };
182 let (_, matched) = resolver.resolve_context(head_without_partial_subcommand);
183 let flag_scope_path = resolver.resolve_flag_scope_path(&matched);
184
185 let arg_tokens: Vec<String> = cmd
186 .head()
187 .iter()
188 .skip(matched.len())
189 .filter(|token| token.as_str() != stub)
190 .cloned()
191 .chain(
192 cmd.positional_args()
193 .filter(|token| token.as_str() != stub)
194 .cloned(),
195 )
196 .collect();
197
198 let context_node = resolver.resolve_exact(&matched).unwrap_or(&self.tree.root);
199 let subcommand_context =
200 context_node.value_key || (has_subcommands && arg_tokens.is_empty());
201
202 CompletionContext {
203 matched_path: matched,
204 flag_scope_path,
205 subcommand_context,
206 }
207 }
208}
209
210fn collect_global_context_flags(root: &CompletionNode) -> BTreeSet<String> {
211 fn walk(node: &CompletionNode, out: &mut BTreeSet<String>) {
212 for (name, flag) in &node.flags {
213 if flag.context_only && flag.context_scope == ContextScope::Global {
214 out.insert(name.clone());
215 }
216 }
217 for child in node.children.values() {
218 walk(child, out);
219 }
220 }
221
222 let mut out = BTreeSet::new();
223 walk(root, &mut out);
224 out
225}
226
227#[cfg(test)]
228mod tests {
229 use std::collections::BTreeMap;
230
231 use crate::completion::{
232 CompletionEngine,
233 model::{
234 CompletionNode, CompletionTree, ContextScope, FlagNode, QuoteStyle, SuggestionEntry,
235 SuggestionOutput,
236 },
237 };
238
239 fn tree() -> CompletionTree {
240 let mut provision = CompletionNode::default();
241 provision.flags.insert(
242 "--provider".to_string(),
243 FlagNode {
244 suggestions: vec![
245 SuggestionEntry::from("vmware"),
246 SuggestionEntry::from("nrec"),
247 ],
248 context_only: true,
249 ..FlagNode::default()
250 },
251 );
252 provision.flags.insert(
253 "--os".to_string(),
254 FlagNode {
255 suggestions_by_provider: BTreeMap::from([
256 ("vmware".to_string(), vec![SuggestionEntry::from("rhel")]),
257 ("nrec".to_string(), vec![SuggestionEntry::from("alma")]),
258 ]),
259 suggestions: vec![SuggestionEntry::from("rhel"), SuggestionEntry::from("alma")],
260 context_only: true,
261 ..FlagNode::default()
262 },
263 );
264
265 let mut orch = CompletionNode::default();
266 orch.children.insert("provision".to_string(), provision);
267
268 CompletionTree {
269 root: CompletionNode::default().with_child("orch", orch),
270 ..CompletionTree::default()
271 }
272 }
273
274 #[test]
275 fn merges_late_provider_flag_into_cursor_context() {
276 let engine = CompletionEngine::new(tree());
277 let line = "orch provision --os --provider vmware";
278 let cursor = line.find("--provider").expect("provider in test line") - 1;
279
280 let (_, suggestions) = engine.complete(line, cursor);
281 let values: Vec<String> = suggestions
282 .into_iter()
283 .filter_map(|entry| match entry {
284 SuggestionOutput::Item(item) => Some(item.text),
285 SuggestionOutput::PathSentinel => None,
286 })
287 .collect();
288
289 assert!(values.contains(&"rhel".to_string()));
290 }
291
292 #[test]
293 fn hides_flags_already_present_later_in_line() {
294 let engine = CompletionEngine::new(tree());
295 let line = "orch provision --provider vmware";
296 let cursor = line.find("--provider").expect("provider in test line") - 2;
297
298 let (_, suggestions) = engine.complete(line, cursor);
299 let values: Vec<String> = suggestions
300 .into_iter()
301 .filter_map(|entry| match entry {
302 SuggestionOutput::Item(item) => Some(item.text),
303 SuggestionOutput::PathSentinel => None,
304 })
305 .collect();
306
307 assert!(!values.contains(&"--provider".to_string()));
308 }
309
310 #[test]
311 fn supports_non_char_boundary_cursor_without_panicking() {
312 let engine = CompletionEngine::new(tree());
313 let line = "orch å";
314 let cursor = line.find('å').expect("multibyte char should exist") + 1;
315 let (_cursor, _suggestions) = engine.complete(line, cursor);
316 }
317
318 #[test]
319 fn equals_flag_without_value_still_requests_suggestions() {
320 let engine = CompletionEngine::new(tree());
321 let line = "orch provision --os=";
322 let (_, suggestions) = engine.complete(line, line.len());
323 let values: Vec<String> = suggestions
324 .into_iter()
325 .filter_map(|entry| match entry {
326 SuggestionOutput::Item(item) => Some(item.text),
327 SuggestionOutput::PathSentinel => None,
328 })
329 .collect();
330 assert!(values.contains(&"rhel".to_string()));
331 assert!(values.contains(&"alma".to_string()));
332 }
333
334 #[test]
335 fn merges_context_flags_from_metadata_even_if_not_in_scope() {
336 let mut provision = CompletionNode::default();
337 provision.flags.insert(
338 "--os".to_string(),
339 FlagNode {
340 suggestions_by_provider: BTreeMap::from([
341 ("vmware".to_string(), vec![SuggestionEntry::from("rhel")]),
342 ("nrec".to_string(), vec![SuggestionEntry::from("alma")]),
343 ]),
344 suggestions: vec![SuggestionEntry::from("rhel"), SuggestionEntry::from("alma")],
345 ..FlagNode::default()
346 },
347 );
348 let mut orch = CompletionNode::default();
349 orch.children.insert("provision".to_string(), provision);
350
351 let mut hidden = CompletionNode::default();
352 hidden.flags.insert(
353 "--provider".to_string(),
354 FlagNode {
355 suggestions: vec![
356 SuggestionEntry::from("vmware"),
357 SuggestionEntry::from("nrec"),
358 ],
359 context_only: true,
360 context_scope: ContextScope::Global,
361 ..FlagNode::default()
362 },
363 );
364
365 let tree = CompletionTree {
366 root: CompletionNode::default()
367 .with_child("orch", orch)
368 .with_child("hidden", hidden),
369 ..CompletionTree::default()
370 };
371 let engine = CompletionEngine::new(tree);
372
373 let line = "orch provision --os --provider vmware";
374 let cursor = line.find("--provider").expect("provider in test line") - 1;
375 let (_, suggestions) = engine.complete(line, cursor);
376 let values: Vec<String> = suggestions
377 .into_iter()
378 .filter_map(|entry| match entry {
379 SuggestionOutput::Item(item) => Some(item.text),
380 SuggestionOutput::PathSentinel => None,
381 })
382 .collect();
383 assert!(values.contains(&"rhel".to_string()));
384 assert!(!values.contains(&"alma".to_string()));
385 }
386
387 #[test]
388 fn subtree_context_flags_do_not_leak_across_branches() {
389 let mut provision = CompletionNode::default();
390 provision.flags.insert(
391 "--os".to_string(),
392 FlagNode {
393 suggestions_by_provider: BTreeMap::from([
394 ("vmware".to_string(), vec![SuggestionEntry::from("rhel")]),
395 ("nrec".to_string(), vec![SuggestionEntry::from("alma")]),
396 ]),
397 suggestions: vec![SuggestionEntry::from("rhel"), SuggestionEntry::from("alma")],
398 ..FlagNode::default()
399 },
400 );
401 let mut orch = CompletionNode::default();
402 orch.children.insert("provision".to_string(), provision);
403
404 let mut hidden = CompletionNode::default();
405 hidden.flags.insert(
406 "--provider".to_string(),
407 FlagNode {
408 suggestions: vec![
409 SuggestionEntry::from("vmware"),
410 SuggestionEntry::from("nrec"),
411 ],
412 context_only: true,
413 context_scope: ContextScope::Subtree,
414 ..FlagNode::default()
415 },
416 );
417
418 let tree = CompletionTree {
419 root: CompletionNode::default()
420 .with_child("orch", orch)
421 .with_child("hidden", hidden),
422 ..CompletionTree::default()
423 };
424 let engine = CompletionEngine::new(tree);
425
426 let line = "orch provision --os --provider vmware";
427 let cursor = line.find("--provider").expect("provider in test line") - 1;
428 let (_, suggestions) = engine.complete(line, cursor);
429 let values: Vec<String> = suggestions
430 .into_iter()
431 .filter_map(|entry| match entry {
432 SuggestionOutput::Item(item) => Some(item.text),
433 SuggestionOutput::PathSentinel => None,
434 })
435 .collect();
436 assert!(values.contains(&"rhel".to_string()));
437 assert!(values.contains(&"alma".to_string()));
438 }
439
440 #[test]
441 fn terminal_command_without_flags_does_not_inherit_root_flags() {
442 let mut root = CompletionNode::default();
443 root.flags
444 .insert("--json".to_string(), FlagNode::default().flag_only());
445 root.children
446 .insert("exit".to_string(), CompletionNode::default());
447 let engine = CompletionEngine::new(CompletionTree {
448 root,
449 ..CompletionTree::default()
450 });
451
452 let analysis = engine.analyze("exit ", 5);
453 assert_eq!(analysis.parsed.cursor_tokens, vec!["exit".to_string()]);
454 assert_eq!(analysis.parsed.cursor_cmd.head(), &["exit".to_string()]);
455 assert_eq!(analysis.context.matched_path, vec!["exit".to_string()]);
456 assert_eq!(analysis.context.flag_scope_path, vec!["exit".to_string()]);
457
458 let suggestions = engine.suggestions_for_analysis(&analysis);
459 assert!(
460 suggestions.is_empty(),
461 "expected no inherited flags, got {suggestions:?}"
462 );
463 }
464
465 #[test]
466 fn subcommand_meta_includes_tooltip_and_preview() {
467 let mut ldap = CompletionNode {
468 tooltip: Some("Directory lookup".to_string()),
469 ..CompletionNode::default()
470 };
471 ldap.children
472 .insert("user".to_string(), CompletionNode::default());
473 ldap.children
474 .insert("host".to_string(), CompletionNode::default());
475
476 let tree = CompletionTree {
477 root: CompletionNode::default().with_child("ldap", ldap),
478 ..CompletionTree::default()
479 };
480 let engine = CompletionEngine::new(tree);
481
482 let (_, suggestions) = engine.complete("ld", 2);
483 let meta = suggestions
484 .into_iter()
485 .find_map(|entry| match entry {
486 SuggestionOutput::Item(item) if item.text == "ldap" => item.meta,
487 SuggestionOutput::PathSentinel => None,
488 _ => None,
489 })
490 .expect("ldap suggestion should have metadata");
491
492 assert!(meta.contains("Directory lookup"));
493 assert!(meta.contains("subcommands:"));
494 assert!(meta.contains("host"));
495 assert!(meta.contains("user"));
496 }
497
498 #[test]
499 fn analyze_exposes_merged_cursor_context() {
500 let engine = CompletionEngine::new(tree());
501 let line = "orch provision --os --provider vmware";
502 let cursor = line.find("--provider").expect("provider in test line") - 1;
503
504 let analysis = engine.analyze(line, cursor);
505
506 assert_eq!(analysis.cursor.token_stub, "");
507 assert_eq!(analysis.context.matched_path, vec!["orch", "provision"]);
508 assert_eq!(analysis.context.flag_scope_path, vec!["orch", "provision"]);
509 assert!(!analysis.context.subcommand_context);
510 assert_eq!(
511 analysis
512 .parsed
513 .cursor_cmd
514 .flag_values("--provider")
515 .expect("provider should merge into cursor context"),
516 &vec!["vmware".to_string()][..]
517 );
518 }
519
520 #[test]
521 fn analyze_preserves_open_quote_context() {
522 let engine = CompletionEngine::new(tree());
523 let line = "orch provision --os \"rh";
524
525 let analysis = engine.analyze(line, line.len());
526
527 assert_eq!(analysis.cursor.token_stub, "rh");
528 assert_eq!(analysis.cursor.quote_style, Some(QuoteStyle::Double));
529 }
530
531 #[test]
532 fn matched_command_len_counts_value_keys_consistently() {
533 let mut set = CompletionNode::default();
534 set.children.insert(
535 "ui.mode".to_string(),
536 CompletionNode {
537 value_key: true,
538 ..CompletionNode::default()
539 },
540 );
541 let mut config = CompletionNode::default();
542 config.children.insert("set".to_string(), set);
543 let tree = CompletionTree {
544 root: CompletionNode::default().with_child("config", config),
545 ..CompletionTree::default()
546 };
547 let engine = CompletionEngine::new(tree);
548
549 let tokens = vec![
550 "config".to_string(),
551 "set".to_string(),
552 "ui.mode".to_string(),
553 ];
554 assert_eq!(engine.matched_command_len_tokens(&tokens), 3);
555 }
556}