use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use crate::types::expr::Expression;
use crate::wal::redo::{CreateMacroRedo, DropMacroRedo};
use crate::wal::traits::WalWriter;
use crate::wal::types::WalOpType;
use crate::StorageError;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MacroParamDef {
pub name: String,
pub default: Option<Expression>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MacroDef {
pub name: String,
pub params: Vec<MacroParamDef>,
pub body: Expression,
}
impl MacroDef {
pub fn new(name: String, params: Vec<MacroParamDef>, body: Expression) -> Self {
Self { name, params, body }
}
pub fn required_params(&self) -> usize {
self.params.iter().filter(|p| p.default.is_none()).count()
}
}
pub trait MacroStorage: Send + Sync {
fn load_all(&self) -> Result<Vec<MacroDef>, StorageError>;
fn save(&self, def: &MacroDef) -> Result<(), StorageError>;
fn delete(&self, name: &str) -> Result<(), StorageError>;
}
#[derive(Debug, Default)]
pub struct MemMacroStorage {
data: RwLock<HashMap<String, MacroDef>>,
}
impl MemMacroStorage {
pub fn new() -> Self {
Self {
data: RwLock::new(HashMap::new()),
}
}
}
impl MacroStorage for MemMacroStorage {
fn load_all(&self) -> Result<Vec<MacroDef>, StorageError> {
Ok(self.data.read().values().cloned().collect())
}
fn save(&self, def: &MacroDef) -> Result<(), StorageError> {
self.data.write().insert(def.name.clone(), def.clone());
Ok(())
}
fn delete(&self, name: &str) -> Result<(), StorageError> {
self.data.write().remove(name);
Ok(())
}
}
pub struct MacroManager {
macros: RwLock<HashMap<String, Arc<MacroDef>>>,
storage: Option<Arc<dyn MacroStorage>>,
wal_writer: Option<Arc<parking_lot::Mutex<dyn WalWriter>>>,
}
impl MacroManager {
pub fn new() -> Self {
Self {
macros: RwLock::new(HashMap::new()),
storage: None,
wal_writer: None,
}
}
pub fn with_storage(storage: Arc<dyn MacroStorage>) -> Self {
Self {
macros: RwLock::new(HashMap::new()),
storage: Some(storage),
wal_writer: None,
}
}
pub fn with_wal(
storage: Arc<dyn MacroStorage>,
wal_writer: Arc<parking_lot::Mutex<dyn WalWriter>>,
) -> Self {
Self {
macros: RwLock::new(HashMap::new()),
storage: Some(storage),
wal_writer: Some(wal_writer),
}
}
fn canonical(name: &str) -> String {
name.to_ascii_uppercase()
}
pub fn initialize(&self) -> Result<(), StorageError> {
if let Some(ref storage) = self.storage {
let defs = storage.load_all()?;
let mut map = self.macros.write();
for def in defs {
map.insert(Self::canonical(&def.name), Arc::new(def));
}
}
Ok(())
}
pub fn validate_definition(def: &MacroDef) -> Result<(), StorageError> {
if def.name.trim().is_empty() {
return Err(StorageError::db_error("Macro name must not be empty"));
}
let mut seen = std::collections::HashSet::new();
for param in &def.params {
if param.name.trim().is_empty() {
return Err(StorageError::db_error(format!(
"Macro '{}' has an empty parameter name",
def.name
)));
}
let key = param.name.to_ascii_uppercase();
if !seen.insert(key) {
return Err(StorageError::db_error(format!(
"Macro '{}' has duplicate parameter '{}'",
def.name, param.name
)));
}
}
Ok(())
}
pub fn create_macro(&self, def: MacroDef) -> Result<(), StorageError> {
Self::validate_definition(&def)?;
let key = Self::canonical(&def.name);
{
let map = self.macros.read();
if map.contains_key(&key) {
return Err(StorageError::db_error(format!(
"Macro '{}' already exists",
def.name
)));
}
}
if let Some(ref storage) = self.storage {
storage.save(&def)?;
}
self.write_wal_create(&def)?;
let mut map = self.macros.write();
if map.contains_key(&key) {
return Err(StorageError::db_error(format!(
"Macro '{}' already exists",
def.name
)));
}
map.insert(key, Arc::new(def));
Ok(())
}
pub fn drop_macro(&self, name: &str) -> Result<(), StorageError> {
let key = Self::canonical(name);
{
let mut map = self.macros.write();
if map.remove(&key).is_none() {
return Err(StorageError::db_error(format!(
"Macro '{}' does not exist",
name
)));
}
}
if let Some(ref storage) = self.storage {
storage.delete(name)?;
}
self.write_wal_drop(name)?;
Ok(())
}
pub fn get_macro(&self, name: &str) -> Option<Arc<MacroDef>> {
let map = self.macros.read();
map.get(&Self::canonical(name)).cloned()
}
pub fn exists(&self, name: &str) -> bool {
let map = self.macros.read();
map.contains_key(&Self::canonical(name))
}
pub fn list_macros(&self) -> Vec<Arc<MacroDef>> {
let map = self.macros.read();
let mut defs: Vec<Arc<MacroDef>> = map.values().cloned().collect();
defs.sort_by(|a, b| a.name.cmp(&b.name));
defs
}
fn wal_timestamp() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn write_wal_create(&self, def: &MacroDef) -> Result<(), StorageError> {
if let Some(ref wal_writer) = self.wal_writer {
let redo = CreateMacroRedo {
space_name: String::new(),
name: def.name.clone(),
params: def.params.clone(),
body: def.body.clone(),
};
let payload = postcard::to_allocvec(&redo)
.map_err(|e| StorageError::serialize_error(e.to_string()))?;
let mut writer = wal_writer.lock();
writer
.append_entry(WalOpType::CreateMacro, Self::wal_timestamp(), &payload)
.map_err(|e| StorageError::db_error(format!("WAL write error: {}", e)))?;
}
Ok(())
}
fn write_wal_drop(&self, name: &str) -> Result<(), StorageError> {
if let Some(ref wal_writer) = self.wal_writer {
let redo = DropMacroRedo {
space_name: String::new(),
name: name.to_string(),
};
let payload = postcard::to_allocvec(&redo)
.map_err(|e| StorageError::serialize_error(e.to_string()))?;
let mut writer = wal_writer.lock();
writer
.append_entry(WalOpType::DropMacro, Self::wal_timestamp(), &payload)
.map_err(|e| StorageError::db_error(format!("WAL write error: {}", e)))?;
}
Ok(())
}
pub fn replay_create(&self, redo: &CreateMacroRedo) -> Result<(), StorageError> {
let def = MacroDef::new(redo.name.clone(), redo.params.clone(), redo.body.clone());
Self::validate_definition(&def)?;
let mut map = self.macros.write();
map.insert(Self::canonical(&def.name), Arc::new(def));
Ok(())
}
pub fn replay_drop(&self, redo: &DropMacroRedo) -> Result<(), StorageError> {
let mut map = self.macros.write();
map.remove(&Self::canonical(&redo.name));
Ok(())
}
}
impl Default for MacroManager {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for MacroManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MacroManager")
.field("count", &self.macros.read().len())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::expr::Expression;
fn sample_def(name: &str) -> MacroDef {
MacroDef::new(
name.to_string(),
vec![MacroParamDef {
name: "x".to_string(),
default: None,
}],
Expression::Binary {
left: Box::new(Expression::variable("x")),
op: crate::types::operators::BinaryOperator::Multiply,
right: Box::new(Expression::literal(crate::Value::Int(2))),
},
)
}
#[test]
fn test_macro_create_get_drop() {
let manager = MacroManager::new();
manager.create_macro(sample_def("double")).unwrap();
assert!(manager.exists("double"));
assert!(manager.exists("DOUBLE"));
let def = manager.get_macro("Double").unwrap();
assert_eq!(def.params.len(), 1);
assert_eq!(def.required_params(), 1);
manager.drop_macro("double").unwrap();
assert!(!manager.exists("double"));
}
#[test]
fn test_macro_duplicate_create() {
let manager = MacroManager::new();
manager.create_macro(sample_def("m")).unwrap();
assert!(manager.create_macro(sample_def("M")).is_err());
}
#[test]
fn test_macro_drop_missing() {
let manager = MacroManager::new();
assert!(manager.drop_macro("nope").is_err());
}
#[test]
fn test_macro_duplicate_params_rejected() {
let mut def = sample_def("bad");
def.params.push(MacroParamDef {
name: "X".to_string(),
default: None,
});
assert!(MacroManager::validate_definition(&def).is_err());
}
#[test]
fn test_macro_list_sorted() {
let manager = MacroManager::new();
manager.create_macro(sample_def("b")).unwrap();
manager.create_macro(sample_def("a")).unwrap();
let names: Vec<String> = manager
.list_macros()
.iter()
.map(|d| d.name.clone())
.collect();
assert_eq!(names, vec!["a".to_string(), "b".to_string()]);
}
#[test]
fn test_macro_with_storage() {
let storage = Arc::new(MemMacroStorage::new());
let manager = MacroManager::with_storage(storage);
manager.create_macro(sample_def("m1")).unwrap();
assert!(manager.exists("m1"));
manager.initialize().unwrap();
assert!(manager.exists("m1"));
}
#[test]
fn test_macro_replay() {
let manager = MacroManager::new();
let def = sample_def("r");
manager
.replay_create(&CreateMacroRedo {
space_name: String::new(),
name: def.name.clone(),
params: def.params.clone(),
body: def.body.clone(),
})
.unwrap();
assert!(manager.exists("r"));
manager
.replay_drop(&DropMacroRedo {
space_name: String::new(),
name: "r".to_string(),
})
.unwrap();
assert!(!manager.exists("r"));
}
}