Skip to main content

spacetimedb_schema/
relation.rs

1use crate::def::error::{RelationError, TypeError};
2use crate::table_name::TableName;
3use core::fmt;
4use core::hash::Hash;
5use derive_more::From;
6use spacetimedb_data_structures::map::HashSet;
7use spacetimedb_lib::db::auth::{StAccess, StTableType};
8use spacetimedb_primitives::{ColId, ColList, ColSet, Constraints, TableId};
9use spacetimedb_sats::algebraic_value::AlgebraicValue;
10use spacetimedb_sats::satn::Satn;
11use spacetimedb_sats::{algebraic_type, AlgebraicType};
12use std::collections::BTreeMap;
13use std::sync::Arc;
14
15#[derive(Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
16pub struct FieldName {
17    pub table: TableId,
18    pub col: ColId,
19}
20
21impl FieldName {
22    pub fn new(table: TableId, col: ColId) -> Self {
23        Self { table, col }
24    }
25
26    pub fn table(&self) -> TableId {
27        self.table
28    }
29
30    pub fn field(&self) -> ColId {
31        self.col
32    }
33}
34
35// TODO(perf): Remove `Clone` derivation.
36#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
37pub enum ColExpr {
38    Col(ColId),
39    Value(AlgebraicValue),
40}
41
42impl ColExpr {
43    /// Returns a borrowed version of `ColExpr`.
44    pub fn borrowed(&self) -> ColExprRef<'_> {
45        match self {
46            Self::Col(x) => ColExprRef::Col(*x),
47            Self::Value(x) => ColExprRef::Value(x),
48        }
49    }
50}
51
52impl fmt::Debug for FieldName {
53    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
54        fmt::Display::fmt(self, f)
55    }
56}
57
58impl fmt::Display for FieldName {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        write!(f, "table#{}.col#{}", self.table, self.col)
61    }
62}
63
64impl fmt::Display for ColExpr {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        match self {
67            ColExpr::Col(x) => write!(f, "{x}"),
68            ColExpr::Value(x) => write!(f, "{}", x.to_satn()),
69        }
70    }
71}
72
73/// A borrowed version of `FieldExpr`.
74#[derive(Clone, Copy)]
75pub enum ColExprRef<'a> {
76    Col(ColId),
77    Value(&'a AlgebraicValue),
78}
79
80// TODO(perf): Remove `Clone` derivation.
81#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
82pub struct Column {
83    pub field: FieldName,
84    pub algebraic_type: AlgebraicType,
85}
86
87impl Column {
88    pub fn new(field: FieldName, algebraic_type: AlgebraicType) -> Self {
89        Self { field, algebraic_type }
90    }
91}
92
93#[derive(Debug, PartialEq, Eq, Hash, Clone)]
94pub struct Header {
95    pub table_id: TableId,
96    pub table_name: TableName,
97    pub fields: Vec<Column>,
98    pub constraints: BTreeMap<ColList, Constraints>,
99}
100
101impl Header {
102    /// Create a new header.
103    ///
104    /// `uncombined_constraints` will be normalized using [`combine_constraints`].
105    pub fn new(
106        table_id: TableId,
107        table_name: TableName,
108        fields: Vec<Column>,
109        uncombined_constraints: impl IntoIterator<Item = (ColList, Constraints)>,
110    ) -> Self {
111        Self {
112            table_id,
113            table_name,
114            fields,
115            constraints: combine_constraints(uncombined_constraints),
116        }
117    }
118
119    /// Equivalent to what [`Clone::clone`] would do.
120    ///
121    /// `Header` intentionally does not implement `Clone`,
122    /// as we can't afford to clone it in normal execution paths.
123    /// However, we don't care about performance in error paths,
124    /// and we need to embed owned `Header`s in error objects to report useful messages.
125    pub fn clone_for_error(&self) -> Self {
126        Header::new(
127            self.table_id,
128            self.table_name.clone(),
129            self.fields.clone(),
130            self.constraints.clone(),
131        )
132    }
133
134    /// Finds the index of the column with a matching `FieldName`.
135    pub fn column_pos(&self, col: FieldName) -> Option<ColId> {
136        self.fields.iter().position(|f| f.field == col).map(Into::into)
137    }
138
139    pub fn column_pos_or_err(&self, col: FieldName) -> Result<ColId, RelationError> {
140        self.column_pos(col)
141            .ok_or_else(|| RelationError::FieldNotFound(self.clone_for_error(), col))
142    }
143
144    pub fn field_name(&self, col: FieldName) -> Option<(ColId, FieldName)> {
145        self.column_pos(col).map(|id| (id, self.fields[id.idx()].field))
146    }
147
148    /// Copy the [Constraints] that are referenced in the list of `for_columns`
149    fn retain_constraints(&self, for_columns: &ColList) -> BTreeMap<ColList, Constraints> {
150        // Copy the constraints of the selected columns and retain the multi-column ones...
151        self.constraints
152            .iter()
153            // Keep constraints with a col list where at least one col is in `for_columns`.
154            .filter(|(cols, _)| cols.iter().any(|c| for_columns.contains(c)))
155            .map(|(cols, constraints)| (cols.clone(), *constraints))
156            .collect()
157    }
158
159    pub fn has_constraint(&self, field: ColId, constraint: Constraints) -> bool {
160        self.constraints
161            .iter()
162            .any(|(col, ct)| col.contains(field) && ct.contains(&constraint))
163    }
164
165    /// Project the [ColExpr]s & the [Constraints] that referenced them
166    pub fn project(&self, cols: &[ColExpr]) -> Result<Self, RelationError> {
167        let mut fields = Vec::with_capacity(cols.len());
168        let mut to_keep = ColList::with_capacity(cols.len() as _);
169
170        for (pos, col) in cols.iter().enumerate() {
171            match col {
172                ColExpr::Col(col) => {
173                    to_keep.push(*col);
174                    fields.push(self.fields[col.idx()].clone());
175                }
176                ColExpr::Value(val) => {
177                    // TODO: why should this field name be relevant?
178                    // We should generate immediate names instead.
179                    let field = FieldName::new(self.table_id, pos.into());
180                    let ty = val.type_of().ok_or_else(|| {
181                        RelationError::TypeInference(field, TypeError::CannotInferType { value: val.clone() })
182                    })?;
183                    fields.push(Column::new(field, ty));
184                }
185            }
186        }
187
188        let constraints = self.retain_constraints(&to_keep);
189
190        Ok(Self::new(self.table_id, self.table_name.clone(), fields, constraints))
191    }
192
193    /// Project the ourself onto the `ColList`, keeping constraints that reference the columns in the ColList.
194    /// Does not change `ColIDs`.
195    pub fn project_col_list(&self, cols: &ColList) -> Self {
196        let mut fields = Vec::with_capacity(cols.len() as usize);
197
198        for col in cols.iter() {
199            fields.push(self.fields[col.idx()].clone());
200        }
201
202        let constraints = self.retain_constraints(cols);
203        Self::new(self.table_id, self.table_name.clone(), fields, constraints)
204    }
205
206    /// Adds the fields &  [Constraints] from `right` to this [`Header`],
207    /// renaming duplicated fields with a counter like `a, a => a, a0`.
208    pub fn extend(&self, right: &Self) -> Self {
209        // Increase the positions of the columns in `right.constraints`, adding the count of fields on `left`
210        let mut constraints = self.constraints.clone();
211        let len_lhs = self.fields.len() as u16;
212        constraints.extend(right.constraints.iter().map(|(cols, c)| {
213            let cols = cols.iter().map(|col| ColId(col.0 + len_lhs)).collect::<ColList>();
214            (cols, *c)
215        }));
216
217        let mut fields = self.fields.clone();
218        fields.extend(right.fields.iter().cloned());
219
220        Self::new(self.table_id, self.table_name.clone(), fields, constraints)
221    }
222}
223
224/// Combine constraints.
225/// The result is a map from `ColList` to `Constraints`.
226/// In particular, it includes all indexes, and tells you which of them are unique.
227/// The result MAY contain multiple `ColList`s with the same columns in different orders.
228/// This corresponds to differently-ordered indices.
229///
230/// Unique constraints are considered logically unordered. Information from them will
231/// be propagated to all indices that contain the same columns.
232pub fn combine_constraints(
233    uncombined: impl IntoIterator<Item = (ColList, Constraints)>,
234) -> BTreeMap<ColList, Constraints> {
235    let mut constraints = BTreeMap::new();
236    for (col_list, constraint) in uncombined {
237        let slot = constraints.entry(col_list).or_insert(Constraints::unset());
238        *slot = slot.push(constraint);
239    }
240
241    let mut uniques: HashSet<ColSet> = HashSet::default();
242    for (col_list, constraint) in &constraints {
243        if constraint.has_unique() {
244            uniques.insert(col_list.into());
245        }
246    }
247
248    for (cols, constraint) in constraints.iter_mut() {
249        if uniques.contains(&ColSet::from(cols)) {
250            *constraint = constraint.push(Constraints::unique());
251        }
252    }
253
254    constraints
255}
256
257impl fmt::Display for Header {
258    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
259        write!(f, "[")?;
260        for (pos, col) in self.fields.iter().enumerate() {
261            write!(
262                f,
263                "{}: {}",
264                col.field,
265                algebraic_type::fmt::fmt_algebraic_type(&col.algebraic_type)
266            )?;
267
268            if pos + 1 < self.fields.len() {
269                write!(f, ", ")?;
270            }
271        }
272        write!(f, "]")
273    }
274}
275
276#[derive(Debug, Clone, Eq, PartialEq, Hash)]
277pub struct DbTable {
278    pub head: Arc<Header>,
279    pub table_id: TableId,
280    pub table_type: StTableType,
281    pub table_access: StAccess,
282}
283
284impl DbTable {
285    pub fn new(head: Arc<Header>, table_id: TableId, table_type: StTableType, table_access: StAccess) -> Self {
286        Self {
287            head,
288            table_id,
289            table_type,
290            table_access,
291        }
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use spacetimedb_primitives::col_list;
299
300    /// Build a [Header] using the initial `start_pos` as the column position for the [Constraints]
301    fn head(id: impl Into<TableId>, name: &str, fields: (ColId, ColId), start_pos: u16) -> Header {
302        let pos_lhs = start_pos;
303        let pos_rhs = start_pos + 1;
304
305        let ct = vec![
306            (ColId(pos_lhs).into(), Constraints::indexed()),
307            (ColId(pos_rhs).into(), Constraints::identity()),
308            (col_list![pos_lhs, pos_rhs], Constraints::primary_key()),
309            (col_list![pos_rhs, pos_lhs], Constraints::unique()),
310        ];
311
312        let id = id.into();
313        let fields = [fields.0, fields.1].map(|col| Column::new(FieldName::new(id, col), AlgebraicType::I8));
314        Header::new(id, TableName::for_test(name), fields.into(), ct)
315    }
316
317    #[test]
318    fn test_project() {
319        let a = 0.into();
320        let b = 1.into();
321
322        let head = head(0, "t1", (a, b), 0);
323        let new = head.project(&[] as &[ColExpr]).unwrap();
324
325        let mut empty = head.clone_for_error();
326        empty.fields.clear();
327        empty.constraints.clear();
328        assert_eq!(empty, new);
329
330        let all = head.clone_for_error();
331        let new = head.project(&[a, b].map(ColExpr::Col)).unwrap();
332        assert_eq!(all, new);
333
334        let mut first = head.clone_for_error();
335        first.fields.pop();
336        first.constraints = first.retain_constraints(&a.into());
337        let new = head.project(&[a].map(ColExpr::Col)).unwrap();
338        assert_eq!(first, new);
339
340        let mut second = head.clone_for_error();
341        second.fields.remove(0);
342        second.constraints = second.retain_constraints(&b.into());
343        let new = head.project(&[b].map(ColExpr::Col)).unwrap();
344        assert_eq!(second, new);
345    }
346
347    #[test]
348    fn test_extend() {
349        let t1 = 0.into();
350        let t2: TableId = 1.into();
351        let a = 0.into();
352        let b = 1.into();
353        let c = 0.into();
354        let d = 1.into();
355
356        let head_lhs = head(t1, "t1", (a, b), 0);
357        let head_rhs = head(t2, "t2", (c, d), 0);
358
359        let new = head_lhs.extend(&head_rhs);
360
361        let lhs = new.project(&[a, b].map(ColExpr::Col)).unwrap();
362        assert_eq!(head_lhs, lhs);
363
364        let mut head_rhs = head(t2, "t2", (c, d), 2);
365        head_rhs.table_id = t1;
366        head_rhs.table_name = head_lhs.table_name.clone();
367        let rhs = new.project(&[2, 3].map(ColId).map(ColExpr::Col)).unwrap();
368        assert_eq!(head_rhs, rhs);
369    }
370
371    #[test]
372    fn test_combine_constraints() {
373        let raw = vec![
374            (col_list![0], Constraints::indexed()),
375            (col_list![0], Constraints::unique()),
376            (col_list![1], Constraints::identity()),
377            (col_list![1, 0], Constraints::primary_key()),
378            (col_list![0, 1], Constraints::unique()),
379            (col_list![2], Constraints::indexed()),
380            (col_list![3], Constraints::unique()),
381        ];
382        let expected = vec![
383            (col_list![0], Constraints::indexed().push(Constraints::unique())),
384            (col_list![1], Constraints::identity()),
385            (col_list![1, 0], Constraints::primary_key().push(Constraints::unique())),
386            (col_list![0, 1], Constraints::unique()),
387            (col_list![2], Constraints::indexed()),
388            (col_list![3], Constraints::unique()),
389        ]
390        .into_iter()
391        .collect::<BTreeMap<_, _>>();
392        assert_eq!(combine_constraints(raw), expected);
393    }
394}