1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175

//! Global schema registry.

use std::{
    alloc::Layout,
    sync::atomic::{AtomicU32, Ordering::SeqCst},
};

use append_only_vec::AppendOnlyVec;
use bones_utils::*;

use crate::prelude::*;

/// A unique identifier for a schema registered in the [`SCHEMA_REGISTRY`].
#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
pub struct SchemaId {
    id: u32,
}

// Note: The schema type is here in the registry module to prevent modification of registered
// schemas by other modules. The idea is that once a schema is registered, it is unchangable and
// "certified" so to speak.
#[doc(hidden)]
/// A schema registered with the [`SCHEMA_REGISTRY`].
///
/// ## Known Limitations
///
/// Currently there isn't a known-safe way to construct a schema for a recursive struct. For
/// example, this struct is troublesome:
///
/// ```rust
/// struct Data {
///     others: Vec<Data>,
/// }
/// ```
///
/// This is because nested schemas are required to be a `&'static Schema`, and it would probalby
/// require unsafe code to create a schema that references itself.
///
/// If this is a problem for your use-case, please open an issue. We would like to remove this
/// limitation or find a suitable workaround in the future.
#[derive(Deref, Clone, Debug)]
pub struct Schema {
    id: SchemaId,
    #[deref]
    data: SchemaData,
    layout: Layout,
    field_offsets: &'static [(Option<String>, usize)],
}

impl PartialEq for Schema {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}
impl Eq for Schema {}

impl Schema {
    /// Get the registered, unique ID of the [`Schema`].
    #[inline]
    pub fn id(&self) -> SchemaId {
        self.id
    }

    /// Get a static reference to the [`Schema`] that was registered.
    #[inline]
    pub fn schema(&self) -> &SchemaData {
        &self.data
    }

    /// Get the [`Layout`] of the [`Schema`].
    #[inline]
    pub fn layout(&self) -> Layout {
        self.layout
    }

    /// If this schema represents a struct, this returns the list of fields, with the names of the
    /// fields, and their byte offsets from the beginning of the struct.
    #[inline]
    pub fn field_offsets(&self) -> &'static [(Option<String>, usize)] {
        self.field_offsets
    }

    /// Helper function to make sure that this schema matches another or return a
    /// [`SchemaMismatchError`].
    pub fn ensure_match(&self, other: &Self) -> Result<(), SchemaMismatchError> {
        if self == other {
            Ok(())
        } else {
            Err(SchemaMismatchError)
        }
    }
}

/// A schema registry that alloates [`SchemaId`]s for [`SchemaData`]s and returns a registered
/// [`&'static Schema`][Schema].
pub struct SchemaRegistry {
    next_id: AtomicU32,
    /// The registered schemas.
    pub schemas: AppendOnlyVec<Schema>,
}

impl SchemaRegistry {
    /// Register a schema with the registry.
    #[track_caller]
    pub fn register(&self, schema_data: SchemaData) -> &Schema {
        // Allocate a new schema ID
        let id = SchemaId {
            id: self.next_id.fetch_add(1, SeqCst),
        };
        assert_ne!(id.id, u32::MAX, "Exhausted all {} schema IDs", u32::MAX);

        // Compute the schema layout info so we can cache it with the Schema.
        let SchemaLayoutInfo {
            layout,
            field_offsets,
        } = schema_data.kind.compute_layout_info();

        // Leak the field offsets to get static references
        let field_offsets: Box<_> = field_offsets
            .into_iter()
            .map(|(name, offset)| (name.map(|n| n.to_string()), offset))
            .collect();
        let field_offsets = Box::leak(field_offsets);

        // Create the schema struct.
        let schema = Schema {
            id,
            data: schema_data,
            layout,
            field_offsets,
        };

        // Insert the schema into the registry.
        let idx = self.schemas.push(schema);

        &self.schemas[idx]
    }
}

/// Global [`SchemaRegistry`] used to register [`SchemaData`]s and produce [`Schema`]s.
pub static SCHEMA_REGISTRY: SchemaRegistry = SchemaRegistry {
    next_id: AtomicU32::new(0),
    schemas: AppendOnlyVec::new(),
};

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn registry_smoke() {
        let mut schemas = Vec::new();
        for i in 0..100 {
            let data = SchemaData {
                name: format!("data{i}").into(),
                full_name: format!("data{i}").into(),
                kind: SchemaKind::Primitive(Primitive::U8),
                type_data: default(),
                type_id: None,
                clone_fn: None,
                drop_fn: None,
                default_fn: None,
                hash_fn: None,
                eq_fn: None,
            };

            let schema = SCHEMA_REGISTRY.register(data.clone());
            schemas.push(schema);
        }

        for (i, schema) in schemas.iter().enumerate() {
            assert_eq!(schema.data.name, format!("data{i}"));
        }
    }
}