use clap::Parser;
use rsmarisa::{Agent, Trie};
use std::io::{self, Write};
use std::path::PathBuf;
use std::process;
#[derive(Parser)]
#[command(name = "rsmarisa-dump")]
#[command(about = "Dump all keys from a MARISA trie dictionary")]
#[command(version)]
struct Args {
#[arg(short = 'd', long, default_value = "\n")]
delimiter: String,
#[arg(short = 'm', long)]
mmap_dictionary: bool,
#[arg(short = 'r', long)]
read_dictionary: bool,
dictionaries: Vec<PathBuf>,
}
fn dump_trie(trie: &Trie, delimiter: &str) -> io::Result<()> {
let mut num_keys = 0;
let mut agent = Agent::new();
let mut stdout = io::stdout();
agent.set_query_str("");
loop {
if trie.predictive_search(&mut agent) {
let key_bytes = agent.key().as_bytes();
stdout.write_all(key_bytes)?;
stdout.write_all(delimiter.as_bytes())?;
num_keys += 1;
} else {
break;
}
}
eprintln!("#keys: {}", num_keys);
Ok(())
}
fn dump_file(path: Option<&PathBuf>, args: &Args) -> io::Result<()> {
let mut trie = Trie::new();
if args.mmap_dictionary {
eprintln!("warning: memory-mapped I/O not yet implemented, using read instead");
}
if let Some(path) = path {
eprintln!("input: {}", path.display());
trie.load(path.to_str().unwrap())?;
} else {
eprintln!("input: <stdin>");
eprintln!("error: reading from stdin not yet supported, please specify a file");
return Err(io::Error::new(io::ErrorKind::Other, "stdin not supported"));
}
dump_trie(&trie, &args.delimiter)
}
fn main() {
let args = Args::parse();
let result = if args.dictionaries.is_empty() {
dump_file(None, &args)
} else {
let mut last_result = Ok(());
for path in &args.dictionaries {
if let Err(e) = dump_file(Some(path), &args) {
eprintln!("error: failed to dump {}: {}", path.display(), e);
last_result = Err(e);
break;
}
}
last_result
};
if result.is_err() {
process::exit(20);
}
}