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    traits: HashMap<ShapeId, Box<dyn Trait>>,
15}
16
17impl Default for TraitMap {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl TraitMap {
24    // TODO(schema) Is there a reasonable with_capacity size for this?
25    /// Creates a new empty TraitMap.
26    pub fn new() -> Self {
27        Self {
28            traits: HashMap::new(),
29        }
30    }
31
32    /// Creates a TraitMap with zero allocated space for Prelude Schemas.
33    pub(crate) fn empty() -> Self {
34        Self {
35            traits: HashMap::with_capacity(0),
36        }
37    }
38
39    /// Inserts a trait into the map.
40    pub fn insert(&mut self, trait_obj: Box<dyn Trait>) {
41        let id = trait_obj.trait_id().clone();
42        self.traits.insert(id, trait_obj);
43    }
44
45    /// Gets a trait by its Shape ID.
46    pub fn get(&self, id: &ShapeId) -> Option<&dyn Trait> {
47        self.traits.get(id).map(|t| t.as_ref())
48    }
49
50    /// Returns true if the map contains a trait with the given Shape ID.
51    pub fn contains(&self, id: &ShapeId) -> bool {
52        self.traits.contains_key(id)
53    }
54
55    /// Returns an iterator over all traits.
56    pub fn iter(&self) -> impl Iterator<Item = (&ShapeId, &dyn Trait)> {
57        self.traits.iter().map(|(s, t)| (s, t.as_ref()))
58    }
59
60    /// Returns the number of traits in the map.
61    pub fn len(&self) -> usize {
62        self.traits.len()
63    }
64
65    /// Returns true if the map is empty.
66    pub fn is_empty(&self) -> bool {
67        self.traits.is_empty()
68    }
69}