saluki_context/

context.rs

use std::{fmt, hash, sync::Arc};

use metrics::Gauge;
use saluki_common::collections::{ContiguousBitSet, PrehashedHashSet};
use stringtheory::MetaString;

use crate::{
    hash::{hash_context, hash_context_with_host_and_seen, ContextKey},
    tags::{Tag, TagSet},
};

const BASE_CONTEXT_SIZE: usize = std::mem::size_of::<Context>() + std::mem::size_of::<ContextInner>();

/// A metric context.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Context {
    inner: Arc<ContextInner>,
}

impl Context {
    /// Creates a new `Context` from the given static name.
    pub fn from_static_name(name: &'static str) -> Self {
        let tags = TagSet::default();
        let origin_tags = TagSet::default();

        let (key, _) = hash_context(name, &tags, &origin_tags);
        Self {
            inner: Arc::new(ContextInner {
                name: MetaString::from_static(name),
                host: None,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Creates a new `Context` from the given static name and given static tags.
    pub fn from_static_parts(name: &'static str, tags: &[&'static str]) -> Self {
        let mut tag_set = TagSet::with_capacity(tags.len());
        for tag in tags {
            tag_set.insert_tag(MetaString::from_static(tag));
        }

        let origin_tags = TagSet::default();

        let (key, _) = hash_context(name, &tag_set, &origin_tags);
        Self {
            inner: Arc::new(ContextInner {
                name: MetaString::from_static(name),
                host: None,
                tags: tag_set,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Creates a new `Context` from the given name and given tags.
    pub fn from_parts<S: Into<MetaString>>(name: S, tags: impl Into<TagSet>) -> Self {
        let name = name.into();
        let tags = tags.into();
        let origin_tags = TagSet::default();
        let (key, _) = hash_context(&name, &tags, &origin_tags);
        Self {
            inner: Arc::new(ContextInner {
                name,
                host: None,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Clones this context, and uses the given name for the cloned context.
    pub fn with_name<S: Into<MetaString>>(&self, name: S) -> Self {
        // Regenerate the context key to account for the new name.
        let name = name.into();
        let host = self.inner.host.clone();
        let tags = self.inner.tags.clone();
        let origin_tags = self.inner.origin_tags.clone();
        let key = ContextInner::calculate_key(&name, host.as_deref(), &tags, &origin_tags);

        Self {
            inner: Arc::new(ContextInner {
                name,
                host,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Clones this context, and uses the given tags for the cloned context.
    ///
    /// The name and origin tags of this context are preserved.
    pub fn with_tags(&self, tags: impl Into<TagSet>) -> Self {
        let name = self.inner.name.clone();
        let host = self.inner.host.clone();
        let tags = tags.into();
        let origin_tags = self.inner.origin_tags.clone();
        let key = ContextInner::calculate_key(&name, host.as_deref(), &tags, &origin_tags);

        Self {
            inner: Arc::new(ContextInner {
                name,
                host,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Clones this context, and uses the given origin tags for the cloned context.
    ///
    /// The name and instrumented tags of this context are preserved.
    pub fn with_origin_tags(&self, origin_tags: impl Into<TagSet>) -> Self {
        let name = self.inner.name.clone();
        let host = self.inner.host.clone();
        let tags = self.inner.tags.clone();
        let origin_tags = origin_tags.into();
        let key = ContextInner::calculate_key(&name, host.as_deref(), &tags, &origin_tags);

        Self {
            inner: Arc::new(ContextInner {
                name,
                host,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Clones this context, replacing both instrumented tags and origin tags in a single allocation.
    ///
    /// Preferred over two separate `with_tags` / `with_origin_tags` calls when both sets need to
    /// be replaced, as it halves the number of `Arc` allocations.
    pub fn with_tags_and_origin_tags(&self, tags: impl Into<TagSet>, origin_tags: impl Into<TagSet>) -> Self {
        let name = self.inner.name.clone();
        let host = self.inner.host.clone();
        let tags = tags.into();
        let origin_tags = origin_tags.into();
        let key = ContextInner::calculate_key(&name, host.as_deref(), &tags, &origin_tags);

        Self {
            inner: Arc::new(ContextInner {
                name,
                host,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    pub(crate) fn from_inner(inner: ContextInner) -> Self {
        Self { inner: Arc::new(inner) }
    }

    #[cfg(test)]
    pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.inner, &other.inner)
    }

    /// Returns the name of this context.
    pub fn name(&self) -> &MetaString {
        &self.inner.name
    }

    /// Returns the host of this context, if one has been set.
    pub fn host(&self) -> Option<&str> {
        self.inner.host.as_deref()
    }

    /// Clones this context, and uses the given host for the cloned context.
    pub fn with_host<S: Into<Option<MetaString>>>(&self, host: S) -> Self {
        let name = self.inner.name.clone();
        let host = host.into();
        let tags = self.inner.tags.clone();
        let origin_tags = self.inner.origin_tags.clone();
        let key = ContextInner::calculate_key(&name, host.as_deref(), &tags, &origin_tags);

        Self {
            inner: Arc::new(ContextInner {
                name,
                host,
                tags,
                origin_tags,
                key,
                active_count: Gauge::noop(),
            }),
        }
    }

    /// Returns the instrumented tags of this context.
    pub fn tags(&self) -> &TagSet {
        &self.inner.tags
    }

    /// Returns the origin tags of this context.
    pub fn origin_tags(&self) -> &TagSet {
        &self.inner.origin_tags
    }

    /// Mutates the instrumented tags of this context via a closure.
    ///
    /// Uses copy-on-write semantics: if this context shares its inner data with other clones, the
    /// inner data is cloned first so that mutations don't affect other holders. If this context is
    /// the sole owner, the mutation happens in place.
    ///
    /// The context key is automatically recomputed after the closure returns.
    pub fn mutate_tags(&mut self, f: impl FnOnce(&mut TagSet)) {
        self.mutate_inner(|inner| f(&mut inner.tags));
    }

    /// Mutates the origin tags of this context via a closure.
    ///
    /// Uses copy-on-write semantics: if this context shares its inner data with other clones, the
    /// inner data is cloned first so that mutations don't affect other holders. If this context is
    /// the sole owner, the mutation happens in place.
    ///
    /// The context key is automatically recomputed after the closure returns.
    pub fn mutate_origin_tags(&mut self, f: impl FnOnce(&mut TagSet)) {
        self.mutate_inner(|inner| f(&mut inner.origin_tags));
    }

    /// Mutates both instrumented tags and origin tags via a single closure.
    ///
    /// Uses copy-on-write semantics: if this context shares its inner data with other clones, the
    /// inner data is cloned first so that mutations don't affect other holders. If this context is
    /// the sole owner, the mutation happens in place.
    ///
    /// The context key is recomputed once after the closure returns.
    pub fn with_tag_sets_mut(&mut self, f: impl FnOnce(&mut TagSet, &mut TagSet)) {
        self.mutate_inner(|inner| f(&mut inner.tags, &mut inner.origin_tags));
    }

    /// Runs the given closure on the inner context data, recomputing the context key afterwards.
    ///
    /// When the inner context state is shared (we aren't the only ones with a strong reference), we clone the inner
    /// data first to have our own copy. Otherwise, we modify the inner data in place.
    fn mutate_inner(&mut self, f: impl FnOnce(&mut ContextInner)) {
        let inner = Arc::make_mut(&mut self.inner);
        f(inner);
        inner.recalculate_key();
    }

    /// Creates a lazy copy-on-write mutable view over this context's tag sets.
    ///
    /// The returned view supports mutations (for example, [`retain_tags`][TagSetMutView::retain_tags])
    /// without immediately triggering an `Arc` clone. The actual clone, mutation, and context key
    /// recomputation only happen when [`TagSetMutView::finish`] is called, and only if changes
    /// were actually recorded.
    ///
    /// `state` provides reusable scratch space for tracking pending changes. Holding a
    /// long-lived [`TagSetMutViewState`] across calls amortizes any vector allocations.
    pub fn tags_mut_view<'a, 'b>(&'a mut self, state: &'b mut TagSetMutViewState) -> TagSetMutView<'a, 'b> {
        TagSetMutView { context: self, state }
    }

    /// Returns the size of this context in bytes.
    ///
    /// A context's size is the sum of the sizes of its fields and the size of the `Context` struct itself, and
    /// includes:
    /// - the context name
    /// - the context host and tags (both instrumented and origin)
    ///
    /// Since origin tags can potentially be expensive to calculate, this method will cache the size of the origin tags
    /// when this method is first called.
    ///
    /// Additionally, the value returned by this method doesn't compensate for externalities such as origin tags
    /// potentially being shared by multiple contexts, or whether or not tags are inlined, interned, or heap
    /// allocated. This means that the value returned is essentially the worst-case usage, and should be used as a rough
    /// estimate.
    pub fn size_of(&self) -> usize {
        let name_size = self.inner.name.len();
        let host_size = self.inner.host.as_ref().map_or(0, |host| host.len());
        let tags_size = self.inner.tags.size_of();
        let origin_tags_size = self.inner.origin_tags.size_of();

        BASE_CONTEXT_SIZE + name_size + host_size + tags_size + origin_tags_size
    }
}

impl From<&'static str> for Context {
    fn from(name: &'static str) -> Self {
        Self::from_static_name(name)
    }
}

impl<'a> From<(&'static str, &'a [&'static str])> for Context {
    fn from((name, tags): (&'static str, &'a [&'static str])) -> Self {
        Self::from_static_parts(name, tags)
    }
}

impl fmt::Display for Context {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.inner.name)?;
        if !self.inner.tags.is_empty() {
            write!(f, "{{")?;

            let mut needs_separator = false;
            for tag in &self.inner.tags {
                if needs_separator {
                    write!(f, ", ")?;
                } else {
                    needs_separator = true;
                }

                write!(f, "{}", tag)?;
            }

            write!(f, "}}")?;
        }

        Ok(())
    }
}

/// Reusable scratch space for [`TagSetMutView`] operations.
///
/// Holding a long-lived instance across calls amortizes bitset allocations. The bitsets are
/// cleared automatically when the associated [`TagSetMutView`] is dropped.
#[derive(Debug, Default)]
pub struct TagSetMutViewState {
    tag_base_removals: ContiguousBitSet,
    tag_addition_removals: ContiguousBitSet,
    origin_base_removals: ContiguousBitSet,
    origin_addition_removals: ContiguousBitSet,
    hash_seen: PrehashedHashSet<u64>,
}

impl TagSetMutViewState {
    /// Creates a new, empty state.
    pub fn new() -> Self {
        Self::default()
    }

    fn clear(&mut self) {
        self.tag_base_removals.clear_all();
        self.tag_addition_removals.clear_all();
        self.origin_base_removals.clear_all();
        self.origin_addition_removals.clear_all();
    }
}

/// A lazy copy-on-write mutable view over a [`Context`]'s tag sets.
///
/// Operations on this view (for example, [`retain_tags`][Self::retain_tags]) are recorded but not
/// applied immediately. The actual `Arc` clone, mutation, and context key recomputation only
/// occur when [`finish`][Self::finish] is called, and only if changes were recorded.
pub struct TagSetMutView<'a, 'b> {
    context: &'a mut Context,
    state: &'b mut TagSetMutViewState,
}

impl<'a, 'b> TagSetMutView<'a, 'b> {
    /// Scan instrumented tags with the given predicate.
    ///
    /// Tags for which `f` returns `false` are flagged for removal. This is a read-only scan;
    /// no mutation occurs until [`finish`][Self::finish] is called.
    pub fn retain_tags(&mut self, f: impl FnMut(&Tag) -> bool) {
        self.context.inner.tags.collect_removals(
            f,
            &mut self.state.tag_base_removals,
            &mut self.state.tag_addition_removals,
        );
    }

    /// Scan origin tags with the given predicate.
    ///
    /// Tags for which `f` returns `false` are flagged for removal. This is a read-only scan;
    /// no mutation occurs until [`finish`][Self::finish] is called.
    pub fn retain_origin_tags(&mut self, f: impl FnMut(&Tag) -> bool) {
        self.context.inner.origin_tags.collect_removals(
            f,
            &mut self.state.origin_base_removals,
            &mut self.state.origin_addition_removals,
        );
    }

    /// Apply all recorded changes and return the total number of tags affected.
    ///
    /// If no changes were recorded, this is a no-op: no `Arc` clone, no rehash, returns 0.
    /// Otherwise, triggers `Arc::make_mut` on the context, applies the changes to both tag sets,
    /// and recomputes the context key.
    ///
    /// Returns the number of tags removed.
    pub fn finish(self) -> usize {
        let total_tags = self.state.tag_base_removals.len() + self.state.tag_addition_removals.len();
        let total_origin = self.state.origin_base_removals.len() + self.state.origin_addition_removals.len();
        let total = total_tags + total_origin;

        if total == 0 {
            return 0;
        }

        let inner = Arc::make_mut(&mut self.context.inner);

        if total_tags > 0 {
            inner
                .tags
                .apply_removals(&self.state.tag_base_removals, &self.state.tag_addition_removals);
        }
        if total_origin > 0 {
            inner
                .origin_tags
                .apply_removals(&self.state.origin_base_removals, &self.state.origin_addition_removals);
        }

        inner.recalculate_key_with_seen(&mut self.state.hash_seen);

        total
    }
}

impl Drop for TagSetMutView<'_, '_> {
    fn drop(&mut self) {
        self.state.clear();
    }
}

pub(super) struct ContextInner {
    key: ContextKey,
    name: MetaString,
    host: Option<MetaString>,
    tags: TagSet,
    origin_tags: TagSet,
    active_count: Gauge,
}

impl ContextInner {
    pub fn from_parts(
        key: ContextKey, name: MetaString, host: Option<MetaString>, tags: TagSet, origin_tags: TagSet,
        active_count: Gauge,
    ) -> Self {
        Self {
            key,
            name,
            host,
            tags,
            origin_tags,
            active_count,
        }
    }

    fn calculate_key(name: &str, host: Option<&str>, tags: &TagSet, origin_tags: &TagSet) -> ContextKey {
        let mut seen = PrehashedHashSet::default();
        let (key, _) = hash_context_with_host_and_seen(name, host, tags, origin_tags, &mut seen);
        key
    }

    fn recalculate_key(&mut self) {
        self.key = Self::calculate_key(&self.name, self.host.as_deref(), &self.tags, &self.origin_tags);
    }

    fn recalculate_key_with_seen(&mut self, seen: &mut PrehashedHashSet<u64>) {
        let (key, _) =
            hash_context_with_host_and_seen(&self.name, self.host.as_deref(), &self.tags, &self.origin_tags, seen);
        self.key = key;
    }
}

impl Clone for ContextInner {
    fn clone(&self) -> Self {
        Self {
            key: self.key,
            name: self.name.clone(),
            host: self.host.clone(),
            tags: self.tags.clone(),
            origin_tags: self.origin_tags.clone(),

            // We're specifically detaching this context from the statistics of the resolver from which `self`
            // originated, as we only want to track the statistics of the contexts created _directly_ through the
            // resolver.
            active_count: Gauge::noop(),
        }
    }
}

impl Drop for ContextInner {
    fn drop(&mut self) {
        self.active_count.decrement(1);
    }
}

impl PartialEq<ContextInner> for ContextInner {
    fn eq(&self, other: &ContextInner) -> bool {
        // TODO: Note about why we consider the hash good enough for equality.
        self.key == other.key
    }
}

impl Eq for ContextInner {}

impl hash::Hash for ContextInner {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        self.key.hash(state);
    }
}

impl fmt::Debug for ContextInner {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ContextInner")
            .field("name", &self.name)
            .field("host", &self.host)
            .field("tags", &self.tags)
            .field("key", &self.key)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tags::Tag;

    const SIZE_OF_CONTEXT_NAME: &str = "size_of_test_metric";
    const SIZE_OF_CONTEXT_CHANGED_NAME: &str = "size_of_test_metric_changed";
    const SIZE_OF_CONTEXT_TAGS: &[&str] = &["size_of_test_tag1", "size_of_test_tag2"];
    const SIZE_OF_CONTEXT_ORIGIN_TAGS: &[&str] = &["size_of_test_origin_tag1", "size_of_test_origin_tag2"];

    fn tag_set(tags: &[&str]) -> TagSet {
        tags.iter().map(|s| Tag::from(*s)).collect::<TagSet>()
    }

    #[test]
    fn size_of_context_from_static_name() {
        let context = Context::from_static_name(SIZE_OF_CONTEXT_NAME);
        assert_eq!(context.size_of(), BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_NAME.len());
    }

    #[test]
    fn size_of_context_from_static_parts() {
        let tags = tag_set(SIZE_OF_CONTEXT_TAGS);

        let context = Context::from_static_parts(SIZE_OF_CONTEXT_NAME, SIZE_OF_CONTEXT_TAGS);
        assert_eq!(
            context.size_of(),
            BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_NAME.len() + tags.size_of()
        );
    }

    #[test]
    fn size_of_context_from_parts() {
        let tags = tag_set(SIZE_OF_CONTEXT_TAGS);

        let context = Context::from_parts(SIZE_OF_CONTEXT_NAME, tags.clone());
        assert_eq!(
            context.size_of(),
            BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_NAME.len() + tags.size_of()
        );
    }

    #[test]
    fn size_of_context_with_name() {
        // Check the check after `with_name` when there's both tags and no tags.
        let context = Context::from_static_name(SIZE_OF_CONTEXT_NAME).with_name(SIZE_OF_CONTEXT_CHANGED_NAME);
        assert_eq!(
            context.size_of(),
            BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_CHANGED_NAME.len()
        );

        let tags = tag_set(SIZE_OF_CONTEXT_TAGS);

        let context = Context::from_static_parts(SIZE_OF_CONTEXT_NAME, SIZE_OF_CONTEXT_TAGS)
            .with_name(SIZE_OF_CONTEXT_CHANGED_NAME);
        assert_eq!(
            context.size_of(),
            BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_CHANGED_NAME.len() + tags.size_of()
        );
    }

    #[test]
    fn size_of_context_origin_tags() {
        let tags = tag_set(SIZE_OF_CONTEXT_TAGS);
        let origin_tags = tag_set(SIZE_OF_CONTEXT_ORIGIN_TAGS);

        let (key, _) = hash_context(SIZE_OF_CONTEXT_NAME, SIZE_OF_CONTEXT_TAGS, SIZE_OF_CONTEXT_ORIGIN_TAGS);

        let context = Context::from_inner(ContextInner {
            key,
            name: MetaString::from_static(SIZE_OF_CONTEXT_NAME),
            host: None,
            tags: tags.clone(),
            origin_tags: origin_tags.clone(),
            active_count: Gauge::noop(),
        });

        // Make sure the size of the context is correct with origin tags.
        assert_eq!(
            context.size_of(),
            BASE_CONTEXT_SIZE + SIZE_OF_CONTEXT_NAME.len() + tags.size_of() + origin_tags.size_of()
        );
    }

    #[test]
    fn with_tags_mut_clones_shared_context() {
        let original = Context::from_static_parts("metric", &["env:prod"]);
        let mut mutated = original.clone();

        // They share the same Arc before mutation.
        assert!(original.ptr_eq(&mutated));

        mutated.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("service:web"));
        });

        // After mutation, they no longer share the same inner.
        assert!(!original.ptr_eq(&mutated));
    }

    #[test]
    fn with_tags_mut_does_not_affect_original() {
        let original = Context::from_static_parts("metric", &["env:prod"]);
        let mut mutated = original.clone();

        mutated.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("service:web"));
        });

        // Original is unchanged.
        assert_eq!(original.tags().len(), 1);
        assert!(original.tags().has_tag("env:prod"));
        assert!(!original.tags().has_tag("service:web"));

        // Mutated has both tags.
        assert_eq!(mutated.tags().len(), 2);
        assert!(mutated.tags().has_tag("env:prod"));
        assert!(mutated.tags().has_tag("service:web"));
    }

    #[test]
    fn with_tags_mut_rehashes() {
        // Build a context and mutate it to add a tag.
        let mut mutated = Context::from_static_parts("metric", &["env:prod"]);
        mutated.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("service:web"));
        });

        // Build an equivalent context from scratch with both tags.
        let expected = Context::from_static_parts("metric", &["env:prod", "service:web"]);

        // The recomputed key should match a freshly-constructed context with the same state.
        assert_eq!(mutated, expected);

        // Modify a tag on the mutated context that _isn't_ shared with `expected` to ensure that there's no asymmetric
        // equality logic.
        mutated.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("cluster:foo"));
        });
        assert_ne!(mutated, expected);
    }

    #[test]
    fn with_origin_tags_mut_clones_shared_context() {
        let original = Context::from_static_name("metric");
        let mut mutated = original.clone();

        assert!(original.ptr_eq(&mutated));

        mutated.mutate_origin_tags(|tags| {
            tags.insert_tag(Tag::from("origin:tag"));
        });

        assert!(!original.ptr_eq(&mutated));
        assert!(original.origin_tags().is_empty());
        assert_eq!(mutated.origin_tags().len(), 1);
        assert!(mutated.origin_tags().has_tag("origin:tag"));
    }

    // --- Helper for contexts with origin tags ---

    fn context_with_origin(name: &'static str, tags: &[&'static str], origin_tags: &[&'static str]) -> Context {
        let (key, _) = hash_context(name, tags, origin_tags);
        Context::from_inner(ContextInner {
            key,
            name: MetaString::from_static(name),
            host: None,
            tags: tag_set(tags),
            origin_tags: tag_set(origin_tags),
            active_count: Gauge::noop(),
        })
    }

    fn context_with_host(
        name: &'static str, host: &'static str, tags: &[&'static str], origin_tags: &[&'static str],
    ) -> Context {
        let tags = tag_set(tags);
        let origin_tags = tag_set(origin_tags);
        let key = ContextInner::calculate_key(name, Some(host), &tags, &origin_tags);
        Context::from_inner(ContextInner {
            key,
            name: MetaString::from_static(name),
            host: Some(MetaString::from_static(host)),
            tags,
            origin_tags,
            active_count: Gauge::noop(),
        })
    }

    #[test]
    fn host_participates_in_context_identity() {
        let host_a = context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]);
        let host_b = context_with_host("metric", "host-b", &["env:prod"], &["origin:a"]);
        let host_a_again = context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]);
        let no_host = context_with_origin("metric", &["env:prod"], &["origin:a"]);

        assert_ne!(host_a, host_b);
        assert_ne!(host_a, no_host);
        assert_eq!(host_a, host_a_again);
        assert_eq!(host_a.host(), Some("host-a"));
        assert_eq!(no_host.host(), None);
    }

    #[test]
    fn host_is_preserved_when_context_is_copied_with_new_parts() {
        let base = context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]);

        let renamed = base.with_name("renamed");
        assert_eq!(
            renamed,
            context_with_host("renamed", "host-a", &["env:prod"], &["origin:a"])
        );
        assert_ne!(renamed, context_with_origin("renamed", &["env:prod"], &["origin:a"]));

        let retagged = base.with_tags(tag_set(&["service:web"]));
        assert_eq!(
            retagged,
            context_with_host("metric", "host-a", &["service:web"], &["origin:a"])
        );
        assert_ne!(retagged, context_with_origin("metric", &["service:web"], &["origin:a"]));

        let reorigined = base.with_origin_tags(tag_set(&["origin:b"]));
        assert_eq!(
            reorigined,
            context_with_host("metric", "host-a", &["env:prod"], &["origin:b"])
        );
        assert_ne!(reorigined, context_with_origin("metric", &["env:prod"], &["origin:b"]));

        let replaced = base.with_tags_and_origin_tags(tag_set(&["service:web"]), tag_set(&["origin:b"]));
        assert_eq!(
            replaced,
            context_with_host("metric", "host-a", &["service:web"], &["origin:b"])
        );
        assert_ne!(replaced, context_with_origin("metric", &["service:web"], &["origin:b"]));
    }

    #[test]
    fn host_is_preserved_when_context_tags_are_mutated() {
        let expected_with_tag = context_with_host("metric", "host-a", &["env:prod", "service:web"], &["origin:a"]);
        let expected_with_origin = context_with_host("metric", "host-a", &["env:prod"], &["origin:a", "origin:b"]);

        let mut tag_mutated = context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]);
        tag_mutated.mutate_tags(|tags| tags.insert_tag(Tag::from("service:web")));
        assert_eq!(tag_mutated, expected_with_tag);
        assert_eq!(tag_mutated.host(), Some("host-a"));

        let mut origin_mutated = context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]);
        origin_mutated.mutate_origin_tags(|tags| tags.insert_tag(Tag::from("origin:b")));
        assert_eq!(origin_mutated, expected_with_origin);
        assert_eq!(origin_mutated.host(), Some("host-a"));
    }

    #[test]
    fn host_is_preserved_when_tag_mut_view_rekeys_context() {
        let mut ctx = context_with_host(
            "metric",
            "host-a",
            &["env:prod", "service:web"],
            &["origin:a", "origin:b"],
        );
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "env");
        view.retain_origin_tags(|tag| tag.as_str() == "origin:a");
        assert_eq!(view.finish(), 2);

        assert_eq!(ctx, context_with_host("metric", "host-a", &["env:prod"], &["origin:a"]));
        assert_ne!(ctx, context_with_origin("metric", &["env:prod"], &["origin:a"]));
        assert_eq!(ctx.host(), Some("host-a"));
    }

    // --- TagSetMutView ---

    #[test]
    fn mut_view_retain_tags_removes_matching() {
        let mut ctx = Context::from_static_parts("metric", &["env:prod", "service:web", "region:us"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "env");
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert_eq!(ctx.tags().len(), 1);
        assert!(ctx.tags().has_tag("env:prod"));
        assert!(!ctx.tags().has_tag("service:web"));
        assert!(!ctx.tags().has_tag("region:us"));
    }

    #[test]
    fn mut_view_retain_origin_tags_removes_matching() {
        let mut ctx = context_with_origin("metric", &[], &["origin:a", "origin:b", "origin:c"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_origin_tags(|tag| tag.as_str() == "origin:a");
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert_eq!(ctx.origin_tags().len(), 1);
        assert!(ctx.origin_tags().has_tag("origin:a"));
        assert!(!ctx.origin_tags().has_tag("origin:b"));
        assert!(!ctx.origin_tags().has_tag("origin:c"));
    }

    #[test]
    fn mut_view_retain_both_tag_sets() {
        let mut ctx = context_with_origin("metric", &["env:prod", "service:web"], &["origin:a", "origin:b"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "env");
        view.retain_origin_tags(|tag| tag.as_str() == "origin:a");
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert_eq!(ctx.tags().len(), 1);
        assert!(ctx.tags().has_tag("env:prod"));
        assert!(!ctx.tags().has_tag("service:web"));
        assert_eq!(ctx.origin_tags().len(), 1);
        assert!(ctx.origin_tags().has_tag("origin:a"));
        assert!(!ctx.origin_tags().has_tag("origin:b"));
    }

    #[test]
    fn mut_view_retain_all_is_noop() {
        let original = Context::from_static_parts("metric", &["env:prod", "service:web"]);
        let mut ctx = original.clone();
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|_| true);
        let removed = view.finish();

        assert_eq!(removed, 0);
        assert!(ctx.ptr_eq(&original));
    }

    #[test]
    fn mut_view_retain_none_removes_all() {
        let mut ctx = Context::from_static_parts("metric", &["env:prod", "service:web", "region:us"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|_| false);
        let removed = view.finish();

        assert_eq!(removed, 3);
        assert!(ctx.tags().is_empty());
    }

    #[test]
    fn mut_view_finish_returns_correct_count() {
        let mut ctx = context_with_origin("metric", &["a:1", "b:2", "c:3"], &["origin:x", "origin:y"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        // Remove b:2 and c:3 (keep a:1).
        view.retain_tags(|tag| tag.name() == "a");
        // Remove origin:y (keep origin:x).
        view.retain_origin_tags(|tag| tag.as_str() == "origin:x");
        let removed = view.finish();

        assert_eq!(removed, 3);
        assert_eq!(ctx.tags().len(), 1);
        assert_eq!(ctx.origin_tags().len(), 1);
    }

    #[test]
    fn mut_view_equivalent_to_direct_mutate_tags() {
        let base = Context::from_static_parts("metric", &["env:prod", "service:web", "region:us"]);
        let predicate = |tag: &Tag| tag.name() == "env";

        // Path A: direct mutation.
        let mut direct = base.clone();
        direct.mutate_tags(|tags| tags.retain(predicate));

        // Path B: mut view.
        let mut via_view = base.clone();
        let mut state = TagSetMutViewState::new();
        let mut view = via_view.tags_mut_view(&mut state);
        view.retain_tags(predicate);
        view.finish();

        assert_eq!(direct, via_view);
        assert_eq!(direct.tags().len(), via_view.tags().len());
        assert!(via_view.tags().has_tag("env:prod"));
        assert!(!via_view.tags().has_tag("service:web"));
    }

    #[test]
    fn mut_view_equivalent_to_direct_mutate_origin_tags() {
        let base = context_with_origin("metric", &["env:prod"], &["origin:a", "origin:b", "origin:c"]);
        let predicate = |tag: &Tag| tag.as_str() == "origin:a";

        // Path A: direct mutation.
        let mut direct = base.clone();
        direct.mutate_origin_tags(|tags| tags.retain(predicate));

        // Path B: mut view.
        let mut via_view = base.clone();
        let mut state = TagSetMutViewState::new();
        let mut view = via_view.tags_mut_view(&mut state);
        view.retain_origin_tags(predicate);
        view.finish();

        assert_eq!(direct, via_view);
        assert_eq!(direct.origin_tags().len(), via_view.origin_tags().len());
    }

    #[test]
    fn mut_view_does_not_affect_cloned_context() {
        let original = Context::from_static_parts("metric", &["env:prod", "service:web"]);
        let mut mutated = original.clone();
        let mut state = TagSetMutViewState::new();

        let mut view = mutated.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "env");
        view.finish();

        // Original is unchanged.
        assert_eq!(original.tags().len(), 2);
        assert!(original.tags().has_tag("env:prod"));
        assert!(original.tags().has_tag("service:web"));

        // Mutated has only the retained tag.
        assert_eq!(mutated.tags().len(), 1);
        assert!(!original.ptr_eq(&mutated));
    }

    #[test]
    fn mut_view_drop_without_finish_discards_changes() {
        let original = Context::from_static_parts("metric", &["env:prod", "service:web"]);
        let mut ctx = original.clone();
        let mut state = TagSetMutViewState::new();

        {
            let mut view = ctx.tags_mut_view(&mut state);
            view.retain_tags(|_| false); // Flag all for removal.
                                         // Drop without calling finish().
        }

        // Nothing changed.
        assert_eq!(ctx.tags().len(), 2);
        assert!(ctx.ptr_eq(&original));
    }

    #[test]
    fn mut_view_state_reuse_across_operations() {
        let mut state = TagSetMutViewState::new();

        // First operation.
        let mut ctx1 = Context::from_static_parts("metric1", &["a:1", "b:2"]);
        let mut view1 = ctx1.tags_mut_view(&mut state);
        view1.retain_tags(|tag| tag.name() == "a");
        let removed1 = view1.finish();

        assert_eq!(removed1, 1);
        assert_eq!(ctx1.tags().len(), 1);
        assert!(ctx1.tags().has_tag("a:1"));

        // Second operation reusing the same state.
        let mut ctx2 = Context::from_static_parts("metric2", &["x:1", "y:2", "z:3"]);
        let mut view2 = ctx2.tags_mut_view(&mut state);
        view2.retain_tags(|tag| tag.name() == "z");
        let removed2 = view2.finish();

        assert_eq!(removed2, 2);
        assert_eq!(ctx2.tags().len(), 1);
        assert!(ctx2.tags().has_tag("z:3"));
    }

    #[test]
    fn mut_view_retain_tags_with_additions() {
        // Start with a base tag, then add one via mutation to create an overlay.
        let mut ctx = Context::from_static_parts("metric", &["base:tag"]);
        ctx.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("added:tag"));
        });
        assert_eq!(ctx.tags().len(), 2);

        let mut state = TagSetMutViewState::new();
        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "added");
        let removed = view.finish();

        assert_eq!(removed, 1);
        assert_eq!(ctx.tags().len(), 1);
        assert!(ctx.tags().has_tag("added:tag"));
        assert!(!ctx.tags().has_tag("base:tag"));
    }

    #[test]
    fn mut_view_retain_tags_removes_only_additions() {
        let mut ctx = Context::from_static_parts("metric", &["base:tag"]);
        ctx.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("added:tag"));
        });

        let mut state = TagSetMutViewState::new();
        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() == "base");
        let removed = view.finish();

        assert_eq!(removed, 1);
        assert_eq!(ctx.tags().len(), 1);
        assert!(ctx.tags().has_tag("base:tag"));
        assert!(!ctx.tags().has_tag("added:tag"));
    }

    #[test]
    fn mut_view_retain_tags_removes_base_and_additions() {
        let mut ctx = Context::from_static_parts("metric", &["base:tag"]);
        ctx.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("added:tag"));
        });

        let mut state = TagSetMutViewState::new();
        let mut view = ctx.tags_mut_view(&mut state);
        view.retain_tags(|_| false);
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert!(ctx.tags().is_empty());
    }

    #[test]
    fn mut_view_multiple_retain_calls_deduplicates() {
        // Two retain calls that both reject the same addition tag must not panic.
        // Semantics: a tag survives only if ALL predicates accept it.
        let mut ctx = Context::from_static_parts("metric", &["base:tag"]);
        ctx.mutate_tags(|tags| {
            tags.insert_tag(Tag::from("added:a"));
            tags.insert_tag(Tag::from("added:b"));
        });
        assert_eq!(ctx.tags().len(), 3);

        let mut state = TagSetMutViewState::new();
        let mut view = ctx.tags_mut_view(&mut state);
        // First predicate removes "added:a" (keeps base:tag and added:b).
        view.retain_tags(|tag| tag.as_str() != "added:a");
        // Second predicate removes "base:tag" (keeps added:a and added:b).
        // Combined effect: only "added:b" survives both predicates.
        // "added:a" is flagged by both calls -- its duplicate index must be deduplicated.
        view.retain_tags(|tag| tag.name() != "base");
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert_eq!(ctx.tags().len(), 1);
        assert!(ctx.tags().has_tag("added:b"));
    }

    #[test]
    fn mut_view_multiple_retain_origin_calls_deduplicates() {
        let mut ctx = context_with_origin("metric", &[], &["origin:a", "origin:b", "origin:c"]);
        let mut state = TagSetMutViewState::new();

        let mut view = ctx.tags_mut_view(&mut state);
        // Both predicates reject "origin:c".
        view.retain_origin_tags(|tag| tag.as_str() != "origin:c");
        view.retain_origin_tags(|tag| tag.as_str() == "origin:a");
        let removed = view.finish();

        assert_eq!(removed, 2);
        assert_eq!(ctx.origin_tags().len(), 1);
        assert!(ctx.origin_tags().has_tag("origin:a"));
    }

    #[test]
    fn mut_view_multiple_retain_equivalent_to_combined_predicate() {
        let base = Context::from_static_parts("metric", &["env:prod", "service:web", "region:us", "cluster:main"]);

        // Path A: two separate retain calls.
        let mut via_two = base.clone();
        let mut state = TagSetMutViewState::new();
        let mut view = via_two.tags_mut_view(&mut state);
        view.retain_tags(|tag| tag.name() != "region");
        view.retain_tags(|tag| tag.name() != "cluster");
        view.finish();

        // Path B: single combined predicate.
        let mut via_one = base.clone();
        let mut state2 = TagSetMutViewState::new();
        let mut view2 = via_one.tags_mut_view(&mut state2);
        view2.retain_tags(|tag| tag.name() != "region" && tag.name() != "cluster");
        view2.finish();

        // Path C: direct mutation.
        let mut via_direct = base.clone();
        via_direct.mutate_tags(|tags| tags.retain(|tag| tag.name() != "region" && tag.name() != "cluster"));

        assert_eq!(via_two, via_one);
        assert_eq!(via_two, via_direct);
    }

    #[test]
    fn with_tag_sets_mut_mutates_both_and_recomputes_key() {
        // `with_tag_sets_mut` mutates both tag sets and recomputes the context key a single time
        // for the combined change. The observable guarantee is that the result matches a context
        // built from scratch with the combined tags, and matches applying the same two mutations
        // via separate `with_tags`/`with_origin_tags` calls (which would each rehash).
        let mut combined = context_with_origin("metric", &["env:prod"], &["origin:a"]);
        combined.with_tag_sets_mut(|tags, origin_tags| {
            tags.insert_tag(Tag::from("service:web"));
            origin_tags.insert_tag(Tag::from("origin:b"));
        });

        // Both tag sets were updated in the single call.
        assert!(combined.tags().has_tag("env:prod"));
        assert!(combined.tags().has_tag("service:web"));
        assert!(combined.origin_tags().has_tag("origin:a"));
        assert!(combined.origin_tags().has_tag("origin:b"));

        // The recomputed key matches a freshly-built context with the same final state.
        let expected = context_with_origin("metric", &["env:prod", "service:web"], &["origin:a", "origin:b"]);
        assert_eq!(combined, expected);

        // ...and matches applying the two mutations separately (i.e. via two rehashes).
        let separate = context_with_origin("metric", &["env:prod"], &["origin:a"])
            .with_tags(tag_set(&["env:prod", "service:web"]))
            .with_origin_tags(tag_set(&["origin:a", "origin:b"]));
        assert_eq!(combined, separate);
    }

    #[test]
    fn display_renders_name_and_instrumented_tags() {
        // With no tags, only the metric name is rendered.
        assert_eq!(Context::from_static_name("metric").to_string(), "metric");

        // With tags, they are rendered in a brace-delimited, comma-space-separated list in
        // insertion order. Origin tags are intentionally not part of the `Display` output.
        let context = Context::from_static_parts("metric", &["env:prod", "service:web"]);
        assert_eq!(context.to_string(), "metric{env:prod, service:web}");
    }
}