weavegraph 0.7.0

//! Collection utilities for Weavegraph.
//!
//! Typed helpers for `FxHashMap<String, Value>` extra data maps and common
//! string-keyed map patterns used throughout the codebase.

use rustc_hash::FxHashMap;
use serde_json::Value;
use std::collections::HashMap;
use thiserror::Error;

/// Errors from typed collection operations.
#[derive(Debug, Error)]
#[cfg_attr(feature = "diagnostics", derive(miette::Diagnostic))]
pub enum CollectionError {
    /// The requested key was not present.
    #[error("Key '{key}' not found in collection")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::collections::missing_key))
    )]
    MissingKey {
        /// The missing key.
        key: String,
    },

    /// The stored value's JSON type did not match the requested type.
    #[error("Invalid type conversion for key '{key}': expected {expected}, found {found}")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::collections::type_mismatch))
    )]
    TypeMismatch {
        /// The key where the mismatch occurred.
        key: String,
        /// The type the caller expected.
        expected: String,
        /// The type that was actually stored.
        found: String,
    },

    /// A `serde_json` serialization or deserialization failure.
    #[error("JSON operation failed: {source}")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::collections::json))
    )]
    Json {
        /// The underlying JSON error.
        #[from]
        source: serde_json::Error,
    },
}

/// Creates a new empty `FxHashMap<String, Value>` for extra data storage.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::new_extra_map;
/// use serde_json::json;
///
/// let mut extra = new_extra_map();
/// extra.insert("key".to_string(), json!("value"));
/// ```
#[must_use]
#[inline]
pub fn new_extra_map() -> FxHashMap<String, Value> {
    FxHashMap::default()
}

/// Creates a new `FxHashMap<String, Value>` pre-allocated for `capacity` entries.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::new_extra_map_with_capacity;
///
/// let mut extra = new_extra_map_with_capacity(10);
/// ```
#[must_use]
#[inline]
pub fn new_extra_map_with_capacity(capacity: usize) -> FxHashMap<String, Value> {
    FxHashMap::with_capacity_and_hasher(capacity, Default::default())
}

/// Builds an extra map from an iterator of `(key, value)` pairs.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::extra_map_from_pairs;
/// use serde_json::json;
///
/// let extra = extra_map_from_pairs([("name", json!("test")), ("count", json!(42))]);
/// ```
#[must_use]
pub fn extra_map_from_pairs<I, K, V>(pairs: I) -> FxHashMap<String, Value>
where
    I: IntoIterator<Item = (K, V)>,
    K: Into<String>,
    V: Into<Value>,
{
    pairs
        .into_iter()
        .map(|(k, v)| (k.into(), v.into()))
        .collect()
}

/// Merges extra maps left-to-right; later entries win on key collision.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::{new_extra_map, merge_extra_maps};
/// use serde_json::json;
///
/// let mut a = new_extra_map();
/// a.insert("x".to_string(), json!(1));
///
/// let mut b = new_extra_map();
/// b.insert("x".to_string(), json!(2));
///
/// let merged = merge_extra_maps([&a, &b]);
/// assert_eq!(merged["x"], json!(2));
/// ```
#[must_use]
pub fn merge_extra_maps<'a, I>(maps: I) -> FxHashMap<String, Value>
where
    I: IntoIterator<Item = &'a FxHashMap<String, Value>>,
{
    maps.into_iter()
        .flat_map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())))
        .collect()
}

/// Typed insert and get methods for `FxHashMap<String, Value>` extra data maps.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::{new_extra_map, ExtraMapExt};
///
/// let mut extra = new_extra_map();
/// extra.insert_string("name", "Alice");
/// extra.insert_number("age", 30);
/// extra.insert_bool("active", true);
///
/// assert_eq!(extra.get_string("name").unwrap(), "Alice");
/// assert_eq!(extra.get_number("age").unwrap(), 30.into());
/// assert!(extra.get_bool("active").unwrap());
/// assert!(extra.has_typed("name", "string"));
/// assert!(!extra.has_typed("name", "number"));
/// ```
pub trait ExtraMapExt {
    /// Inserts a string value.
    fn insert_string(&mut self, key: impl Into<String>, value: impl Into<String>);

    /// Inserts a numeric value.
    fn insert_number(&mut self, key: impl Into<String>, value: impl Into<serde_json::Number>);

    /// Inserts a boolean value.
    fn insert_bool(&mut self, key: impl Into<String>, value: bool);

    /// Serializes and inserts any `Serialize` value; returns `Err` on serialization failure.
    fn insert_json<T: serde::Serialize>(
        &mut self,
        key: impl Into<String>,
        value: T,
    ) -> Result<(), CollectionError>;

    /// Returns the string at `key`, or `Err` on missing key or type mismatch.
    fn get_string(&self, key: &str) -> Result<String, CollectionError>;

    /// Returns the number at `key`, or `Err` on missing key or type mismatch.
    fn get_number(&self, key: &str) -> Result<serde_json::Number, CollectionError>;

    /// Returns the bool at `key`, or `Err` on missing key or type mismatch.
    fn get_bool(&self, key: &str) -> Result<bool, CollectionError>;

