spacetimedb_lib/
relation.rs

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