reifydb_testing/testscript/

command.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 ReifyDB

// This file includes and modifies code from the toydb project (https://github.com/erikgrinaker/toydb),
// originally licensed under the Apache License, Version 2.0.
// Original copyright:
//   Copyright (c) 2024 Erik Grinaker
//
// The original Apache License can be found at:
//   http://www.apache.org/licenses/LICENSE-2.0

use std::{
	collections::{BTreeSet, HashSet, VecDeque},
	error::Error,
	fmt,
	str::FromStr,
};

#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub(crate) struct Block {
	pub commands: Vec<Command>,

	pub literal: String,

	pub line_number: u32,
}

#[derive(Clone, PartialEq)]
#[non_exhaustive]
pub struct Command {
	pub name: String,

	pub args: Vec<Argument>,

	pub prefix: Option<String>,

	pub tags: HashSet<String>,

	pub silent: bool,

	pub fail: bool,

	pub line_number: u32,
}

impl fmt::Debug for Command {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_struct("Command")
			.field("name", &self.name)
			.field("args", &self.args)
			.field("prefix", &self.prefix)
			.field("tags", &BTreeSet::from_iter(&self.tags))
			.field("silent", &self.silent)
			.field("fail", &self.fail)
			.field("line_number", &self.line_number)
			.finish()
	}
}

impl Command {
	pub fn consume_args(&self) -> ArgumentConsumer<'_> {
		ArgumentConsumer::new(&self.args)
	}
}

#[derive(Clone, Debug, PartialEq)]
#[non_exhaustive]
pub struct Argument {
	pub key: Option<String>,

	pub value: String,
}

impl Argument {
	pub fn name(&self) -> &str {
		match self.key.as_deref() {
			Some(key) => key,
			None => &self.value,
		}
	}

	pub fn parse<T>(&self) -> Result<T, Box<dyn Error>>
	where
		T: FromStr,
		<T as FromStr>::Err: fmt::Display,
	{
		self.value.parse().map_err(|e| format!("invalid argument '{}': {e}", self.value).into())
	}
}

pub struct ArgumentConsumer<'a> {
	args: VecDeque<&'a Argument>,
}

impl<'a> Iterator for ArgumentConsumer<'a> {
	type Item = &'a Argument;

	fn next(&mut self) -> Option<Self::Item> {
		self.args.pop_front()
	}
}

impl<'a> ArgumentConsumer<'a> {
	fn new(args: &'a [Argument]) -> Self {
		Self {
			args: VecDeque::from_iter(args.iter()),
		}
	}

	pub fn lookup(&mut self, key: &str) -> Option<&'a Argument> {
		let arg = self.args.iter().rev().find(|a| a.key.as_deref() == Some(key)).copied();
		if arg.is_some() {
			self.args.retain(|a| a.key.as_deref() != Some(key))
		}
		arg
	}

	pub fn lookup_parse<T>(&mut self, key: &str) -> Result<Option<T>, Box<dyn Error>>
	where
		T: FromStr,
		<T as FromStr>::Err: fmt::Display,
	{
		let value = self
			.args
			.iter()
			.rev()
			.find(|a| a.key.as_deref() == Some(key))
			.map(|a| a.parse())
			.transpose()?;
		if value.is_some() {
			self.args.retain(|a| a.key.as_deref() != Some(key))
		}
		Ok(value)
	}

	pub fn next_key(&mut self) -> Option<&'a Argument> {
		self.args.iter().position(|a| a.key.is_some()).map(|i| self.args.remove(i).unwrap())
	}

	pub fn next_pos(&mut self) -> Option<&'a Argument> {
		self.args.iter().position(|a| a.key.is_none()).map(|i| self.args.remove(i).unwrap())
	}

	pub fn reject_rest(&self) -> Result<(), Box<dyn Error>> {
		if let Some(arg) = self.args.front() {
			return Err(format!("invalid argument '{}'", arg.name()).into());
		}
		Ok(())
	}

	pub fn rest(&mut self) -> Vec<&'a Argument> {
		self.args.drain(..).collect()
	}

	pub fn rest_key(&mut self) -> Vec<&'a Argument> {
		let keyed: Vec<_> = self.args.iter().filter(|a| a.key.is_some()).copied().collect();
		if !keyed.is_empty() {
			self.args.retain(|a| a.key.is_none());
		}
		keyed
	}

	pub fn rest_pos(&mut self) -> Vec<&'a Argument> {
		let pos: Vec<_> = self.args.iter().filter(|a| a.key.is_none()).copied().collect();
		if !pos.is_empty() {
			self.args.retain(|a| a.key.is_some());
		}
		pos
	}
}