    /// Deserializes the value at `key` to `T`.
    fn get_typed<T: serde::de::DeserializeOwned>(&self, key: &str) -> Result<T, CollectionError>;

    /// Returns `true` if `key` exists with the given JSON type name
    /// (`"string"`, `"number"`, `"bool"`, `"array"`, `"object"`, `"null"`).
    fn has_typed(&self, key: &str, expected_type: &str) -> bool;
}

impl ExtraMapExt for FxHashMap<String, Value> {
    fn insert_string(&mut self, key: impl Into<String>, value: impl Into<String>) {
        self.insert(key.into(), Value::String(value.into()));
    }

    fn insert_number(&mut self, key: impl Into<String>, value: impl Into<serde_json::Number>) {
        self.insert(key.into(), Value::Number(value.into()));
    }

    fn insert_bool(&mut self, key: impl Into<String>, value: bool) {
        self.insert(key.into(), Value::Bool(value));
    }

    fn insert_json<T: serde::Serialize>(
        &mut self,
        key: impl Into<String>,
        value: T,
    ) -> Result<(), CollectionError> {
        self.insert(key.into(), serde_json::to_value(value)?);
        Ok(())
    }

    fn get_string(&self, key: &str) -> Result<String, CollectionError> {
        match self.get(key) {
            Some(Value::String(s)) => Ok(s.clone()),
            Some(other) => Err(CollectionError::TypeMismatch {
                key: key.to_owned(),
                expected: "string".to_owned(),
                found: format!("{other:?}"),
            }),
            None => Err(CollectionError::MissingKey {
                key: key.to_owned(),
            }),
        }
    }

    fn get_number(&self, key: &str) -> Result<serde_json::Number, CollectionError> {
        match self.get(key) {
            Some(Value::Number(n)) => Ok(n.clone()),
            Some(other) => Err(CollectionError::TypeMismatch {
                key: key.to_owned(),
                expected: "number".to_owned(),
                found: format!("{other:?}"),
            }),
            None => Err(CollectionError::MissingKey {
                key: key.to_owned(),
            }),
        }
    }

    fn get_bool(&self, key: &str) -> Result<bool, CollectionError> {
        match self.get(key) {
            Some(Value::Bool(b)) => Ok(*b),
            Some(other) => Err(CollectionError::TypeMismatch {
                key: key.to_owned(),
                expected: "boolean".to_owned(),
                found: format!("{other:?}"),
            }),
            None => Err(CollectionError::MissingKey {
                key: key.to_owned(),
            }),
        }
    }

    fn get_typed<T: serde::de::DeserializeOwned>(&self, key: &str) -> Result<T, CollectionError> {
        let value = self.get(key).ok_or_else(|| CollectionError::MissingKey {
            key: key.to_owned(),
        })?;
        Ok(serde_json::from_value(value.clone())?)
    }

    fn has_typed(&self, key: &str, expected_type: &str) -> bool {
        matches!(
            (self.get(key), expected_type),
            (Some(Value::String(_)), "string")
                | (Some(Value::Number(_)), "number")
                | (Some(Value::Bool(_)), "bool")
                | (Some(Value::Array(_)), "array")
                | (Some(Value::Object(_)), "object")
                | (Some(Value::Null), "null")
        )
    }
}

/// Ergonomic methods for string-keyed maps.
///
/// # Examples
///
/// ```rust
/// use weavegraph::utils::collections::StringMapExt;
/// use rustc_hash::FxHashMap;
///
/// let mut counts: FxHashMap<String, u32> = FxHashMap::default();
/// counts.insert_or_update("hits".to_string(), 1, |n| *n += 1);
/// counts.insert_or_update("hits".to_string(), 1, |n| *n += 1);
/// assert_eq!(counts.get_or_default("hits", 0), 2);
/// assert_eq!(counts.get_or_default("misses", 0), 0);
/// ```
pub trait StringMapExt<V> {
    /// Returns the value at `key`, or `default` if the key is absent.
    fn get_or_default(&self, key: &str, default: V) -> V
    where
        V: Clone;

    /// Inserts `value` if `key` is absent; otherwise calls `update_fn` on the existing value.
    fn insert_or_update<F>(&mut self, key: String, value: V, update_fn: F)
    where
        F: FnOnce(&mut V);
}

impl<V> StringMapExt<V> for HashMap<String, V> {
    fn get_or_default(&self, key: &str, default: V) -> V
    where
        V: Clone,
    {
        self.get(key).cloned().unwrap_or(default)
    }

    fn insert_or_update<F>(&mut self, key: String, value: V, update_fn: F)
    where
        F: FnOnce(&mut V),
    {
        self.entry(key).and_modify(update_fn).or_insert(value);
    }
}

impl<V> StringMapExt<V> for FxHashMap<String, V> {
    fn get_or_default(&self, key: &str, default: V) -> V
    where
        V: Clone,
    {
        self.get(key).cloned().unwrap_or(default)
    }

    fn insert_or_update<F>(&mut self, key: String, value: V, update_fn: F)
    where
        F: FnOnce(&mut V),
    {
        self.entry(key).and_modify(update_fn).or_insert(value);
    }
}