use std::borrow::Cow;
use std::io::{stdout, IsTerminal, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::{collections::BTreeMap, env::var_os};
use clap::{parser::ValueSource, Arg, ArgAction, Command};
use rustyline::DefaultEditor;
use luallaby::{read_lua_file, Error, LuaError, Value, LUA_VERSION, VM};
macro_rules! ok {
( $res:expr ) => {
if let Err(e) = $res {
eprintln!("{}: {}", std::env::args().next().expect("no argv[0]"), e);
std::process::exit(1);
}
};
}
enum Action {
Execute(String),
Load(String),
}
fn main() -> Result<(), Error> {
let interrupt = Arc::new(AtomicBool::new(false));
let i = interrupt.clone();
let mut double = false;
ctrlc::set_handler(move || {
if double {
std::process::exit(130);
}
i.store(true, Ordering::Release);
double = true;
})
.expect("could not set SIGINT handler");
let args = cli().get_matches();
if args.get_flag(ARG_VERSION) {
println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
}
let mut vm = VM::default();
vm.set_interrupt(interrupt);
vm.set_warn(args.get_flag(ARG_WARNINGS));
vm.set_args(args.get_one::<String>(ARG_SCRIPT).map(String::as_str));
if !args.get_flag(ARG_NOENV) {
vm.load_env();
handle_init(&mut vm);
}
let mut actions = BTreeMap::new();
for id in args.ids() {
if matches!(id.as_str(), ARG_EXECUTE | ARG_LOAD) {
for (value, index) in args
.get_many::<String>(id.as_str())
.expect("could not get arguments")
.zip(
args.indices_of(id.as_str())
.expect("could not get argument indices"),
)
{
match id.as_str() {
ARG_EXECUTE => actions.insert(index, Action::Execute(value.to_string())),
ARG_LOAD => actions.insert(index, Action::Load(value.to_string())),
_ => unreachable!(),
};
}
}
}
for action in actions.values() {
match action {
Action::Execute(stat) => {
run_stat(&mut vm, stat, Some("(command line)"));
}
Action::Load(str) => {
let (glob, module) = str.split_once('=').unwrap_or((str, str));
ok!(vm.load_lib(glob, module));
}
}
}
if args.get_flag(ARG_INTERACTIVE) {
repl(&mut vm)
} else if args.contains_id(ARG_SCRIPT) {
let filename = args
.get_one::<String>(ARG_SCRIPT)
.expect("could not get arguments");
run_script(&mut vm, filename);
Ok(())
} else if !matches!(
args.value_source(ARG_EXECUTE),
Some(ValueSource::CommandLine)
) {
if std::io::stdout().is_terminal() {
print_version();
repl(&mut vm)
} else {
run_script(&mut vm, "-");
Ok(())
}
} else {
Ok(())
}
}
const ARG_SCRIPT: &str = "script";
const ARG_EXECUTE: &str = "execute";
const ARG_LOAD: &str = "load";
const ARG_INTERACTIVE: &str = "interactive";
const ARG_NOENV: &str = "noenv";
const ARG_WARNINGS: &str = "warnings";
const ARG_VERSION: &str = "version";
fn cli() -> Command {
Command::new(env!("CARGO_PKG_NAME"))
.version(env!("CARGO_PKG_VERSION"))
.disable_version_flag(true)
.disable_help_flag(true)
.arg(
Arg::new(ARG_SCRIPT)
.num_args(..)
.trailing_var_arg(true)
.value_name("FILE")
.help("Path (and arguments) to the Lua script to run"),
)
.arg(
Arg::new(ARG_EXECUTE)
.short('e')
.long("execute")
.value_name("STATEMENT")
.action(ArgAction::Append)
.help("Execute string statement"),
)
.arg(
Arg::new(ARG_LOAD)
.short('l')
.long("load")
.value_name("MODULE")
.action(ArgAction::Append)
.help("Load a module"),
)
.arg(
Arg::new(ARG_INTERACTIVE)
.short('i')
.long("interactive")
.action(ArgAction::SetTrue)
.help("Run in interactive mode"),
)
.arg(
Arg::new(ARG_NOENV)
.short('E')
.long("no-env")
.action(ArgAction::SetTrue)
.help("Ignore environment variables"),
)
.arg(
Arg::new(ARG_WARNINGS)
.short('W')
.long("warnings")
.action(ArgAction::SetTrue)
.help("Enable warnings"),
)
.arg(
Arg::new(ARG_VERSION)
.short('v')
.long("version")
.action(ArgAction::SetTrue)
.help("Print version"),
)
.arg(
Arg::new("help")
.short('h')
.long("help")
.action(ArgAction::Help)
.help("Print help"),
)
}
fn print_version() {
println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
}
fn handle_init(vm: &mut VM) {
let (major, minor) = LUA_VERSION;
let env = var_os(format!("LUA_INIT_{}_{}", major, minor)).or(var_os("LUA_INIT"));
if let Some(init) = env {
let init = init.to_string_lossy();
let init = init.as_ref();
if matches!(init.as_bytes().get(0), Some(b'@')) {
run_script(vm, &init[1..]);
} else {
run_stat(vm, init, Some("LUA_INIT"));
}
}
}
fn run_stat(vm: &mut VM, stat: &str, source: Option<&str>) {
ok!(vm.run_str(stat.as_bytes(), source))
}
fn run_script(vm: &mut VM, filename: &str) {
let input = read_lua_file(filename).expect(&format!("unable to read file '{}'", filename));
let source = if filename == "-" { "stdin" } else { filename };
ok!(vm.run_file(&input, Some(source)));
}
fn repl(vm: &mut VM) -> Result<(), Error> {
let mut code = String::new();
let mut rl = DefaultEditor::new()?;
loop {
let default_prompt = if code.is_empty() { "> " } else { ">> " };
let global = vm.get_global(if code.is_empty() {
"_PROMPT"
} else {
"_PROMPT2"
});
let prompt = match global {
Value::Nil => Cow::from(default_prompt),
val => {
let str = vm.to_string(val)?;
Cow::from(String::from_utf8_lossy(&str).into_owned())
}
};
if !stdout().is_terminal() {
stdout().write_all(prompt.as_bytes())?;
}
let line = rl.readline(prompt.as_ref());
match line {
Ok(line) => {
if !stdout().is_terminal() {
stdout().write_all(line.as_bytes())?;
writeln!(stdout())?;
}
code.push_str(&line);
code.push('\n');
if code.trim().is_empty() {
code.clear();
} else {
match vm.run_str_interactive(code.trim().as_bytes()) {
Err(Error::Lua {
typ: LuaError::UnexpectedEof,
..
}) => continue,
Err(e) => eprintln!("{}", e),
Ok(_) => {}
}
rl.add_history_entry(code.trim())?;
code = String::new();
}
}
Err(..) => break,
}
}
if !stdout().is_terminal() {
stdout().write_all(&[b'\n'])?;
}
Ok(())
}