hamelin_lib 0.15.6

use std::collections::HashSet;

use crate::tree::ast::identifier::{Identifier, SimpleIdentifier};
use crate::tree::ast::query::{Query, QueryKind};

use super::CombineQueriesError;

/// Tracks reserved DEF names while preparing inline queries for hoisting.
#[derive(Debug, Default)]
pub struct DefNameRegistry {
    used: HashSet<Identifier>,
}

impl DefNameRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn extend_from_query(&mut self, query: &Query) {
        if let QueryKind::Valid(valid) = &query.kind {
            for d in &valid.defs {
                if let Ok(id) = d.name.valid_ref() {
                    self.used.insert(id.clone());
                }
            }
        }
    }

    /// Allocate a hoisted CTE name derived from `base`, reserving it in the registry.
    pub fn allocate_cte_name(
        &mut self,
        base: &str,
    ) -> Result<SimpleIdentifier, CombineQueriesError> {
        let candidate = pick_available_def_name(base, &self.used)
            .map_err(CombineQueriesError::InvalidCteName)?;
        self.used.insert(candidate.clone().into());
        Ok(candidate)
    }

    pub fn used_names(&self) -> &HashSet<Identifier> {
        &self.used
    }
}

/// Pick a DEF name derived from `base` that does not appear in `blocked`.
pub fn pick_available_def_name(
    base: &str,
    blocked: &HashSet<Identifier>,
) -> Result<SimpleIdentifier, String> {
    let mut candidate = SimpleIdentifier::parse(base).map_err(|e| e.to_string())?;
    if !blocked.contains(&candidate.clone().into()) {
        return Ok(candidate);
    }

    let mut suffix = 2;
    loop {
        candidate =
            SimpleIdentifier::parse(&format!("{base}_{suffix}")).map_err(|e| e.to_string())?;
        if !blocked.contains(&candidate.clone().into()) {
            return Ok(candidate);
        }
        suffix += 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree::ast::ParseWithErrors;

    #[test]
    fn allocate_avoids_collision() {
        let mut registry = DefNameRegistry::new();
        let (seed, errs) =
            Query::parse_with_errors("DEF events_okta = 1 AS int;\nFROM events:base");
        assert!(errs.is_empty(), "{errs:?}");
        registry.extend_from_query(&seed);
        let name = registry.allocate_cte_name("events_okta").unwrap();
        assert_eq!(name.as_str(), "events_okta_2");
    }

    #[test]
    fn allocate_rejects_invalid_identifier() {
        let mut registry = DefNameRegistry::new();
        let err = registry.allocate_cte_name("123bad").unwrap_err();
        assert!(matches!(err, CombineQueriesError::InvalidCteName(_)));
    }
}