From 7eadbd938dda5a63cf81ef36b5322df589d86c5b Mon Sep 17 00:00:00 2001
From: JT <jonathandturner@users.noreply.github.com>
Date: Tue, 8 Jun 2021 14:31:39 +1200
Subject: [PATCH] Add support for subcommand completions (#3571)

* Add support for subcommand completions

* Update test

* WIP

* Fix prepend for completions

* Fix test
---
 crates/nu-cli/src/completion/engine.rs | 51 ++++++++++++++++++++++----
 crates/nu-cli/src/shell/completer.rs   |  9 ++++-
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/crates/nu-cli/src/completion/engine.rs b/crates/nu-cli/src/completion/engine.rs
index f99a6c1fab..7bece6993a 100644
--- a/crates/nu-cli/src/completion/engine.rs
+++ b/crates/nu-cli/src/completion/engine.rs
@@ -203,7 +203,8 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
     } else {
         let mut command = None;
         let mut prev = None;
-        for loc in locations {
+
+        for loc in &locations {
             // We don't use span.contains because we want to include the end. This handles the case
             // where the cursor is just after the text (i.e., no space between cursor and text)
             if loc.span.start() <= pos && pos <= loc.span.end() {
@@ -215,14 +216,31 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
                             let cmd = cmd.clone();
                             let span = loc.span;
                             vec![
-                                loc,
+                                loc.clone(),
                                 LocationType::Flag(cmd.unwrap_or_default()).spanned(span),
                             ]
                         } else {
-                            vec![loc]
+                            let mut output = vec![];
+
+                            for rloc in locations.iter().rev() {
+                                if let Spanned {
+                                    span,
+                                    item: LocationType::Command,
+                                } = &rloc
+                                {
+                                    output.push(LocationType::Command.spanned(Span::new(
+                                        span.start(),
+                                        locations[locations.len() - 1].span.end(),
+                                    )));
+                                    break;
+                                }
+                            }
+
+                            output.push(loc.clone());
+                            output
                         }
                     }
-                    _ => vec![loc],
+                    _ => vec![loc.clone()],
                 };
             } else if pos < loc.span.start() {
                 break;
@@ -236,14 +254,25 @@ pub fn completion_location(line: &str, block: &Block, pos: usize) -> Vec<Complet
         }
 
         if let Some(prev) = prev {
+            let mut locations = vec![];
             // Cursor is between locations (or at the end). Look at the line to see if the cursor
             // is after some character that would imply we're in the command position.
             let start = prev.span.end();
+
+            if let Spanned {
+                item: LocationType::Command,
+                span,
+            } = &prev
+            {
+                locations.push(LocationType::Command.spanned(Span::new(span.start(), pos)));
+            }
             if line[start..pos].contains(BEFORE_COMMAND_CHARS) {
-                vec![LocationType::Command.spanned(Span::new(pos, pos))]
+                locations.push(LocationType::Command.spanned(Span::new(pos, pos)));
+                locations
             } else {
                 // TODO this should be able to be mapped to a command
-                vec![LocationType::Argument(command, None).spanned(Span::new(pos, pos))]
+                locations.push(LocationType::Argument(command, None).spanned(Span::new(pos, pos)));
+                locations
             }
         } else {
             // Cursor is before any possible completion location, so must be a command
@@ -399,7 +428,10 @@ mod tests {
 
             assert_eq!(
                 completion_location(line, &registry, 3),
-                vec![LocationType::Argument(Some("cd".to_string()), None)],
+                vec![
+                    LocationType::Command,
+                    LocationType::Argument(Some("cd".to_string()), None)
+                ],
             );
         }
 
@@ -429,7 +461,10 @@ mod tests {
 
             assert_eq!(
                 completion_location(line, &registry, 6),
-                vec![LocationType::Argument(Some("echo".to_string()), None)],
+                vec![
+                    LocationType::Command,
+                    LocationType::Argument(Some("echo".to_string()), None)
+                ],
             );
         }
     }
diff --git a/crates/nu-cli/src/shell/completer.rs b/crates/nu-cli/src/shell/completer.rs
index 085df33df5..e7d70ae4b9 100644
--- a/crates/nu-cli/src/shell/completer.rs
+++ b/crates/nu-cli/src/shell/completer.rs
@@ -6,7 +6,7 @@ use crate::completion::path::{PathCompleter, PathSuggestion};
 use crate::completion::{self, Completer, Suggestion};
 use nu_engine::EvaluationContext;
 use nu_parser::ParserScope;
-use nu_source::Tag;
+use nu_source::{Span, Tag};
 
 use std::borrow::Cow;
 
@@ -72,6 +72,7 @@ impl NuCompleter {
 
                         LocationType::Argument(cmd, _arg_name) => {
                             let path_completer = PathCompleter;
+                            let prepend = Span::new(pos, location.span.start()).slice(line);
 
                             const QUOTE_CHARS: &[char] = &['\'', '"', '`'];
 
@@ -103,7 +104,11 @@ impl NuCompleter {
                             }
                             .into_iter()
                             .map(|s| Suggestion {
-                                replacement: requote(s.suggestion.replacement),
+                                replacement: format!(
+                                    "{}{}",
+                                    prepend,
+                                    requote(s.suggestion.replacement)
+                                ),
                                 display: s.suggestion.display,
                             })
                             .collect()