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#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
36pub enum ColExpr {
37 Col(ColId),
38 Value(AlgebraicValue),
39}
40
41impl ColExpr {
42 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#[derive(Clone, Copy)]
74pub enum ColExprRef<'a> {
75 Col(ColId),
76 Value(&'a AlgebraicValue),
77}
78
79#[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 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 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 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 fn retain_constraints(&self, for_columns: &ColList) -> BTreeMap<ColList, Constraints> {
149 self.constraints
151 .iter()
152 .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 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 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 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 pub fn extend(&self, right: &Self) -> Self {
208 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
223pub 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 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}