Implemented completion for variable names

Fixed 'w' and 'b' motions in vi mode stopping on underscores

src/prompt/readline/complete.rs

@@ -1,6 +1,6 @@
 use std::{env, os::unix::fs::PermissionsExt, path::PathBuf, sync::Arc};
 
-use crate::{builtin::BUILTINS, libsh::error::ShResult, parse::lex::{self, LexFlags, Tk, TkFlags}, prompt::readline::{annotate_input, annotate_input_recursive, get_insertions, markers::{self, is_marker}}, state::read_logic};
+use crate::{builtin::BUILTINS, libsh::error::{ShErr, ShErrKind, ShResult}, parse::lex::{self, LexFlags, Tk, TkFlags}, prompt::readline::{Marker, annotate_input, annotate_input_recursive, get_insertions, markers::{self, is_marker}}, state::{read_logic, read_vars}};
 
 pub enum CompCtx {
 	CmdName,
@@ -53,11 +53,11 @@ impl Completer {
 		(before_cursor, after_cursor)
 	}
 
-	pub fn get_completion_context(&self, line: &str, cursor_pos: usize) -> (bool, usize) {
+	pub fn get_completion_context(&self, line: &str, cursor_pos: usize) -> (Vec<Marker>, usize) {
 		let annotated = annotate_input_recursive(line);
 		log::debug!("Annotated input for completion context: {:?}", annotated);
-		let mut in_cmd = false;
+		let mut ctx = vec![markers::NULL];
-		let mut same_position = false; // so that arg markers do not overwrite command markers if they are in the same spot
+		let mut last_priority = 0;
 		let mut ctx_start = 0;
 		let mut pos = 0;
 
@@ -67,32 +67,47 @@ impl Completer {
 					match ch {
 						markers::COMMAND | markers::BUILTIN => {
 							log::debug!("Found command marker at position {}", pos);
+							if last_priority < 2 {
+								if last_priority > 0 {
+									ctx.pop();
+								}
 								ctx_start = pos;
-							same_position = true;
+								last_priority = 2;
-							in_cmd = true;
+								ctx.push(markers::COMMAND);
+							}
+						}
+						markers::VAR_SUB => {
+							log::debug!("Found variable substitution marker at position {}", pos);
+							if last_priority < 3 {
+								if last_priority > 0 {
+									ctx.pop();
+								}
+								ctx_start = pos;
+								last_priority = 3;
+								ctx.push(markers::VAR_SUB);
+							}
 						}
 						markers::ARG => {
 							log::debug!("Found argument marker at position {}", pos);
-							if !same_position {
+							if last_priority < 1 {
 								ctx_start = pos;
-								in_cmd = false;
+								ctx.push(markers::ARG);
 							}
 						}
 						_ => {}
 					}
 				}
 				_ => {
-					same_position = false;
+					last_priority = 0; // reset priority on normal characters
 					pos += 1; // we hit a normal character, advance our position
 					if pos >= cursor_pos {
-						log::debug!("Cursor is at position {}, current context: {}", pos, if in_cmd { "command" } else { "argument" });
 						break;
 					}
 				}
 			}
 		}
 
-		(in_cmd, ctx_start)
+		(ctx, ctx_start)
 	}
 
 	pub fn reset(&mut self) {
@@ -150,6 +165,54 @@ impl Completer {
 		}
 	}
 
+	pub fn extract_var_name(text: &str) -> Option<(String,usize,usize)> {
+		let mut chars = text.chars().peekable();
+		let mut name = String::new();
+		let mut reading_name = false;
+		let mut pos = 0;
+		let mut name_start = 0;
+		let mut name_end = 0;
+
+		while let Some(ch) = chars.next() {
+			match ch {
+				'$' => {
+					if chars.peek() == Some(&'{') {
+						continue;
+					}
+
+					reading_name = true;
+					name_start = pos + 1; // Start after the '$'
+				}
+				'{' if !reading_name => {
+					reading_name = true;
+					name_start = pos + 1;
+				}
+				ch if ch.is_alphanumeric() || ch == '_' => {
+					if reading_name {
+						name.push(ch);
+					}
+				}
+				_ => {
+					if reading_name {
+						name_end = pos; // End before the non-alphanumeric character
+						break;
+					}
+				}
+			}
+			pos += 1;
+		}
+
+		if !reading_name {
+			return None;
+		}
+
+		if name_end == 0 {
+			name_end = pos;
+		}
+
+		Some((name, name_start, name_end))
+	}
+
 	pub fn get_completed_line(&self) -> String {
 		if self.candidates.is_empty() {
 			return self.original_input.clone();
@@ -180,9 +243,8 @@ impl Completer {
 
 
 		// Look for marker at the START of what we're completing, not at cursor
-		let (is_cmd, token_start) = self.get_completion_context(&line, cursor_pos);
+		let (mut ctx, token_start) = self.get_completion_context(&line, cursor_pos);
 		self.token_span.0 = token_start; // Update start of token span based on context
-		log::debug!("Completion context: {}, token span: {:?}, token_start: {}", if is_cmd { "command" } else { "argument" }, self.token_span, token_start);
 		cur_token.span.set_range(self.token_span.0..self.token_span.1); // Update token span to reflect context
 
 		// If token contains '=', only complete after the '='
@@ -192,18 +254,81 @@ impl Completer {
 			self.token_span.0 = cur_token.span.start + eq_pos + 1;
 		}
 
+		if ctx.last().is_some_and(|m| *m == markers::VAR_SUB) {
+			let var_sub = &cur_token.as_str();
+			if let Some((var_name,start,end)) = Self::extract_var_name(var_sub) {
+				log::debug!("Extracted variable name for completion: {}", var_name);
+				if read_vars(|v| v.get_var(&var_name)).is_empty() {
+					// if we are here, we have a variable substitution that isn't complete
+					// so let's try to complete it
+					let ret: ShResult<CompResult> = read_vars(|v| {
+						let var_matches = v.flatten_vars()
+							.keys()
+							.filter(|k| k.starts_with(&var_name) && *k != &var_name)
+							.map(|k| k.to_string())
+							.collect::<Vec<_>>();
+
+						if !var_matches.is_empty() {
+							let name_start = cur_token.span.start + start;
+							let name_end = cur_token.span.start + end;
+							self.token_span = (name_start, name_end);
+							cur_token.span.set_range(self.token_span.0..self.token_span.1);
+							Ok(CompResult::from_candidates(var_matches))
+						} else {
+							Ok(CompResult::NoMatch)
+						}
+					});
+
+					if !matches!(ret, Ok(CompResult::NoMatch)) {
+						return ret;
+					} else {
+						ctx.pop();
+					}
+				} else {
+					ctx.pop();
+				}
+			}
+		}
+
+		let raw_tk = cur_token.as_str().to_string();
 		let expanded_tk = cur_token.expand()?;
 		let expanded_words = expanded_tk.get_words().into_iter().collect::<Vec<_>>();
 		let expanded = expanded_words.join("\\ ");
 
-		let candidates = if is_cmd {
+		let mut candidates = match ctx.pop() {
+			Some(markers::COMMAND) => {
 				log::debug!("Completing command: {}", &expanded);
 				Self::complete_command(&expanded)?
-		} else {
+			}
+			Some(markers::ARG) => {
 				log::debug!("Completing filename: {}", &expanded);
 				Self::complete_filename(&expanded)
+			}
+			Some(m) => {
+				log::warn!("Unknown marker {:?} in completion context", m);
+				return Ok(CompResult::NoMatch);
+			}
+			None => {
+				log::warn!("No marker found in completion context");
+				return Ok(CompResult::NoMatch);
+			}
 		};
 
+		// Now we are just going to graft the completed text
+		// onto the original token. This prevents something like
+		// $SOME_PATH/
+		// from being completed into
+		// /path/to/some_path/file.txt
+		// and instead returns
+		// $SOME_PATH/file.txt
+		candidates = candidates.into_iter()
+			.map(|c| match c.strip_prefix(&expanded) {
+				Some(suffix) => format!("{raw_tk}{suffix}"),
+				None => c
+			})
+			.collect();
+
+
 		Ok(CompResult::from_candidates(candidates))
 	}
 

src/prompt/readline/linebuf.rs

@@ -53,7 +53,7 @@ impl From<&str> for CharClass {
       return CharClass::Alphanum;
     }
 
-    if value.chars().all(char::is_alphanumeric) {
+    if value.chars().all(|c| c.is_alphanumeric() || c == '_') {
       CharClass::Alphanum
     } else if value.chars().all(char::is_whitespace) {
       CharClass::Whitespace

src/prompt/readline/mod.rs

@@ -24,37 +24,41 @@ pub mod highlight;
 pub mod complete;
 
 pub mod markers {
+	use super::Marker;
+
 	// token-level (derived from token class)
-	pub const COMMAND: char = '\u{fdd0}';
+	pub const COMMAND: Marker = '\u{fdd0}';
-	pub const BUILTIN: char = '\u{fdd1}';
+	pub const BUILTIN: Marker = '\u{fdd1}';
-	pub const ARG: char = '\u{fdd2}';
+	pub const ARG: Marker = '\u{fdd2}';
-	pub const KEYWORD: char = '\u{fdd3}';
+	pub const KEYWORD: Marker = '\u{fdd3}';
-	pub const OPERATOR: char = '\u{fdd4}';
+	pub const OPERATOR: Marker = '\u{fdd4}';
-	pub const REDIRECT: char = '\u{fdd5}';
+	pub const REDIRECT: Marker = '\u{fdd5}';
-	pub const COMMENT: char = '\u{fdd6}';
+	pub const COMMENT: Marker = '\u{fdd6}';
-	pub const ASSIGNMENT: char = '\u{fdd7}';
+	pub const ASSIGNMENT: Marker = '\u{fdd7}';
-	pub const CMD_SEP: char = '\u{fde0}';
+	pub const CMD_SEP: Marker = '\u{fde0}';
-	pub const CASE_PAT: char = '\u{fde1}';
+	pub const CASE_PAT: Marker = '\u{fde1}';
-	pub const SUBSH: char = '\u{fde7}';
+	pub const SUBSH: Marker = '\u{fde7}';
-	pub const SUBSH_END: char = '\u{fde8}';
+	pub const SUBSH_END: Marker = '\u{fde8}';
 
 	// sub-token (needs scanning)
-	pub const VAR_SUB: char = '\u{fdda}';
+	pub const VAR_SUB: Marker = '\u{fdda}';
-	pub const VAR_SUB_END: char = '\u{fde3}';
+	pub const VAR_SUB_END: Marker = '\u{fde3}';
-	pub const CMD_SUB: char = '\u{fdd8}';
+	pub const CMD_SUB: Marker = '\u{fdd8}';
-	pub const CMD_SUB_END: char = '\u{fde4}';
+	pub const CMD_SUB_END: Marker = '\u{fde4}';
-	pub const PROC_SUB: char = '\u{fdd9}';
+	pub const PROC_SUB: Marker = '\u{fdd9}';
-	pub const PROC_SUB_END: char = '\u{fde9}';
+	pub const PROC_SUB_END: Marker = '\u{fde9}';
-	pub const STRING_DQ: char = '\u{fddb}';
+	pub const STRING_DQ: Marker = '\u{fddb}';
-	pub const STRING_DQ_END: char = '\u{fde5}';
+	pub const STRING_DQ_END: Marker = '\u{fde5}';
-	pub const STRING_SQ: char = '\u{fddc}';
+	pub const STRING_SQ: Marker = '\u{fddc}';
-	pub const STRING_SQ_END: char = '\u{fde6}';
+	pub const STRING_SQ_END: Marker = '\u{fde6}';
-	pub const ESCAPE: char = '\u{fddd}';
+	pub const ESCAPE: Marker = '\u{fddd}';
-	pub const GLOB: char = '\u{fdde}';
+	pub const GLOB: Marker = '\u{fdde}';
 
-	pub const RESET: char = '\u{fde2}';
+	pub const RESET: Marker = '\u{fde2}';
 
-	pub const END_MARKERS: [char;7] = [
+	pub const NULL: Marker = '\u{fdef}';
+
+	pub const END_MARKERS: [Marker;7] = [
 		VAR_SUB_END,
 		CMD_SUB_END,
 		PROC_SUB_END,
@@ -63,7 +67,7 @@ pub mod markers {
 		SUBSH_END,
 		RESET
 	];
-	pub const TOKEN_LEVEL: [char;10] = [
+	pub const TOKEN_LEVEL: [Marker;10] = [
 		SUBSH,
 		COMMAND,
 		BUILTIN,
@@ -75,7 +79,7 @@ pub mod markers {
 		CASE_PAT,
 		ASSIGNMENT,
 	];
-	pub const SUB_TOKEN: [char;6] = [
+	pub const SUB_TOKEN: [Marker;6] = [
 		VAR_SUB,
 		CMD_SUB,
 		PROC_SUB,
@@ -84,10 +88,11 @@ pub mod markers {
 		GLOB,
 	];
 
-	pub fn is_marker(c: char) -> bool {
+	pub fn is_marker(c: Marker) -> bool {
 		TOKEN_LEVEL.contains(&c) || SUB_TOKEN.contains(&c) || END_MARKERS.contains(&c)
 	}
 }
+type Marker = char;
 
 /// Non-blocking readline result
 pub enum ReadlineEvent {
@@ -639,7 +644,7 @@ pub fn annotate_input_recursive(input: &str) -> String {
 	annotated
 }
 
-pub fn get_insertions(input: &str) -> Vec<(usize, char)> {
+pub fn get_insertions(input: &str) -> Vec<(usize, Marker)> {
 	let input = Arc::new(input.to_string());
 	let tokens: Vec<Tk> = lex::LexStream::new(input, LexFlags::LEX_UNFINISHED)
 		.flatten()
@@ -662,7 +667,7 @@ pub fn get_insertions(input: &str) -> Vec<(usize, char)> {
 /// - String tokens (which need sub-token scanning for variables, quotes, etc.)
 /// - Structural markers (SOI, EOI, Null)
 /// - Unimplemented features (comments, brace groups)
-pub fn marker_for(class: &TkRule) -> Option<char> {
+pub fn marker_for(class: &TkRule) -> Option<Marker> {
 	match class {
 		TkRule::Pipe |
 		TkRule::ErrPipe |
@@ -683,17 +688,17 @@ pub fn marker_for(class: &TkRule) -> Option<char> {
 	}
 }
 
-pub fn annotate_token(token: Tk) -> Vec<(usize, char)> {
+pub fn annotate_token(token: Tk) -> Vec<(usize, Marker)> {
 	// Sort by position descending, with priority ordering at same position:
 	// - RESET first (inserted first, ends up rightmost)
 	// - Regular markers middle
 	// - END markers last (inserted last, ends up leftmost)
 	// Result: [END][TOGGLE][RESET]
-	let sort_insertions = |insertions: &mut Vec<(usize, char)>| {
+	let sort_insertions = |insertions: &mut Vec<(usize, Marker)>| {
 		insertions.sort_by(|a, b| {
 			match b.0.cmp(&a.0) {
 				std::cmp::Ordering::Equal => {
-					let priority = |m: char| -> u8 {
+					let priority = |m: Marker| -> u8 {
 						match m {
 							markers::RESET => 0,
 							markers::VAR_SUB |
@@ -718,12 +723,12 @@ pub fn annotate_token(token: Tk) -> Vec<(usize, char)> {
 		});
 	};
 
-	let in_context = |c: char, insertions: &[(usize, char)]| -> bool {
+	let in_context = |c: Marker, insertions: &[(usize, Marker)]| -> bool {
 		let mut stack = insertions.to_vec();
 		stack.sort_by(|a, b| {
 			match b.0.cmp(&a.0) {
 				std::cmp::Ordering::Equal => {
-					let priority = |m: char| -> u8 {
+					let priority = |m: Marker| -> u8 {
 						match m {
 							markers::RESET => 0,
 							markers::VAR_SUB |
@@ -755,7 +760,7 @@ pub fn annotate_token(token: Tk) -> Vec<(usize, char)> {
 		ctx.1 == c
 	};
 
-	let mut insertions: Vec<(usize, char)> = vec![];
+	let mut insertions: Vec<(usize, Marker)> = vec![];
 
 
 	if token.class != TkRule::Str