Added prompt escape code expansion flag to echo, -p

Added non-formatted runtime to prompt escape codes

Added prompt escape code that expands to the output of a shell function

Reworked internal logic for termios control

src/builtin/echo.rs

@@ -1,14 +1,7 @@
 use std::sync::LazyLock;
 
 use crate::{
-  builtin::setup_builtin,
+  builtin::setup_builtin, expand::expand_prompt, getopt::{Opt, OptSet, get_opts_from_tokens}, jobs::JobBldr, libsh::error::{ShErr, ShErrKind, ShResult, ShResultExt}, parse::{NdRule, Node}, prelude::*, procio::{IoStack, borrow_fd}, state
-  getopt::{get_opts_from_tokens, Opt, OptSet},
-  jobs::JobBldr,
-  libsh::error::{ShErr, ShErrKind, ShResult, ShResultExt},
-  parse::{NdRule, Node},
-  prelude::*,
-  procio::{borrow_fd, IoStack},
-  state,
 };
 
 pub static ECHO_OPTS: LazyLock<OptSet> = LazyLock::new(|| {
@@ -16,7 +9,7 @@ pub static ECHO_OPTS: LazyLock<OptSet> = LazyLock::new(|| {
     Opt::Short('n'),
     Opt::Short('E'),
     Opt::Short('e'),
-    Opt::Short('r'),
+    Opt::Short('p'),
   ]
   .into()
 });
@@ -26,6 +19,7 @@ bitflags! {
     const NO_NEWLINE = 0b000001;
     const USE_STDERR = 0b000010;
     const USE_ESCAPE = 0b000100;
+		const USE_PROMPT = 0b001000;
   }
 }
 
@@ -49,11 +43,13 @@ pub fn echo(node: Node, io_stack: &mut IoStack, job: &mut JobBldr) -> ShResult<(
     borrow_fd(STDOUT_FILENO)
   };
 
