reifydb_macro_impl/
key.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
5
6use crate::{generate::compile_error, key_conformance::expand_tests};
7
8pub struct KeyField {
9	pub name: String,
10	pub column: KeyColumn,
11	pub ty: String,
12}
13
14#[derive(Clone, Copy)]
15pub enum KeyColumn {
16	U8,
17	U16,
18	U32,
19	U64,
20	U128,
21	RowNumber,
22	GroupId,
23	Blob16,
24	TableId,
25	ColumnId,
26	ObjectId,
27	StorageId,
28	FlowId,
29	FlowEdgeId,
30	OperatorId,
31	HandlerId,
32	NamespaceId,
33	SumTypeId,
34	SequenceId,
35	EpochSeconds,
36	IdentityId,
37	ViewId,
38	SeriesId,
39	SinkId,
40	SourceId,
41	QueueId,
42	RingBufferId,
43	BindingId,
44	DictionaryId,
45	ProcedureId,
46	DateTime,
47	Partition,
48	ColumnPropertyId,
49	RowShapeFingerprint,
50	RelationshipId,
51	MigrationId,
52	MigrationEventId,
53	PrimaryKeyId,
54	ColumnSnapshotId,
55	IndexId,
56	OptionU8,
57	ReprU8,
58}
59
60impl KeyColumn {
61	fn width(self) -> usize {
62		match self {
63			KeyColumn::U8 | KeyColumn::ReprU8 => 1,
64			KeyColumn::U16 => 2,
65			KeyColumn::U32 => 4,
66			KeyColumn::U64
67			| KeyColumn::RowNumber
68			| KeyColumn::TableId
69			| KeyColumn::ColumnId
70			| KeyColumn::FlowId
71			| KeyColumn::FlowEdgeId
72			| KeyColumn::OperatorId
73			| KeyColumn::HandlerId
74			| KeyColumn::NamespaceId
75			| KeyColumn::SumTypeId
76			| KeyColumn::SequenceId
77			| KeyColumn::EpochSeconds
78			| KeyColumn::ViewId
79			| KeyColumn::SeriesId
80			| KeyColumn::SinkId
81			| KeyColumn::SourceId
82			| KeyColumn::QueueId
83			| KeyColumn::RingBufferId
84			| KeyColumn::BindingId
85			| KeyColumn::DictionaryId
86			| KeyColumn::ProcedureId
87			| KeyColumn::DateTime
88			| KeyColumn::ColumnPropertyId
89			| KeyColumn::RowShapeFingerprint
90			| KeyColumn::RelationshipId
91			| KeyColumn::MigrationId
92			| KeyColumn::MigrationEventId
93			| KeyColumn::ColumnSnapshotId
94			| KeyColumn::PrimaryKeyId
95			| KeyColumn::IndexId => 8,
96			KeyColumn::ObjectId | KeyColumn::StorageId => 9,
97			KeyColumn::OptionU8 => 2,
98			KeyColumn::U128 | KeyColumn::Blob16 | KeyColumn::Partition | KeyColumn::IdentityId => 16,
99			KeyColumn::GroupId => 24,
100		}
101	}
102
103	fn field_exprs(self, field: &str) -> Vec<String> {
104		let u = |width: &str, expr: String| vec![format!("Field::UDesc(Width::{width}, {expr} as u128)")];
105		match self {
106			KeyColumn::U8 => u("U8", format!("self.{field}")),
107			KeyColumn::U16 => u("U16", format!("self.{field}")),
108			KeyColumn::U32 => u("U32", format!("self.{field}")),
109			KeyColumn::U64 => u("U64", format!("self.{field}")),
110			KeyColumn::U128 => u("U128", format!("self.{field}")),
111			KeyColumn::ReprU8 => u("U8", format!("self.{field} as u8")),
112			KeyColumn::RowNumber => u("U64", format!("self.{field}.0")),
113			KeyColumn::Partition => u("U128", format!("self.{field}.0")),
114			KeyColumn::ProcedureId => u("U64", format!("*self.{field}")),
115			KeyColumn::EpochSeconds => u("U64", format!("self.{field}.seconds()")),
116			KeyColumn::DateTime => u("U64", format!("self.{field}.to_nanos()")),
117			KeyColumn::RowShapeFingerprint | KeyColumn::IndexId => {
118				u("U64", format!("self.{field}.as_u64()"))
119			}
120			KeyColumn::GroupId => vec![format!(
121				"Field::BytesDesc(ByteEncoding::Fixed, ::std::borrow::Cow::Borrowed(self.{field}.as_bytes()))"
122			)],
123			KeyColumn::IdentityId => vec![format!(
124				"Field::BytesDesc(ByteEncoding::Escaped, ::std::borrow::Cow::Borrowed(self.{field}.as_bytes()))"
125			)],
126			KeyColumn::Blob16 => vec![format!(
127				"Field::RawAsc(RawEncoding::Verbatim, ::std::borrow::Cow::Borrowed(&self.{field}))"
128			)],
129			KeyColumn::ObjectId | KeyColumn::StorageId => vec![
130				format!("Field::UAsc(Width::U8, self.{field}.type_tag() as u128)"),
131				format!("Field::UDesc(Width::U64, self.{field}.as_u64() as u128)"),
132			],
133			KeyColumn::OptionU8 => vec![
134				format!("Field::UDesc(Width::U8, self.{field}.is_some() as u128)"),
135				format!("Field::UDesc(Width::U8, self.{field}.unwrap_or(0u8) as u128)"),
136			],
137			KeyColumn::TableId
138			| KeyColumn::ColumnId
139			| KeyColumn::FlowId
140			| KeyColumn::FlowEdgeId
141			| KeyColumn::OperatorId
142			| KeyColumn::HandlerId
143			| KeyColumn::NamespaceId
144			| KeyColumn::SumTypeId
145			| KeyColumn::SequenceId
146			| KeyColumn::ViewId
147			| KeyColumn::SeriesId
148			| KeyColumn::SinkId
149			| KeyColumn::SourceId
150			| KeyColumn::QueueId
151			| KeyColumn::RingBufferId
152			| KeyColumn::BindingId
153			| KeyColumn::DictionaryId
154			| KeyColumn::ColumnPropertyId
155			| KeyColumn::RelationshipId
156			| KeyColumn::MigrationId
157			| KeyColumn::MigrationEventId
158			| KeyColumn::ColumnSnapshotId
159			| KeyColumn::PrimaryKeyId => u("U64", format!("self.{field}.0")),
160		}
161	}
162
163	fn encode_stmt(self, field: &str) -> String {
164		match self {
165			KeyColumn::U8 => format!("serializer.extend_u8(self.{field});"),
166			KeyColumn::ReprU8 => format!("serializer.extend_u8(self.{field} as u8);"),
167			KeyColumn::U16 => format!("serializer.extend_u16(self.{field});"),
168			KeyColumn::U32 => format!("serializer.extend_u32(self.{field});"),
169			KeyColumn::U64 => format!("serializer.extend_u64(self.{field});"),
170			KeyColumn::U128 => format!("serializer.extend_u128(self.{field});"),
171			KeyColumn::RowNumber => format!("serializer.extend_u64(self.{field}.0);"),
172			KeyColumn::GroupId => format!("serializer.extend_fixed(*self.{field}.as_bytes());"),
173			KeyColumn::Blob16 => format!("serializer.extend_raw(&self.{field});"),
174			KeyColumn::IdentityId => format!("serializer.extend_identity_id(&self.{field});"),
175			KeyColumn::TableId
176			| KeyColumn::ColumnId
177			| KeyColumn::FlowId
178			| KeyColumn::FlowEdgeId
179			| KeyColumn::OperatorId
180			| KeyColumn::HandlerId
181			| KeyColumn::NamespaceId
182			| KeyColumn::SumTypeId
183			| KeyColumn::SequenceId
184			| KeyColumn::ViewId
185			| KeyColumn::SeriesId
186			| KeyColumn::SinkId
187			| KeyColumn::SourceId
188			| KeyColumn::QueueId
189			| KeyColumn::RingBufferId
190			| KeyColumn::BindingId
191			| KeyColumn::DictionaryId
192			| KeyColumn::ColumnPropertyId
193			| KeyColumn::RelationshipId
194			| KeyColumn::MigrationId
195			| KeyColumn::MigrationEventId
196			| KeyColumn::ColumnSnapshotId
197			| KeyColumn::PrimaryKeyId => format!("serializer.extend_u64(self.{field}.0);"),
198			KeyColumn::ProcedureId => format!("serializer.extend_u64(*self.{field});"),
199			KeyColumn::EpochSeconds => format!("serializer.extend_u64(self.{field}.seconds());"),
200			KeyColumn::DateTime => format!("serializer.extend_datetime(&self.{field});"),
201			KeyColumn::RowShapeFingerprint => format!("serializer.extend_u64(self.{field}.as_u64());"),
202			KeyColumn::Partition => format!("serializer.extend_u128(self.{field}.0);"),
203			KeyColumn::IndexId => format!("serializer.extend_u64(self.{field}.as_u64());"),
204			KeyColumn::ObjectId | KeyColumn::StorageId => {
205				format!("serializer.extend_object_id(self.{field});")
206			}
207			KeyColumn::OptionU8 => format!(
208				"match self.{field} {{ Some(v) => {{ serializer.extend_u8(1u8).extend_u8(v); }} \
209				 None => {{ serializer.extend_u8(0u8).extend_u8(0u8); }} }}"
210			),
211		}
212	}
213
214	fn decode_expr(self, field_type: &str) -> String {
215		match self {
216			KeyColumn::U8 => "de.read_u8().ok()?".to_string(),
217			KeyColumn::ReprU8 => format!("{field_type}::try_from(de.read_u8().ok()?).ok()?"),
218			KeyColumn::U16 => "de.read_u16().ok()?".to_string(),
219			KeyColumn::U32 => "de.read_u32().ok()?".to_string(),
220			KeyColumn::U64 => "de.read_u64().ok()?".to_string(),
221			KeyColumn::U128 => "de.read_u128().ok()?".to_string(),
222			KeyColumn::RowNumber => "RowNumber(de.read_u64().ok()?)".to_string(),
223			KeyColumn::GroupId => "GroupId::from_bytes(de.read_fixed().ok()?)".to_string(),
224			KeyColumn::Blob16 => "{ let bytes = de.read_raw(16).ok()?; let mut buf = [0u8; 16]; \
225				 buf.copy_from_slice(bytes); buf }"
226				.to_string(),
227			KeyColumn::IdentityId => "de.read_identity_id().ok()?".to_string(),
228			KeyColumn::TableId
229			| KeyColumn::ColumnId
230			| KeyColumn::FlowId
231			| KeyColumn::FlowEdgeId
232			| KeyColumn::OperatorId
233			| KeyColumn::HandlerId
234			| KeyColumn::NamespaceId
235			| KeyColumn::SumTypeId
236			| KeyColumn::SequenceId
237			| KeyColumn::ViewId
238			| KeyColumn::SeriesId
239			| KeyColumn::SinkId
240			| KeyColumn::SourceId
241			| KeyColumn::QueueId
242			| KeyColumn::RingBufferId
243			| KeyColumn::BindingId
244			| KeyColumn::DictionaryId
245			| KeyColumn::ColumnPropertyId
246			| KeyColumn::RelationshipId
247			| KeyColumn::MigrationId
248			| KeyColumn::MigrationEventId
249			| KeyColumn::ColumnSnapshotId
250			| KeyColumn::PrimaryKeyId => format!("{field_type}(de.read_u64().ok()?)"),
251			KeyColumn::ProcedureId => "ProcedureId::from_raw(de.read_u64().ok()?)".to_string(),
252			KeyColumn::EpochSeconds => "EpochSeconds::new(de.read_u64().ok()?)".to_string(),
253			KeyColumn::DateTime => "de.read_datetime().ok()?".to_string(),
254			KeyColumn::RowShapeFingerprint => "RowShapeFingerprint::new(de.read_u64().ok()?)".to_string(),
255			KeyColumn::Partition => "Partition(de.read_u128().ok()?)".to_string(),
256			KeyColumn::IndexId => "IndexId::Primary(PrimaryKeyId(de.read_u64().ok()?))".to_string(),
257			KeyColumn::ObjectId => "de.read_object_id().ok()?".to_string(),
258			KeyColumn::StorageId => "StorageId::from_object(de.read_object_id().ok()?)?".to_string(),
259			KeyColumn::OptionU8 => "match de.read_u8().ok()? { 1u8 => Some(de.read_u8().ok()?), 0u8 => { \
260				de.read_u8().ok()?; None }, _ => return None }"
261				.to_string(),
262		}
263	}
264}
265
266pub fn derive_key(input: TokenStream) -> TokenStream {
267	let tokens: Vec<TokenTree> = input.into_iter().collect();
268	let mut iter = tokens.iter().peekable();
269
270	let mut tag: Option<String> = None;
271	while let Some(TokenTree::Punct(p)) = iter.peek() {
272		if p.as_char() != '#' {
273			break;
274		}
275		iter.next();
276		match iter.next() {
277			Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
278				if let Some(found) = parse_key_attribute(g) {
279					tag = Some(found);
280				}
281			}
282			_ => return compile_error("expected an attribute after '#'"),
283		}
284	}
285
286	let tag = match tag {
287		Some(tag) => tag,
288		None => return compile_error("KeyCodec requires #[key(tag = Variant)] naming its KeyTag"),
289	};
290
291	if let Some(TokenTree::Ident(i)) = iter.peek()
292		&& *i == "pub"
293	{
294		iter.next();
295		if let Some(TokenTree::Group(g)) = iter.peek()
296			&& g.delimiter() == Delimiter::Parenthesis
297		{
298			iter.next();
299		}
300	}
301
302	match iter.next() {
303		Some(TokenTree::Ident(i)) if *i == "struct" => {}
304		_ => return compile_error("KeyCodec can only be derived for structs"),
305	}
306
307	let name = match iter.next() {
308		Some(TokenTree::Ident(i)) => i.to_string(),
309		_ => return compile_error("expected struct name"),
310	};
311
312	let body = match iter.next() {
313		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Brace => g.clone(),
314		Some(TokenTree::Punct(p)) if p.as_char() == '<' => {
315			return compile_error("KeyCodec cannot be derived for a generic struct");
316		}
317		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => {
318			return compile_error("KeyCodec requires named fields, so a tuple struct has no field order");
319		}
320		_ => return compile_error("expected struct body"),
321	};
322
323	let fields = match parse_fields(&body) {
324		Ok(fields) => fields,
325		Err(err) => return err,
326	};
327
328	expand(&name, &tag, &fields)
329}
330
331fn parse_key_attribute(group: &Group) -> Option<String> {
332	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
333	let mut iter = tokens.iter();
334
335	match iter.next() {
336		Some(TokenTree::Ident(i)) if *i == "key" => {}
337		_ => return None,
338	}
339
340	let inner = match iter.next() {
341		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => g.clone(),
342		_ => return None,
343	};
344
345	let inner_tokens: Vec<TokenTree> = inner.stream().into_iter().collect();
346	match inner_tokens.as_slice() {
347		[TokenTree::Ident(tag_ident), TokenTree::Punct(eq), TokenTree::Ident(variant)]
348			if *tag_ident == "tag" && eq.as_char() == '=' =>
349		{
350			Some(variant.to_string())
351		}
352		_ => None,
353	}
354}
355
356fn parse_fields(body: &Group) -> Result<Vec<KeyField>, TokenStream> {
357	let tokens: Vec<TokenTree> = body.stream().into_iter().collect();
358	let mut iter = tokens.iter().peekable();
359	let mut fields = Vec::new();
360
361	while iter.peek().is_some() {
362		let mut repr_u8 = false;
363		while let Some(TokenTree::Punct(p)) = iter.peek() {
364			if p.as_char() == '#' {
365				iter.next();
366				if let Some(TokenTree::Group(g)) = iter.peek() {
367					if g.delimiter() == Delimiter::Bracket && is_repr_u8_attribute(g) {
368						repr_u8 = true;
369					}
370					iter.next();
371				}
372			} else {
373				break;
374			}
375		}
376		if iter.peek().is_none() {
377			break;
378		}
379
380		if let Some(TokenTree::Ident(i)) = iter.peek()
381			&& *i == "pub"
382		{
383			iter.next();
384			if let Some(TokenTree::Group(g)) = iter.peek()
385				&& g.delimiter() == Delimiter::Parenthesis
386			{
387				iter.next();
388			}
389		}
390
391		let field_name = match iter.next() {
392			Some(TokenTree::Ident(i)) => i.to_string(),
393			_ => return Err(compile_error("expected field name")),
394		};
395
396		match iter.next() {
397			Some(TokenTree::Punct(p)) if p.as_char() == ':' => {}
398			_ => return Err(compile_error("expected ':' after field name")),
399		}
400
401		let mut ty_tokens: Vec<TokenTree> = Vec::new();
402		loop {
403			match iter.peek() {
404				Some(TokenTree::Punct(p)) if p.as_char() == ',' => {
405					iter.next();
406					break;
407				}
408				Some(t) => {
409					ty_tokens.push((*t).clone());
410					iter.next();
411				}
412				None => break,
413			}
414		}
415
416		if ty_tokens.is_empty() {
417			return Err(compile_error(&format!("field '{}' has no type", field_name)));
418		}
419
420		let ty = render(&ty_tokens);
421		let column = if repr_u8 {
422			KeyColumn::ReprU8
423		} else {
424			match column_type(&ty_tokens) {
425				Some(column) => column,
426				None => {
427					return Err(compile_error(&format!(
428						"field '{}' has type '{}', which has no flat key column type",
429						field_name, ty
430					)));
431				}
432			}
433		};
434
435		fields.push(KeyField {
436			name: field_name,
437			column,
438			ty,
439		});
440	}
441
442	Ok(fields)
443}
444
445fn is_repr_u8_attribute(group: &Group) -> bool {
446	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
447	let mut iter = tokens.iter();
448
449	match iter.next() {
450		Some(TokenTree::Ident(i)) if *i == "key" => {}
451		_ => return false,
452	}
453
454	let inner = match iter.next() {
455		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => g.clone(),
456		_ => return false,
457	};
458
459	let inner_tokens: Vec<TokenTree> = inner.stream().into_iter().collect();
460	matches!(
461		inner_tokens.as_slice(),
462		[TokenTree::Ident(repr), TokenTree::Punct(eq), TokenTree::Ident(width)]
463			if *repr == "repr" && eq.as_char() == '=' && *width == "u8"
464	)
465}
466
467fn column_type(tokens: &[TokenTree]) -> Option<KeyColumn> {
468	if let [TokenTree::Group(group)] = tokens {
469		return (group.delimiter() == Delimiter::Bracket && is_byte_array_16(group))
470			.then_some(KeyColumn::Blob16);
471	}
472
473	if is_option_u8(tokens) {
474		return Some(KeyColumn::OptionU8);
475	}
476
477	let unqualified = strip_path(tokens);
478	if unqualified.len() != 1 {
479		return None;
480	}
481
482	let head = match unqualified.first() {
483		Some(TokenTree::Ident(i)) => i.to_string(),
484		_ => return None,
485	};
486
487	match head.as_str() {
488		"u8" => Some(KeyColumn::U8),
489		"u16" => Some(KeyColumn::U16),
490		"u32" => Some(KeyColumn::U32),
491		"u64" => Some(KeyColumn::U64),
492		"u128" => Some(KeyColumn::U128),
493		"RowNumber" => Some(KeyColumn::RowNumber),
494		"GroupId" => Some(KeyColumn::GroupId),
495		"TableId" => Some(KeyColumn::TableId),
496		"ColumnId" => Some(KeyColumn::ColumnId),
497		"ObjectId" => Some(KeyColumn::ObjectId),
498		"StorageId" => Some(KeyColumn::StorageId),
499		"QueueId" => Some(KeyColumn::QueueId),
500		"RingBufferId" => Some(KeyColumn::RingBufferId),
501		"SeriesId" => Some(KeyColumn::SeriesId),
502		"DateTime" => Some(KeyColumn::DateTime),
503		"Partition" => Some(KeyColumn::Partition),
504		"FlowId" => Some(KeyColumn::FlowId),
505		"FlowEdgeId" => Some(KeyColumn::FlowEdgeId),
506		"OperatorId" => Some(KeyColumn::OperatorId),
507		"HandlerId" => Some(KeyColumn::HandlerId),
508		"NamespaceId" => Some(KeyColumn::NamespaceId),
509		"SumTypeId" => Some(KeyColumn::SumTypeId),
510		"SequenceId" => Some(KeyColumn::SequenceId),
511		"EpochSeconds" => Some(KeyColumn::EpochSeconds),
512		"IdentityId" => Some(KeyColumn::IdentityId),
513		"AuthenticationId" | "IdentityAttributeId" | "RoleId" | "PolicyId" | "TokenId" => Some(KeyColumn::U64),
514		"ViewId" => Some(KeyColumn::ViewId),
515		"SinkId" => Some(KeyColumn::SinkId),
516		"SourceId" => Some(KeyColumn::SourceId),
517		"BindingId" => Some(KeyColumn::BindingId),
518		"DictionaryId" => Some(KeyColumn::DictionaryId),
519		"ProcedureId" => Some(KeyColumn::ProcedureId),
520		"ColumnPropertyId" => Some(KeyColumn::ColumnPropertyId),
521		"RowShapeFingerprint" => Some(KeyColumn::RowShapeFingerprint),
522		"IndexId" => Some(KeyColumn::IndexId),
523		"RelationshipId" => Some(KeyColumn::RelationshipId),
524		"MigrationId" => Some(KeyColumn::MigrationId),
525		"MigrationEventId" => Some(KeyColumn::MigrationEventId),
526		"PrimaryKeyId" => Some(KeyColumn::PrimaryKeyId),
527		"ColumnSnapshotId" => Some(KeyColumn::ColumnSnapshotId),
528		_ => None,
529	}
530}
531
532fn is_option_u8(tokens: &[TokenTree]) -> bool {
533	matches!(
534		tokens,
535		[TokenTree::Ident(o), TokenTree::Punct(lt), TokenTree::Ident(u), TokenTree::Punct(gt)]
536			if *o == "Option" && lt.as_char() == '<' && *u == "u8" && gt.as_char() == '>'
537	)
538}
539
540fn is_byte_array_16(group: &Group) -> bool {
541	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
542	match tokens.as_slice() {
543		[TokenTree::Ident(element), TokenTree::Punct(semi), TokenTree::Literal(len)] => {
544			*element == "u8" && semi.as_char() == ';' && len.to_string() == "16"
545		}
546		_ => false,
547	}
548}
549
550fn strip_path(tokens: &[TokenTree]) -> &[TokenTree] {
551	let mut at = 0;
552	if is_colon(tokens.first()) && is_colon(tokens.get(1)) {
553		at = 2;
554	}
555	while matches!(tokens.get(at), Some(TokenTree::Ident(_)))
556		&& is_colon(tokens.get(at + 1))
557		&& is_colon(tokens.get(at + 2))
558	{
559		at += 3;
560	}
561	&tokens[at..]
562}
563
564fn is_colon(token: Option<&TokenTree>) -> bool {
565	matches!(token, Some(TokenTree::Punct(p)) if p.as_char() == ':')
566}
567
568fn render(tokens: &[TokenTree]) -> String {
569	tokens.iter().map(|token| token.to_string()).collect()
570}
571
572fn expand(name: &str, tag: &str, fields: &[KeyField]) -> TokenStream {
573	let capacity: usize = 1 + fields.iter().map(|f| f.column.width()).sum::<usize>();
574
575	let mut encode_body = String::new();
576	for field in fields {
577		encode_body.push_str("\t\t");
578		encode_body.push_str(&field.column.encode_stmt(&field.name));
579		encode_body.push('\n');
580	}
581
582	let mut decode_body = String::new();
583	for field in fields {
584		decode_body.push_str(&format!("\t\t\t{}: {},\n", field.name, field.column.decode_expr(&field.ty)));
585	}
586
587	let mut out = String::new();
588	out.push_str(&format!("impl {name} {{\n"));
589	out.push_str(&format!("\tpub const TAG: KeyTag = KeyTag::{tag};\n\n"));
590	out.push_str("\tpub fn encode(&self) -> EncodedKey {\n");
591	out.push_str(&format!(
592		"\t\tlet mut serializer = ::reifydb_codec::key::serializer::KeySerializer::with_capacity({capacity});\n"
593	));
594	out.push_str("\t\tserializer.extend_u8(Self::TAG as u8);\n");
595	out.push_str(&encode_body);
596	out.push_str("\t\tserializer.to_encoded_key()\n\t}\n\n");
597	out.push_str("\tpub fn decode(key: &EncodedKey) -> Option<Self> {\n");
598	out.push_str(
599		"\t\tlet mut de = ::reifydb_codec::key::deserializer::KeyDeserializer::from_bytes(key.as_slice());\n",
600	);
601	out.push_str("\t\tlet found: KeyTag = de.read_u8().ok()?.try_into().ok()?;\n");
602	out.push_str("\t\tif found != Self::TAG {\n\t\t\treturn None;\n\t\t}\n");
603	out.push_str(&format!("\t\tlet decoded = Self {{\n{decode_body}\t\t}};\n"));
604	out.push_str("\t\tif !de.is_empty() {\n\t\t\treturn None;\n\t\t}\n");
605	out.push_str("\t\tSome(decoded)\n\t}\n}\n\n");
606
607	out.push_str(&format!("#[automatically_derived]\nimpl KeyFields for {name} {{\n"));
608	out.push_str("\tfn fields(&self) -> ::smallvec::SmallVec<[Field<'_>; 6]> {\n");
609	out.push_str("\t\t::smallvec::smallvec![\n");
610	for field in fields {
611		for expr in field.column.field_exprs(&field.name) {
612			out.push_str(&format!("\t\t\t{expr},\n"));
613		}
614	}
615	out.push_str("\t\t]\n\t}\n}\n\n");
616	out.push_str(&expand_tests(name, fields));
617
618	out.parse().expect("derived key impl must be valid Rust")
619}
620
621#[cfg(test)]
622mod tests {
623	use super::derive_key;
624
625	fn expand(source: &str) -> String {
626		derive_key(source.parse().unwrap()).to_string()
627	}
628
629	#[test]
630	fn missing_attribute_is_rejected() {
631		let out = expand("struct RowKey { row: u64 }");
632		assert!(out.contains("compile_error"), "{out}");
633		assert!(out.contains("key"), "{out}");
634	}
635
636	#[test]
637	fn a_well_formed_struct_expands_without_an_error() {
638		let out = expand("#[key(tag = Row)] struct RowKey { table: u64, row: RowNumber }");
639		assert!(!out.contains("compile_error"), "{out}");
640		assert!(out.contains("impl RowKey { pub const TAG"), "{out}");
641		assert!(out.contains("KeyTag :: Row"), "{out}");
642	}
643
644	#[test]
645	fn fields_are_encoded_in_declaration_order() {
646		let out = expand("#[key(tag = Row)] struct RowKey { table: u64, row: RowNumber }");
647		let table = out.find("self . table").expect("table field encoded");
648		let row = out.find("self . row").expect("row field encoded");
649		assert!(table < row, "{out}");
650	}
651
652	#[test]
653	fn a_repr_u8_field_encodes_the_discriminant_and_decodes_through_try_from() {
654		// without the attribute a repr(u8) enum has no column type at all, and a decode that skipped
655		// try_from would hand back a discriminant the enum never declared
656		let out = expand("#[key(tag = SystemVersion)] struct K { #[key(repr = u8)] version: SystemVersion }");
657		assert!(!out.contains("compile_error"), "{out}");
658		assert!(out.contains("extend_u8 (self . version as u8)"), "{out}");
659		assert!(out.contains("SystemVersion :: try_from"), "{out}");
660	}
661
662	#[test]
663	fn a_repr_u8_field_still_rejects_a_type_the_attribute_was_not_put_on() {
664		let out = expand("#[key(tag = SystemVersion)] struct K { version: SystemVersion }");
665		assert!(out.contains("compile_error"), "{out}");
666	}
667
668	#[test]
669	fn an_unrecognised_field_type_is_rejected() {
670		let out = expand("#[key(tag = Row)] struct RowKey { at: String }");
671		assert!(out.contains("compile_error"), "{out}");
672		assert!(out.contains("at"), "{out}");
673		assert!(out.contains("String"), "{out}");
674	}
675
676	#[test]
677	fn a_tuple_struct_is_rejected() {
678		let out = expand("#[key(tag = Row)] struct RowKey(u64);");
679		assert!(out.contains("compile_error"), "{out}");
680	}
681
682	#[test]
683	fn a_generic_struct_is_rejected() {
684		let out = expand("#[key(tag = Row)] struct RowKey<T> { row: T }");
685		assert!(out.contains("compile_error"), "{out}");
686	}
687
688	#[test]
689	fn a_byte_array_field_maps_to_blob16() {
690		let out = expand("#[key(tag = Row)] struct RowKey { blob: [u8; 16] }");
691		assert!(!out.contains("compile_error"), "{out}");
692		assert!(out.contains("extend_raw"));
693	}
694
695	#[test]
696	fn a_byte_array_of_the_wrong_width_is_rejected() {
697		let out = expand("#[key(tag = Row)] struct RowKey { blob: [u8; 32] }");
698		assert!(out.contains("compile_error"), "{out}");
699	}
700
701	#[test]
702	fn table_id_and_column_id_encode_via_their_wrapped_u64() {
703		let out = expand("#[key(tag = Table)] struct TableKey { table: TableId, column: ColumnId }");
704		assert!(!out.contains("compile_error"), "{out}");
705		assert!(out.contains("self . table . 0"), "{out}");
706		assert!(out.contains("TableId (de . read_u64 () . ok () ?)"), "{out}");
707		assert!(out.contains("ColumnId (de . read_u64 () . ok () ?)"), "{out}");
708	}
709
710	#[test]
711	fn an_option_u8_field_encodes_as_a_fixed_two_byte_presence_and_value_pair() {
712		let out = expand("#[key(tag = Row)] struct RowKey { tag: Option<u8> }");
713		assert!(!out.contains("compile_error"), "{out}");
714		assert!(out.contains("with_capacity (3)"), "{out}: width must stay fixed at 2 bytes for None and Some");
715		assert!(out.contains("Some (v) => { serializer . extend_u8 (1u8) . extend_u8 (v) ; }"), "{out}");
716		assert!(out.contains("None => { serializer . extend_u8 (0u8) . extend_u8 (0u8) ; }"), "{out}");
717		assert!(out.contains("1u8 => Some (de . read_u8 () . ok () ?)"), "{out}");
718		assert!(out.contains("0u8 => { de . read_u8 () . ok () ?; None }"), "{out}");
719	}
720
721	#[test]
722	fn object_id_and_storage_id_encode_via_the_catalog_ext() {
723		let out = expand("#[key(tag = Row)] struct RowKey { storage: StorageId, object: ObjectId }");
724		assert!(!out.contains("compile_error"), "{out}");
725		assert!(out.contains("extend_object_id (self . storage)"), "{out}");
726		assert!(out.contains("extend_object_id (self . object)"), "{out}");
727		assert!(out.contains("StorageId :: from_object (de . read_object_id () . ok () ?) ?"), "{out}");
728		assert!(out.contains("de . read_object_id () . ok () ?"), "{out}");
729	}
730}
reifydb_macro_impl/key.rs

reifydb_macro_impl/
key.rs