#[cfg(test)]
pub mod tests {
	use super::*;

	/// Constructs an Argument from a string value or key => value.
	macro_rules! arg {
		($value:expr) => {
			Argument {
				key: None,
				value: $value.to_string(),
			}
		};
		($key:expr => $value:expr) => {
			Argument {
				key: Some($key.to_string()),
				value: $value.to_string(),
			}
		};
	}

	/// Constructs a Command by parsing the given input string.
	macro_rules! cmd {
		($input:expr) => {{ crate::testscript::parser::parse_command(&format!("{}\n", $input)).expect("invalid command") }};
	}

	/// Tests Argument.name().
	#[test]
	fn test_argument_name() {
		assert_eq!(arg!("value").name(), "value");
		assert_eq!(arg!("key" => "value").name(), "key");
	}

	/// Basic tests of Argument.parse(). Not comprehensive, since it
	/// dispatches to flow::str::parse().
	#[test]
	fn test_argument_parse() {
		assert_eq!(arg!("-1").parse::<i64>().unwrap(), -1_i64);
		assert_eq!(arg!("0").parse::<i64>().unwrap(), 0_i64);
		assert_eq!(arg!("1").parse::<i64>().unwrap(), 1_i64);

		assert_eq!(
			arg!("").parse::<i64>().unwrap_err().to_string(),
			"invalid argument '': cannot parse integer from empty string"
		);
		assert_eq!(
			arg!("foo").parse::<i64>().unwrap_err().to_string(),
			"invalid argument 'foo': invalid digit found in string"
		);

		assert!(!arg!("false").parse::<bool>().unwrap());
		assert!(arg!("true").parse::<bool>().unwrap());

		assert_eq!(
			arg!("").parse::<bool>().unwrap_err().to_string(),
			"invalid argument '': provided string was not `true` or `false`"
		);
	}

	/// Tests Command.consume_args(). ArgumentConsumer is tested separately.
	#[test]
	fn test_command_consume_args() {
		let cmd = cmd!("cmd foo key=value bar");
		assert_eq!(cmd.consume_args().rest(), vec![&cmd.args[0], &cmd.args[1], &cmd.args[2]]);
	}