-  let mut echo_output = argv
+  let mut echo_output = prepare_echo_args(argv
     .into_iter()
     .map(|a| a.0) // Extract the String from the tuple of (String,Span)
-    .collect::<Vec<_>>()
+    .collect::<Vec<_>>(),
-    .join(" ");
+		flags.contains(EchoFlags::USE_ESCAPE),
+		flags.contains(EchoFlags::USE_PROMPT)
+	)?.join(" ");
 
   if !flags.contains(EchoFlags::NO_NEWLINE) {
     echo_output.push('\n')
@@ -66,6 +62,122 @@ pub fn echo(node: Node, io_stack: &mut IoStack, job: &mut JobBldr) -> ShResult<(
   Ok(())
 }
 
+pub fn prepare_echo_args(argv: Vec<String>, use_escape: bool, use_prompt: bool) -> ShResult<Vec<String>> {
+	if !use_escape {
+		if use_prompt {
+			let expanded: ShResult<Vec<String>> = argv
+				.into_iter()
+				.map(|s| expand_prompt(s.as_str()))
+				.collect();
+			return expanded
+		}
+		return Ok(argv);
+	}
+
+	let mut prepared_args = Vec::with_capacity(argv.len());
+
+	for arg in argv {
+		let mut prepared_arg = String::new();
+		if use_prompt {
+			prepared_arg = expand_prompt(&prepared_arg)?;
+		}
+
+		let mut chars = arg.chars().peekable();
+
+		while let Some(c) = chars.next() {
+			if c == '\\' {
+				if let Some(&next_char) = chars.peek() {
+					match next_char {
+						'n' => {
+							prepared_arg.push('\n');
+							chars.next();
+						}
+						't' => {
+							prepared_arg.push('\t');
+							chars.next();
+						}
+						'r' => {
+							prepared_arg.push('\r');
+							chars.next();
+						}
+						'a' => {
+							prepared_arg.push('\x07');
+							chars.next();
+						}
+						'b' => {
+							prepared_arg.push('\x08');
+							chars.next();
+						}
+						'e' | 'E' => {
+							prepared_arg.push('\x1b');
+							chars.next();
+						}
+						'x' => {
+							chars.next(); // consume 'x'
+							let mut hex_digits = String::new();
+							for _ in 0..2 {
+								if let Some(&hex_char) = chars.peek() {
+									if hex_char.is_ascii_hexdigit() {
+										hex_digits.push(hex_char);
+										chars.next();
+									} else {
+										break;
+									}
+								} else {
+									break;
+								}
+							}
+							if let Ok(value) = u8::from_str_radix(&hex_digits, 16) {
+								prepared_arg.push(value as char);
+							} else {
+								prepared_arg.push('\\');
+								prepared_arg.push('x');
+								prepared_arg.push_str(&hex_digits);
+							}
+						}
+						'0' => {
+							chars.next(); // consume '0'
+							let mut octal_digits = String::new();
+							for _ in 0..3 {
+								if let Some(&octal_char) = chars.peek() {
+									if ('0'..='7').contains(&octal_char) {
+										octal_digits.push(octal_char);
+										chars.next();
+									} else {
+										break;
+									}
+								} else {
+									break;
+								}
+							}
+							if let Ok(value) = u8::from_str_radix(&octal_digits, 8) {
+								prepared_arg.push(value as char);
+							} else {
+								prepared_arg.push('\\');
+								prepared_arg.push('0');
+								prepared_arg.push_str(&octal_digits);
+							}
+						}
+						'\\' => {
+							prepared_arg.push('\\');
+							chars.next();
+						}
+						_ => prepared_arg.push(c),
+					}
+				} else {
+					prepared_arg.push(c);
+				}
+			} else {
+				prepared_arg.push(c);
+			}
+		}
+
+		prepared_args.push(prepared_arg);
+	}
+
+	Ok(prepared_args)
+}
+
 pub fn get_echo_flags(mut opts: Vec<Opt>) -> ShResult<EchoFlags> {
   let mut flags = EchoFlags::empty();
 
@@ -82,6 +194,7 @@ pub fn get_echo_flags(mut opts: Vec<Opt>) -> ShResult<EchoFlags> {
       'n' => flags |= EchoFlags::NO_NEWLINE,
       'r' => flags |= EchoFlags::USE_STDERR,
       'e' => flags |= EchoFlags::USE_ESCAPE,
+			'p' => flags |= EchoFlags::USE_PROMPT,
       _ => unreachable!(),
     }
   }

src/expand.rs

@@ -11,7 +11,7 @@ use crate::parse::lex::{is_field_sep, is_hard_sep, LexFlags, LexStream, Tk, TkFl
 use crate::parse::{Redir, RedirType};
 use crate::prelude::*;
 use crate::procio::{IoBuf, IoFrame, IoMode, IoStack};
-use crate::state::{LogTab, VarFlags, read_jobs, read_vars, write_jobs, write_meta, write_vars};
+use crate::state::{LogTab, VarFlags, read_jobs, read_logic, read_vars, write_jobs, write_meta, write_vars};
 
 const PARAMETERS: [char; 7] = ['@', '*', '#', '$', '?', '!', '0'];
 
@@ -789,13 +789,15 @@ pub fn expand_cmd_sub(raw: &str) -> ShResult<String> {
       io_stack.push_frame(cmd_sub_io_frame);
       if let Err(e) = exec_input(raw.to_string(), Some(io_stack)) {
         eprintln!("{e}");
-        exit(1);
+        unsafe { libc::_exit(1) };
       }
-      exit(0);
+      unsafe { libc::_exit(0) };
     }
     ForkResult::Parent { child } => {
       std::mem::drop(cmd_sub_io_frame); // Closes the write pipe
       let status = waitpid(child, Some(WtFlag::WSTOPPED))?;
+      // Reclaim terminal foreground in case child changed it
+      crate::jobs::take_term()?;
       match status {
         WtStat::Exited(_, _) => {
           flog!(DEBUG, "filling buffer");
@@ -1423,9 +1425,11 @@ pub enum PromptTk {
   AsciiOct(i32),
   Text(String),
   AnsiSeq(String),
+	Function(String), // Expands to the output of any defined shell function
   VisGrp,
   UserSeq,
-  Runtime,
+	RuntimeMillis,
+  RuntimeFormatted,
   Weekday,
   Dquote,
   Squote,
@@ -1442,6 +1446,8 @@ pub enum PromptTk {
   SuccessSymbol,
   FailureSymbol,
   JobCount,
+	VisGroupOpen,
+	VisGroupClose,
 }
 
 pub fn format_cmd_runtime(dur: std::time::Duration) -> String {
@@ -1646,10 +1652,46 @@ fn tokenize_prompt(raw: &str) -> Vec<PromptTk> {
             '$' => tokens.push(PromptTk::PromptSymbol),
             'n' => tokens.push(PromptTk::Text("\n".into())),
             'r' => tokens.push(PromptTk::Text("\r".into())),
-            'T' => tokens.push(PromptTk::Runtime),
+            't' => tokens.push(PromptTk::RuntimeMillis),
+            'T' => tokens.push(PromptTk::RuntimeFormatted),
             '\\' => tokens.push(PromptTk::Text("\\".into())),
             '"' => tokens.push(PromptTk::Text("\"".into())),
             '\'' => tokens.push(PromptTk::Text("'".into())),
+						'(' => tokens.push(PromptTk::VisGroupOpen),
+						')' => tokens.push(PromptTk::VisGroupClose),
+						'!' => {
+							let mut func_name = String::new();
+							let is_braced = chars.peek() == Some(&'{');
+							while let Some(ch) = chars.peek() {
+
+								match ch {
+									'}' if is_braced => {
+										chars.next();
+										break;
+									}
+									'A'..='Z' | 'a'..='z' | '0'..='9' | '_' => {
+										func_name.push(*ch);
+										chars.next();
+									}
+									_ => {
+										if is_braced {
+											// Invalid character in braced function name
+											tokens.push(PromptTk::Text(format!("\\!{{{func_name}")));
+											break;
+										} else {
+											// End of unbraced function name
+											let func_exists = read_logic(|l| l.get_func(&func_name).is_some());
+											if func_exists {
+												tokens.push(PromptTk::Function(func_name));
+											} else {
+												tokens.push(PromptTk::Text(format!("\\!{func_name}")));
+											}
+											break;
+										}
+									}
+								}
+							}
+						}
             'e' => {
               if chars.next() == Some('[') {
                 let mut params = String::new();
@@ -1737,7 +1779,13 @@ pub fn expand_prompt(raw: &str) -> ShResult<String> {
 			PromptTk::AsciiOct(_) => todo!(),
 			PromptTk::Text(txt) => result.push_str(&txt),
 			PromptTk::AnsiSeq(params) => result.push_str(&params),
-      PromptTk::Runtime => {
+			PromptTk::RuntimeMillis => {
+				if let Some(runtime) = write_meta(|m| m.stop_timer()) {
+					let runtime_millis = runtime.as_millis().to_string();
+					result.push_str(&runtime_millis);
+				}
+			}
+			PromptTk::RuntimeFormatted => {
 				if let Some(runtime) = write_meta(|m| m.stop_timer()) {
 					let runtime_fmt = format_cmd_runtime(runtime);
 					result.push_str(&runtime_fmt);
@@ -1790,7 +1838,20 @@ pub fn expand_prompt(raw: &str) -> ShResult<String> {
 			PromptTk::SuccessSymbol => todo!(),
 			PromptTk::FailureSymbol => todo!(),
 			PromptTk::JobCount => todo!(),
-      _ => unimplemented!(),
+			PromptTk::Function(f) => {
+				flog!(DEBUG, "Expanding prompt function: {}", f);
+				let output = expand_cmd_sub(&f)?;
+				result.push_str(&output);
+			}
+			PromptTk::VisGrp => todo!(),
+			PromptTk::UserSeq => todo!(),
+			PromptTk::Weekday => todo!(),
+			PromptTk::Dquote => todo!(),
+			PromptTk::Squote => todo!(),
+			PromptTk::Return => todo!(),
+			PromptTk::Newline => todo!(),
+			PromptTk::VisGroupOpen => todo!(),
+			PromptTk::VisGroupClose => todo!(),
 		}
   }
 

src/libsh/sys.rs

@@ -1,6 +1,6 @@
 use termios::{LocalFlags, Termios};
 
-use crate::{prelude::*, state::write_jobs};
+use crate::{prelude::*};
 ///
 /// The previous state of the terminal options.
 ///
@@ -31,64 +31,46 @@ use crate::{prelude::*, state::write_jobs};
 /// lifecycle could lead to undefined behavior.
 pub(crate) static mut SAVED_TERMIOS: Option<Option<Termios>> = None;
 
-pub fn save_termios() {
+#[derive(Debug)]
-  unsafe {
+pub struct TermiosGuard {
-    SAVED_TERMIOS = Some(if isatty(std::io::stdin().as_raw_fd()).unwrap() {
+	saved_termios: Option<Termios>
-      let mut termios = termios::tcgetattr(std::io::stdin()).unwrap();
-      termios.local_flags &= !LocalFlags::ECHOCTL;
-      termios::tcsetattr(
-        std::io::stdin(),
-        nix::sys::termios::SetArg::TCSANOW,
-        &termios,
-      )
-      .unwrap();
-      Some(termios)
-    } else {
-      None
-    });
-  }
 }
-#[allow(static_mut_refs)]
-///Access the saved termios
-///
-///# Safety
-///This function is unsafe because it accesses a public mutable static value.
-/// This function should only ever be called after save_termios() has already
-/// been called.
-pub unsafe fn get_saved_termios() -> Option<Termios> { unsafe {
-  // SAVED_TERMIOS should *only ever* be set once and accessed once
-  // Set at the start of the program, and accessed during the exit of the program
-  // to reset the termios. Do not use this variable anywhere else
-  SAVED_TERMIOS.clone().flatten()
-}}
 
-/// Set termios to not echo control characters, like ^Z for instance
+impl TermiosGuard {
-pub fn set_termios() {
+	pub fn new(new_termios: Termios) -> Self {
+		let mut new = Self { saved_termios: None };
+
 		if isatty(std::io::stdin().as_raw_fd()).unwrap() {
-    let mut termios = termios::tcgetattr(std::io::stdin()).unwrap();
+			let current_termios = termios::tcgetattr(std::io::stdin()).unwrap();
-    termios.local_flags &= !LocalFlags::ECHOCTL;
+			new.saved_termios = Some(current_termios);
+
 			termios::tcsetattr(
 				std::io::stdin(),
 				nix::sys::termios::SetArg::TCSANOW,
-      &termios,
+				&new_termios,
-    )
+			).unwrap();
-    .unwrap();
+		}
+
+		new
 	}
 }
 
-pub fn sh_quit(code: i32) -> ! {
+impl Default for TermiosGuard {
-  write_jobs(|j| {
+	fn default() -> Self {
-    for job in j.jobs_mut().iter_mut().flatten() {
+	  let mut termios = termios::tcgetattr(std::io::stdin()).unwrap();
-      job.killpg(Signal::SIGTERM).ok();
+		termios.local_flags &= !LocalFlags::ECHOCTL;
+		Self::new(termios)
+	}
+}
+
+impl Drop for TermiosGuard {
+	fn drop(&mut self) {
+		if let Some(saved) = &self.saved_termios {
+			termios::tcsetattr(
+				std::io::stdin(),
+				nix::sys::termios::SetArg::TCSANOW,
+				saved,
+			).unwrap();
+		}
 	}
-  });
-  if let Some(termios) = unsafe { get_saved_termios() } {
-    termios::tcsetattr(std::io::stdin(), termios::SetArg::TCSANOW, &termios).unwrap();
-  }
-  if code == 0 {
-    eprintln!("exit");
-  } else {
-    eprintln!("exit {code}");
-  }
-  exit(code);
 }

src/main.rs

@@ -18,11 +18,14 @@ pub mod state;
 #[cfg(test)]
 pub mod tests;
 
+use std::process::ExitCode;
+use std::sync::atomic::Ordering;
+
 use crate::libsh::error::ShErrKind;
-use crate::libsh::sys::{save_termios, set_termios};
+use crate::libsh::sys::TermiosGuard;
 use crate::parse::execute::exec_input;
 use crate::prelude::*;
-use crate::signal::{check_signals, sig_setup, signals_pending};
+use crate::signal::{QUIT_CODE, check_signals, sig_setup, signals_pending};
 use crate::state::source_rc;
 use clap::Parser;
 use shopt::FernEditMode;
@@ -53,12 +56,12 @@ fn kickstart_lazy_evals() {
   read_vars(|_| {});
 }
 
-fn main() {
+fn main() -> ExitCode {
   kickstart_lazy_evals();
   let args = FernArgs::parse();
   if args.version {
     println!("fern {}", env!("CARGO_PKG_VERSION"));
-    return;
+    return ExitCode::SUCCESS;
   }
 
   if let Some(path) = args.script {
@@ -66,17 +69,21 @@ fn main() {
   } else {
     fern_interactive();
   }
+
+	ExitCode::from(QUIT_CODE.load(Ordering::SeqCst) as u8)
 }
 
 fn run_script<P: AsRef<Path>>(path: P, args: Vec<String>) {
   let path = path.as_ref();
   if !path.is_file() {
     eprintln!("fern: Failed to open input file: {}", path.display());
-    exit(1);
+		QUIT_CODE.store(1, Ordering::SeqCst);
+		return;
   }
   let Ok(input) = fs::read_to_string(path) else {
     eprintln!("fern: Failed to read input file: {}", path.display());
-    exit(1);
+		QUIT_CODE.store(1, Ordering::SeqCst);
+		return;
   };
 
   write_vars(|v| v.cur_scope_mut().bpush_arg(path.to_string_lossy().to_string()));
@@ -86,13 +93,19 @@ fn run_script<P: AsRef<Path>>(path: P, args: Vec<String>) {
 
   if let Err(e) = exec_input(input, None) {
     eprintln!("{e}");
-    exit(1);
+		match e.kind() {
+			ShErrKind::CleanExit(code) => {
+				QUIT_CODE.store(*code, Ordering::SeqCst);
+			}
+			_ => {
+				QUIT_CODE.store(1, Ordering::SeqCst);
+			}
+		}
   }
 }
 
 fn fern_interactive() {
-  save_termios();
+	let _termios_guard = TermiosGuard::default(); // sets raw mode, restores termios on drop
-  set_termios();
   sig_setup();
 
   if let Err(e) = source_rc() {
@@ -129,9 +142,10 @@ fn fern_interactive() {
         line
       }
       Err(e) => {
-				if let ShErrKind::ReadlineIntr(partial) = e.kind()  {
+				match e.kind() {
+					ShErrKind::ReadlineIntr(partial) => {
 						// Did we get signaled? Check signal flags
-					// If nothing to worry about, retry the readline
+						// If nothing to worry about, retry the readline with the unfinished input
 						while signals_pending() {
 							if let Err(e) = check_signals() {
 								if let ShErrKind::ClearReadline = e.kind() {
@@ -145,7 +159,12 @@ fn fern_interactive() {
 						}
 						partial_input = partial.to_string();
 						continue;
-				} else {
+					}
+					ShErrKind::CleanExit(code) => {
+						QUIT_CODE.store(*code, Ordering::SeqCst);
+						return;
+					}
+					_ => {
 						eprintln!("{e}");
 						readline_err_count += 1;
 						if readline_err_count == 20 {
@@ -156,10 +175,19 @@ fn fern_interactive() {
 						}
 					}
 				}
+      }
     };
 
     if let Err(e) = exec_input(input, None) {
+			match e.kind() {
+				ShErrKind::CleanExit(code) => {
+					QUIT_CODE.store(*code, Ordering::SeqCst);
+					return;
+				}
+				_ => {
 					eprintln!("{e}");
 				}
 			}
+    }
+  }
 }

src/parse/mod.rs

@@ -1282,11 +1282,30 @@ impl ParseStream {
         assignments.push(assign)
       } else if is_keyword {
         return Ok(None);
+      } else if prefix_tk.class == TkRule::Sep {
+        // Separator ends the prefix section - add it so commit() consumes it
+        node_tks.push(prefix_tk.clone());
+        break;
+      } else {
+        // Other non-prefix token ends the prefix section
+        break;
       }
     }
 
-    if argv.is_empty() && assignments.is_empty() {
+    if argv.is_empty() {
+			if assignments.is_empty() {
 				return Ok(None);
+			} else {
+				// If we have assignments but no command word,
+				// return the assignment-only command without parsing more tokens
+				self.commit(node_tks.len());
+				return Ok(Some(Node {
+					class: NdRule::Command { assignments, argv },
+					tokens: node_tks,
+					flags: NdFlags::empty(),
+					redirs,
+				}));
+			}
     }
 
     while let Some(tk) = tk_iter.next() {

src/prompt/mod.rs

@@ -11,17 +11,19 @@ use crate::{
 /// Initialize the line editor
 fn get_prompt() -> ShResult<String> {
   let Ok(prompt) = env::var("PS1") else {
-    // prompt expands to:
+    // default prompt expands to:
     //
     // username@hostname
     // short/path/to/pwd/
     // $ _
     let default =
-      "\\n\\e[1;0m\\u\\e[1;36m@\\e[1;31m\\h\\n\\e[1;36m\\W\\e[1;32m/\\n\\e[1;32m\\$\\e[0m ";
+      "\\e[0m\\n\\e[1;0m\\u\\e[1;36m@\\e[1;31m\\h\\n\\e[1;36m\\W\\e[1;32m/\\n\\e[1;32m\\$\\e[0m ";
     return expand_prompt(default);
   };
+	let sanitized = format!("\\e[0m{prompt}");
+	flog!(DEBUG, "Using prompt: {}", sanitized.replace("\n", "\\n"));
 
-  expand_prompt(&prompt)
+  expand_prompt(&sanitized)
 }
 
 pub fn readline(edit_mode: FernEditMode, initial: Option<&str>) -> ShResult<String> {

src/prompt/readline/mod.rs

@@ -8,7 +8,6 @@ use vimode::{CmdReplay, ModeReport, ViInsert, ViMode, ViNormal, ViReplace, ViVis
 
 use crate::libsh::{
   error::{ShErr, ShErrKind, ShResult},
-  sys::sh_quit,
   term::{Style, Styled},
 };
 use crate::prelude::*;
@@ -95,7 +94,7 @@ impl Readline for FernVi {
       if cmd.verb().is_some_and(|v| v.1 == Verb::EndOfFile) {
         if self.editor.buffer.is_empty() {
           std::mem::drop(raw_mode_guard);
-          sh_quit(0);
+					return Err(ShErr::simple(ShErrKind::CleanExit(0), "exit"));
         } else {
           self.editor.buffer.clear();
           continue;

src/prompt/readline/term.rs

@@ -687,7 +687,7 @@ impl LineWriter for TermWriter {
     for _ in 0..rows_to_clear {
       self.buffer.push_str("\x1b[2K\x1b[A");
     }
-    self.buffer.push_str("\x1b[2K");
+    self.buffer.push_str("\x1b[2K\r"); // Clear line and return to column 0
     write_all(self.out, self.buffer.as_str())?;
     self.buffer.clear();
     Ok(())

src/signal.rs

@@ -15,8 +15,8 @@ static GOT_SIGTSTP: AtomicBool = AtomicBool::new(false);
 static GOT_SIGCHLD: AtomicBool = AtomicBool::new(false);
 static REAPING_ENABLED: AtomicBool = AtomicBool::new(true);
 
-static SHOULD_QUIT: AtomicBool = AtomicBool::new(false);
+pub static SHOULD_QUIT: AtomicBool = AtomicBool::new(false);
-static QUIT_CODE: AtomicI32 = AtomicI32::new(0);
+pub static QUIT_CODE: AtomicI32 = AtomicI32::new(0);
 
 pub fn signals_pending() -> bool {
 	GOT_SIGINT.load(Ordering::SeqCst)

src/state.rs

@@ -686,7 +686,6 @@ impl MetaTab {
   pub fn stop_timer(&mut self) -> Option<Duration> {
     self
       .runtime_start
-      .take() // runtime_start returns to None
       .map(|start| start.elapsed()) // return the duration, if any
   }
 }