Skip to main content

ixa/network/
edge.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{LazyLock, Mutex};
4
5use crate::entity::{Entity, EntityId};
6
7#[derive(Copy, Debug, PartialEq)]
8/// An edge in the network graph. Edges are directed, so the
9/// source person is implicit.
10pub struct Edge<E: Entity, ET: EdgeType<E>> {
11    /// The person this edge points to.
12    pub neighbor: EntityId<E>,
13    /// The weight associated with the edge.
14    pub weight: f32,
15    /// An inner value defined by type `T`. Often a ZST.
16    pub inner: ET,
17}
18
19// Generics prevent the compiler from "seeing" that `Edge` always satisfies these
20// traits if they are derived.
21impl<E: Entity, ET: EdgeType<E>> Clone for Edge<E, ET> {
22    fn clone(&self) -> Self {
23        Self {
24            neighbor: self.neighbor,
25            weight: self.weight,
26            inner: self.inner.clone(),
27        }
28    }
29}
30
31pub trait EdgeType<E: Entity>: Clone + 'static {
32    #[must_use]
33    fn name() -> &'static str {
34        let full = std::any::type_name::<Self>();
35        // Splitting any string, including an empty one, always yields at least one segment.
36        full.rsplit("::").next().unwrap()
37    }
38
39    /// The index of this item in the owner, which is initialized globally per type
40    /// upon first access. We explicitly initialize this in a `ctor` in order to know
41    /// how many [`EdgeType<E>`] types exist globally when we construct any `NetworkStore<E>`.
42    #[must_use]
43    fn id() -> usize;
44}
45
46/// A map from Entity ID to a count of the edge types already associated with the entity. The value for the key is
47/// equivalent to the next edge type ID that will be assigned to the next edge type that requests an ID. Each `Entity`
48/// type has its own series of increasing edge type IDs.
49static NEXT_EDGE_TYPE_ID_BY_ENTITY: LazyLock<Mutex<HashMap<usize, usize>>> =
50    LazyLock::new(|| Mutex::new(HashMap::default()));
51
52/// Returns the number of registered edge types for the entity type `E`.
53#[must_use]
54pub fn get_registered_edge_type_count<E: Entity>() -> usize {
55    let map = NEXT_EDGE_TYPE_ID_BY_ENTITY.lock().unwrap();
56    *map.get(&E::id()).unwrap_or(&0)
57}
58
59/// Adds a new edge type to the registry. The job of this method is to create whatever
60/// "singleton" data/metadata is associated with the [`EdgeType`] if it doesn't already
61/// exist, which in this case is only the value of `EdgeType::id()`.
62pub fn add_to_edge_type_to_registry<E: Entity, ET: EdgeType<E>>() {
63    let _ = ET::id();
64}
65
66/// Encapsulates the synchronization logic for initializing an [`EdgeType<E>`]'s ID.
67///
68/// Acquires a global lock on the next available edge type ID for the given entity type `E`,
69/// but only increments it if we successfully initialize the provided ID. The ID of an
70/// edge type is
71/// assigned at runtime but only once per type. It's possible for a single
72/// type to attempt to initialize its index multiple times from different threads,
73/// which is why all this synchronization is required. However, the overhead
74/// is negligible, as this initialization only happens once upon first access.
75///
76/// In fact, for our use case we know we are calling this function
77/// once for each type in each `EdgeType<E>`'s `ctor` function, which
78/// should be the only time this method is ever called for the type.
79#[must_use]
80pub fn initialize_edge_type_id<E: Entity>(edge_type_id: &AtomicUsize) -> usize {
81    // Acquire a global lock.
82    let mut guard = NEXT_EDGE_TYPE_ID_BY_ENTITY.lock().unwrap();
83    let candidate = guard.entry(E::id()).or_insert_with(|| 0);
84
85    // Try to claim the candidate index. Here we guard against the potential race condition that
86    // another instance of this plugin in another thread just initialized the index prior to us
87    // obtaining the lock. If the index has been initialized beneath us, we do not update
88    // NEXT_EDGE_TYPE_ID_BY_ENTITY, we just return the value `edge_type_id` was initialized to.
89    // For a justification of the data ordering, see:
90    //     https://github.com/CDCgov/ixa/pull/477#discussion_r2244302872
91    match edge_type_id.compare_exchange(usize::MAX, *candidate, Ordering::AcqRel, Ordering::Acquire)
92    {
93        Ok(_) => {
94            // We won the race — increment the global next edge-type ID and return the new ID.
95            *candidate += 1;
96            *candidate - 1
97        }
98        Err(existing) => {
99            // Another thread beat us — don’t increment the global next edge-type ID,
100            // just return the existing one.
101            existing
102        }
103    }
104}