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