Skip to main content

osp_cli/completion/
engine.rs

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        // Context-only flags can appear later in the line than the cursor.
76        // Merge them before scope resolution so completion reflects the user's
77        // effective command state, not just the prefix before the cursor.
78        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        // When the user is still typing a subcommand, the partial token is the
174        // last `head` element but not yet a resolvable child node. Drop that
175        // partial token so context resolution stays on the parent command.
176        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}