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///
13/// # Cross-lifetime lookup
14///
15/// Internal storage is keyed by `ShapeId<'static>` because traits come from
16/// codegen-emitted schemas (themselves `'static`). The public `get` and
17/// `contains` accessors accept `&ShapeId<'_>` of any lifetime — runtime-
18/// built shape IDs (e.g. parsed from a wire-format `__type` field) are
19/// looked up against the `'static`-keyed table via `ShapeId`'s
20/// `Borrow<str>` impl, which keys on the fully qualified name.
21#[derive(Debug)]
22pub struct TraitMap {
23    // Wrapped in `Option` because `HashMap::new()` is not `const fn` in stable Rust,
24    // allowing `TraitMap::EMPTY` to be used in const contexts (e.g. prelude schemas).
25    traits: Option<HashMap<ShapeId<'static>, Box<dyn Trait>>>,
26}
27
28impl Default for TraitMap {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl TraitMap {
35    /// An empty trait map for use in const contexts (e.g. prelude schemas).
36    pub const EMPTY: Self = Self { traits: None };
37
38    /// Creates a new empty TraitMap.
39    pub fn new() -> Self {
40        Self {
41            traits: Some(HashMap::new()),
42        }
43    }
44
45    /// Inserts a trait into the map.
46    pub fn insert(&mut self, trait_obj: Box<dyn Trait>) {
47        let id = trait_obj.trait_id().clone();
48        self.traits
49            .get_or_insert_with(HashMap::new)
50            .insert(id, trait_obj);
51    }
52
53    /// Gets a trait by its Shape ID.
54    ///
55    /// Accepts a [`ShapeId`] of any lifetime — lookup is by fully
56    /// qualified name. A `ShapeId<'static>` constructed via the
57    /// `shape_id!` macro and a runtime-built `ShapeId<'_>` with the
58    /// same FQN resolve to the same entry.
59    pub fn get(&self, id: &ShapeId<'_>) -> Option<&dyn Trait> {
60        // `id.as_str()` returns `&str`; the map's `Borrow<str>` impl
61        // (via `ShapeId`'s `Borrow<str>`) makes this a direct lookup.
62        self.traits.as_ref()?.get(id.as_str()).map(|t| t.as_ref())
63    }
64
65    /// Looks up a trait by its fully qualified name (e.g.
66    /// `"smithy.api#sensitive"`).
67    ///
68    /// Equivalent to [`Self::get`] when the caller already has the FQN
69    /// as a string slice — for example, after extracting the
70    /// `__type` field from a JSON document.
71    ///
72    /// # Example
73    ///
74    /// ```
75    /// use aws_smithy_schema::traits::JsonNameTrait;
76    /// use aws_smithy_schema::{shape_id, ShapeId, TraitMap};
77    ///
78    /// let mut traits = TraitMap::new();
79    /// traits.insert(Box::new(JsonNameTrait::new("custom_name")));
80    ///
81    /// // Cross-lifetime: a static-lifetime ShapeId from `shape_id!` and a
82    /// // runtime-built ShapeId<'_> with the same FQN both resolve via
83    /// // `Borrow<str>`. The companion `_fqn` accessor takes `&str` directly.
84    /// assert!(traits.get(&shape_id!("smithy.api", "jsonName")).is_some());
85    /// assert!(traits.get_fqn("smithy.api#jsonName").is_some());
86    /// ```
87    pub fn get_fqn(&self, fqn: &str) -> Option<&dyn Trait> {
88        self.traits.as_ref()?.get(fqn).map(|t| t.as_ref())
89    }
90
91    /// Returns true if the map contains a trait with the given Shape ID.
92    ///
93    /// Accepts a [`ShapeId`] of any lifetime; see [`Self::get`].
94    pub fn contains(&self, id: &ShapeId<'_>) -> bool {
95        self.traits
96            .as_ref()
97            .is_some_and(|m| m.contains_key(id.as_str()))
98    }
99
100    /// Returns true if the map contains a trait with the given fully
101    /// qualified name. See [`Self::get_fqn`].
102    pub fn contains_fqn(&self, fqn: &str) -> bool {
103        self.traits.as_ref().is_some_and(|m| m.contains_key(fqn))
104    }
105
106    /// Returns the number of traits in the map.
107    pub fn len(&self) -> usize {
108        self.traits.as_ref().map_or(0, |m| m.len())
109    }
110
111    /// Returns true if the map is empty.
112    pub fn is_empty(&self) -> bool {
113        self.traits.as_ref().is_none_or(|m| m.is_empty())
114    }
115}
aws_smithy_schema/schema/trait_map.rs

aws_smithy_schema/schema/
trait_map.rs
