aws_smithy_schema/schema/
trait_map.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use crate::{ShapeId, Trait};
7use std::collections::HashMap;
8
9/// A map of traits keyed by their Shape ID.
10///
11/// This provides efficient lookup of traits during serialization and deserialization.
12#[derive(Debug)]
13pub struct TraitMap {
14    // Wrapped in `Option` because `HashMap::new()` is not `const fn` in stable Rust,
15    // allowing `TraitMap::EMPTY` to be used in const contexts (e.g. prelude schemas).
16    traits: Option<HashMap<ShapeId, Box<dyn Trait>>>,
17}
18
19impl Default for TraitMap {
20    fn default() -> Self {
21        Self::new()
22    }
23}
24
25impl TraitMap {
26    /// An empty trait map for use in const contexts (e.g. prelude schemas).
27    pub const EMPTY: Self = Self { traits: None };
28
29    /// Creates a new empty TraitMap.
30    pub fn new() -> Self {
31        Self {
32            traits: Some(HashMap::new()),
33        }
34    }
35
36    /// Inserts a trait into the map.
37    pub fn insert(&mut self, trait_obj: Box<dyn Trait>) {
38        let id = *trait_obj.trait_id();
39        self.traits
40            .get_or_insert_with(HashMap::new)
41            .insert(id, trait_obj);
42    }
43
44    /// Gets a trait by its Shape ID.
45    pub fn get(&self, id: &ShapeId) -> Option<&dyn Trait> {
46        self.traits.as_ref()?.get(id).map(|t| t.as_ref())
47    }
48
49    /// Returns true if the map contains a trait with the given Shape ID.
50    pub fn contains(&self, id: &ShapeId) -> bool {
51        self.traits.as_ref().is_some_and(|m| m.contains_key(id))
52    }
53
54    /// Returns the number of traits in the map.
55    pub fn len(&self) -> usize {
56        self.traits.as_ref().map_or(0, |m| m.len())
57    }
58
59    /// Returns true if the map is empty.
60    pub fn is_empty(&self) -> bool {
61        self.traits.as_ref().is_none_or(|m| m.is_empty())
62    }
63}