	/// Tests ArgumentConsumer.lookup().
	#[test]
	fn test_argument_consumer_lookup() {
		let cmd = cmd!("cmd value key=value foo=bar key=other");

		// lookup() returns None on unknown keys, including ones that
		// match a value argument.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup("unknown"), None);
		assert_eq!(args.lookup("value"), None);
		assert_eq!(args.rest().len(), 4);

		// lookup() removes duplicate keys, returning the last.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup("key"), Some(&cmd.args[3]));
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[2]]);

		// lookup() removes single keys.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup("foo"), Some(&cmd.args[2]));
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[1], &cmd.args[3]]);
	}

	/// Tests ArgumentConsumer.lookup_parse().
	#[test]
	fn test_argument_consumer_lookup_parse() {
		let cmd = cmd!("cmd value key=1 foo=bar key=2");

		// lookup_parse() returns None on unknown keys, including ones
		// that match a value argument.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup_parse::<String>("unknown").unwrap(), None);
		assert_eq!(args.lookup_parse::<String>("value").unwrap(), None);
		assert_eq!(args.rest().len(), 4);

		// lookup_parse() parses and removes duplicate keys, returning
		// the last.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup_parse("key").unwrap(), Some(2));
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[2]]);

		// lookup_parse() parses and removes single keys, with string
		// parsing being a noop.
		let mut args = cmd.consume_args();
		assert_eq!(args.lookup_parse("foo").unwrap(), Some("bar".to_string()));
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[1], &cmd.args[3]]);

		// lookup_parse() does not remove arguments on parse errors,
		// even with duplicate keys.
		let mut args = cmd.consume_args();
		assert!(args.lookup_parse::<bool>("key").is_err());
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[1], &cmd.args[2], &cmd.args[3]]);
	}

	/// Tests ArgumentConsumer.next(), next_pos(), and next_key().
	#[test]
	fn test_argument_consumer_next() {
		let cmd = cmd!("cmd foo key=1 key=2 bar");

		// next() returns references to all arguments and consumes them.
		let mut args = cmd.consume_args();
		assert_eq!(args.next(), Some(&cmd.args[0]));
		assert_eq!(args.next(), Some(&cmd.args[1]));
		assert_eq!(args.next(), Some(&cmd.args[2]));
		assert_eq!(args.next(), Some(&cmd.args[3]));
		assert_eq!(args.next(), None);
		assert!(args.rest().is_empty());

		// next_key() returns references to key/value arguments and
		// consumes them.
		let mut args = cmd.consume_args();
		assert_eq!(args.next_key(), Some(&cmd.args[1]));
		assert_eq!(args.next_key(), Some(&cmd.args[2]));
		assert_eq!(args.next_key(), None);
		assert_eq!(args.next(), Some(&cmd.args[0]));
		assert_eq!(args.next(), Some(&cmd.args[3]));
		assert_eq!(args.next(), None);
		assert!(args.rest().is_empty());

		// next_pos() returns references to key/value arguments and
		// consumes them.
		let mut args = cmd.consume_args();
		assert_eq!(args.next_pos(), Some(&cmd.args[0]));
		assert_eq!(args.next_pos(), Some(&cmd.args[3]));
		assert_eq!(args.next_pos(), None);
		assert_eq!(args.next(), Some(&cmd.args[1]));
		assert_eq!(args.next(), Some(&cmd.args[2]));
		assert_eq!(args.next(), None);
		assert!(args.rest().is_empty());
	}

	/// Tests ArgumentConsumer.reject_rest().
	#[test]
	fn test_argument_consumer_reject_rest() {
		// Empty args return Ok.
		let cmd = cmd!("cmd");
		assert!(cmd.consume_args().reject_rest().is_ok());

		// Positional argument fails. It does not consume the arg.
		let cmd = cmd!("cmd value");
		let mut args = cmd.consume_args();
		assert_eq!(args.reject_rest().unwrap_err().to_string(), "invalid argument 'value'");
		assert!(!args.rest().is_empty());

		// Key/value argument fails.
		let cmd = cmd!("cmd key=value");
		let mut args = cmd.consume_args();
		assert_eq!(args.reject_rest().unwrap_err().to_string(), "invalid argument 'key'");
		assert!(!args.rest().is_empty());
	}

	/// Tests ArgumentConsumer.rest(), rest_pos() and rest_key().
	#[test]
	fn test_argument_consumer_rest() {
		let cmd = cmd!("cmd foo key=1 key=2 bar");

		// rest() returns references to all arguments and consumes them.
		let mut args = cmd.consume_args();
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[1], &cmd.args[2], &cmd.args[3]]);
		assert!(args.rest().is_empty());

		// rest_pos() returns and consumes positional arguments.
		let mut args = cmd.consume_args();
		assert_eq!(args.rest_pos(), vec![&cmd.args[0], &cmd.args[3]]);
		assert!(args.rest_pos().is_empty());
		assert_eq!(args.rest(), vec![&cmd.args[1], &cmd.args[2]]);

		// rest_key() returns and consumes key/value arguments.
		let mut args = cmd.consume_args();
		assert_eq!(args.rest_key(), vec![&cmd.args[1], &cmd.args[2]]);
		assert!(args.rest_key().is_empty());
		assert_eq!(args.rest(), vec![&cmd.args[0], &cmd.args[3]]);
	}
}