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	ColumnBlockId,
55	IndexId,
56	OptionU8,
57	ReprU8,
58	Str,
59	DisplayStr,
60	EscapedBytes,
61	KeyspaceId,
62	RawTail,
63	TaggedIndexId,
64	VarintU128,
65	SortRun,
66	AscRowNumber,
67}
68
69impl KeyColumn {
70	fn width(self) -> usize {
71		match self {
72			KeyColumn::U8 | KeyColumn::ReprU8 | KeyColumn::KeyspaceId => 1,
73			KeyColumn::Str
74			| KeyColumn::DisplayStr
75			| KeyColumn::EscapedBytes
76			| KeyColumn::RawTail
77			| KeyColumn::SortRun => 16,
78			KeyColumn::TaggedIndexId => 9,
79			KeyColumn::VarintU128 => 17,
80			KeyColumn::AscRowNumber => 8,
81			KeyColumn::U16 => 2,
82			KeyColumn::U32 => 4,
83			KeyColumn::U64
84			| KeyColumn::RowNumber
85			| KeyColumn::TableId
86			| KeyColumn::ColumnId
87			| KeyColumn::FlowId
88			| KeyColumn::FlowEdgeId
89			| KeyColumn::OperatorId
90			| KeyColumn::HandlerId
91			| KeyColumn::NamespaceId
92			| KeyColumn::SumTypeId
93			| KeyColumn::SequenceId
94			| KeyColumn::EpochSeconds
95			| KeyColumn::ViewId
96			| KeyColumn::SeriesId
97			| KeyColumn::SinkId
98			| KeyColumn::SourceId
99			| KeyColumn::QueueId
100			| KeyColumn::RingBufferId
101			| KeyColumn::BindingId
102			| KeyColumn::DictionaryId
103			| KeyColumn::ProcedureId
104			| KeyColumn::DateTime
105			| KeyColumn::ColumnPropertyId
106			| KeyColumn::RowShapeFingerprint
107			| KeyColumn::RelationshipId
108			| KeyColumn::MigrationId
109			| KeyColumn::MigrationEventId
110			| KeyColumn::ColumnBlockId
111			| KeyColumn::PrimaryKeyId
112			| KeyColumn::IndexId => 8,
113			KeyColumn::ObjectId | KeyColumn::StorageId => 9,
114			KeyColumn::OptionU8 => 2,
115			KeyColumn::U128 | KeyColumn::Blob16 | KeyColumn::Partition | KeyColumn::IdentityId => 16,
116			KeyColumn::GroupId => 24,
117		}
118	}
119
120	fn field_exprs(self, field: &str) -> Vec<String> {
121		let u = |width: &str, expr: String| vec![format!("Field::UDesc(Width::{width}, {expr} as u128)")];
122		let cow = |expr: String| format!("::std::borrow::Cow::Borrowed({expr})");
123		match self {
124			KeyColumn::Str => vec![format!(
125				"Field::BytesDesc(ByteEncoding::Escaped, {})",
126				cow(format!("self.{field}.as_ref().as_bytes()"))
127			)],
128			KeyColumn::DisplayStr => vec![format!(
129				"Field::BytesDesc(ByteEncoding::Escaped, ::std::borrow::Cow::Owned(self.{field}.to_string().into_bytes()))"
130			)],
131			KeyColumn::EscapedBytes => vec![format!(
132				"Field::BytesDesc(ByteEncoding::Escaped, {})",
133				cow(format!("self.{field}.as_slice()"))
134			)],
135			KeyColumn::KeyspaceId => u("U8", format!("self.{field}.0")),
136			KeyColumn::RawTail => {
137				vec![format!(
138					"Field::RawAsc(RawEncoding::Verbatim, {})",
139					cow(format!("self.{field}.as_slice()"))
140				)]
141			}
142			KeyColumn::TaggedIndexId => vec![
143				format!("Field::UAsc(Width::U8, index_tag(&self.{field}) as u128)"),
144				format!("Field::UDesc(Width::U64, self.{field}.as_u64() as u128)"),
145			],
146			KeyColumn::VarintU128 => u("Varint", format!("self.{field}")),
147			KeyColumn::SortRun => {
148				vec![format!(
149					"Field::RawAsc(RawEncoding::SortRun, {})",
150					cow(format!("self.{field}.as_slice()"))
151				)]
152			}
153			KeyColumn::AscRowNumber => vec![format!("Field::UAsc(Width::U64, self.{field}.0.0 as u128)")],
154			KeyColumn::U8 => u("U8", format!("self.{field}")),
155			KeyColumn::U16 => u("U16", format!("self.{field}")),
156			KeyColumn::U32 => u("U32", format!("self.{field}")),
157			KeyColumn::U64 => u("U64", format!("self.{field}")),
158			KeyColumn::U128 => u("U128", format!("self.{field}")),
159			KeyColumn::ReprU8 => u("U8", format!("self.{field} as u8")),
160			KeyColumn::RowNumber => u("U64", format!("self.{field}.0")),
161			KeyColumn::Partition => u("U128", format!("self.{field}.0")),
162			KeyColumn::ProcedureId => u("U64", format!("*self.{field}")),
163			KeyColumn::EpochSeconds => u("U64", format!("self.{field}.seconds()")),
164			KeyColumn::DateTime => u("U64", format!("self.{field}.to_order()")),
165			KeyColumn::RowShapeFingerprint | KeyColumn::IndexId => {
166				u("U64", format!("self.{field}.as_u64()"))
167			}
168			KeyColumn::GroupId => vec![format!(
169				"Field::BytesDesc(ByteEncoding::Fixed, ::std::borrow::Cow::Borrowed(self.{field}.as_bytes()))"
170			)],
171			KeyColumn::IdentityId => vec![format!(
172				"Field::BytesDesc(ByteEncoding::Escaped, ::std::borrow::Cow::Borrowed(self.{field}.as_bytes()))"
173			)],
174			KeyColumn::Blob16 => vec![format!(
175				"Field::RawAsc(RawEncoding::Verbatim, ::std::borrow::Cow::Borrowed(&self.{field}))"
176			)],
177			KeyColumn::ObjectId | KeyColumn::StorageId => vec![
178				format!("Field::UAsc(Width::U8, self.{field}.type_tag() as u128)"),
179				format!("Field::UDesc(Width::U64, self.{field}.as_u64() as u128)"),
180			],
181			KeyColumn::OptionU8 => vec![
182				format!("Field::UDesc(Width::U8, self.{field}.is_some() as u128)"),
183				format!("Field::UDesc(Width::U8, self.{field}.unwrap_or(0u8) as u128)"),
184			],
185			KeyColumn::TableId
186			| KeyColumn::ColumnId
187			| KeyColumn::FlowId
188			| KeyColumn::FlowEdgeId
189			| KeyColumn::OperatorId
190			| KeyColumn::HandlerId
191			| KeyColumn::NamespaceId
192			| KeyColumn::SumTypeId
193			| KeyColumn::SequenceId
194			| KeyColumn::ViewId
195			| KeyColumn::SeriesId
196			| KeyColumn::SinkId
197			| KeyColumn::SourceId
198			| KeyColumn::QueueId
199			| KeyColumn::RingBufferId
200			| KeyColumn::BindingId
201			| KeyColumn::DictionaryId
202			| KeyColumn::ColumnPropertyId
203			| KeyColumn::RelationshipId
204			| KeyColumn::MigrationId
205			| KeyColumn::MigrationEventId
206			| KeyColumn::ColumnBlockId
207			| KeyColumn::PrimaryKeyId => u("U64", format!("self.{field}.0")),
208		}
209	}
210
211	fn encode_stmt(self, field: &str) -> String {
212		match self {
213			KeyColumn::Str => format!("serializer.extend_str(&self.{field});"),
214			KeyColumn::DisplayStr => format!("serializer.extend_str(self.{field}.to_string());"),
215			KeyColumn::EscapedBytes => format!("serializer.extend_bytes(&self.{field});"),
216			KeyColumn::KeyspaceId => format!("serializer.extend_u8(self.{field}.0);"),
217			KeyColumn::RawTail => format!("serializer.extend_raw(self.{field}.as_slice());"),
218			KeyColumn::TaggedIndexId => format!("serializer.extend_index_id(self.{field});"),
219			KeyColumn::VarintU128 => format!("serializer.extend_u128_varint(self.{field});"),
220			KeyColumn::SortRun => format!("extend_sort_run(&mut serializer, &self.{field});"),
221			KeyColumn::AscRowNumber => format!("serializer.extend_raw(&self.{field}.0.0.to_be_bytes());"),
222			KeyColumn::U8 => format!("serializer.extend_u8(self.{field});"),
223			KeyColumn::ReprU8 => format!("serializer.extend_u8(self.{field} as u8);"),
224			KeyColumn::U16 => format!("serializer.extend_u16(self.{field});"),
225			KeyColumn::U32 => format!("serializer.extend_u32(self.{field});"),
226			KeyColumn::U64 => format!("serializer.extend_u64(self.{field});"),
227			KeyColumn::U128 => format!("serializer.extend_u128(self.{field});"),
228			KeyColumn::RowNumber => format!("serializer.extend_u64(self.{field}.0);"),
229			KeyColumn::GroupId => format!("serializer.extend_fixed(*self.{field}.as_bytes());"),
230			KeyColumn::Blob16 => format!("serializer.extend_raw(&self.{field});"),
231			KeyColumn::IdentityId => format!("serializer.extend_identity_id(&self.{field});"),
232			KeyColumn::TableId
233			| KeyColumn::ColumnId
234			| KeyColumn::FlowId
235			| KeyColumn::FlowEdgeId
236			| KeyColumn::OperatorId
237			| KeyColumn::HandlerId
238			| KeyColumn::NamespaceId
239			| KeyColumn::SumTypeId
240			| KeyColumn::SequenceId
241			| KeyColumn::ViewId
242			| KeyColumn::SeriesId
243			| KeyColumn::SinkId
244			| KeyColumn::SourceId
245			| KeyColumn::QueueId
246			| KeyColumn::RingBufferId
247			| KeyColumn::BindingId
248			| KeyColumn::DictionaryId
249			| KeyColumn::ColumnPropertyId
250			| KeyColumn::RelationshipId
251			| KeyColumn::MigrationId
252			| KeyColumn::MigrationEventId
253			| KeyColumn::ColumnBlockId
254			| KeyColumn::PrimaryKeyId => format!("serializer.extend_u64(self.{field}.0);"),
255			KeyColumn::ProcedureId => format!("serializer.extend_u64(*self.{field});"),
256			KeyColumn::EpochSeconds => format!("serializer.extend_u64(self.{field}.seconds());"),
257			KeyColumn::DateTime => format!("serializer.extend_datetime(&self.{field});"),
258			KeyColumn::RowShapeFingerprint => format!("serializer.extend_u64(self.{field}.as_u64());"),
259			KeyColumn::Partition => format!("serializer.extend_u128(self.{field}.0);"),
260			KeyColumn::IndexId => format!("serializer.extend_u64(self.{field}.as_u64());"),
261			KeyColumn::ObjectId | KeyColumn::StorageId => {
262				format!("serializer.extend_object_id(self.{field});")
263			}
264			KeyColumn::OptionU8 => format!(
265				"match self.{field} {{ Some(v) => {{ serializer.extend_u8(1u8).extend_u8(v); }} \
266				 None => {{ serializer.extend_u8(0u8).extend_u8(0u8); }} }}"
267			),
268		}
269	}
270
271	fn decode_expr(self, field_type: &str) -> String {
272		match self {
273			KeyColumn::Str => format!("{field_type}(de.read_str().ok()?)"),
274			KeyColumn::DisplayStr => {
275				format!("<{field_type} as ::core::str::FromStr>::from_str(&de.read_str().ok()?).ok()?")
276			}
277			KeyColumn::EscapedBytes => format!("{field_type}::new(de.read_bytes().ok()?)"),
278			KeyColumn::KeyspaceId => "KeyspaceId(de.read_u8().ok()?)".to_string(),
279			KeyColumn::RawTail if field_type == "Vec<u8>" => {
280				"de.read_raw(de.remaining()).ok()?.to_vec()".to_string()
281			}
282			KeyColumn::RawTail => format!("{field_type}::new(de.read_raw(de.remaining()).ok()?)"),
283			KeyColumn::TaggedIndexId => "de.read_index_id().ok()?".to_string(),
284			KeyColumn::VarintU128 => "de.read_u128_varint().ok()?".to_string(),
285			KeyColumn::SortRun => "read_sort_run(&mut de)?".to_string(),
286			KeyColumn::AscRowNumber => "read_row_tail(&mut de)?".to_string(),
287			KeyColumn::U8 => "de.read_u8().ok()?".to_string(),
288			KeyColumn::ReprU8 => format!("{field_type}::try_from(de.read_u8().ok()?).ok()?"),
289			KeyColumn::U16 => "de.read_u16().ok()?".to_string(),
290			KeyColumn::U32 => "de.read_u32().ok()?".to_string(),
291			KeyColumn::U64 => "de.read_u64().ok()?".to_string(),
292			KeyColumn::U128 => "de.read_u128().ok()?".to_string(),
293			KeyColumn::RowNumber => "RowNumber(de.read_u64().ok()?)".to_string(),
294			KeyColumn::GroupId => "GroupId::from_bytes(de.read_fixed().ok()?)".to_string(),
295			KeyColumn::Blob16 => "{ let bytes = de.read_raw(16).ok()?; let mut buf = [0u8; 16]; \
296				 buf.copy_from_slice(bytes); buf }"
297				.to_string(),
298			KeyColumn::IdentityId => "de.read_identity_id().ok()?".to_string(),
299			KeyColumn::TableId
300			| KeyColumn::ColumnId
301			| KeyColumn::FlowId
302			| KeyColumn::FlowEdgeId
303			| KeyColumn::OperatorId
304			| KeyColumn::HandlerId
305			| KeyColumn::NamespaceId
306			| KeyColumn::SumTypeId
307			| KeyColumn::SequenceId
308			| KeyColumn::ViewId
309			| KeyColumn::SeriesId
310			| KeyColumn::SinkId
311			| KeyColumn::SourceId
312			| KeyColumn::QueueId
313			| KeyColumn::RingBufferId
314			| KeyColumn::BindingId
315			| KeyColumn::DictionaryId
316			| KeyColumn::ColumnPropertyId
317			| KeyColumn::RelationshipId
318			| KeyColumn::MigrationId
319			| KeyColumn::MigrationEventId
320			| KeyColumn::ColumnBlockId
321			| KeyColumn::PrimaryKeyId => format!("{field_type}(de.read_u64().ok()?)"),
322			KeyColumn::ProcedureId => "ProcedureId::from_raw(de.read_u64().ok()?)".to_string(),
323			KeyColumn::EpochSeconds => "EpochSeconds::new(de.read_u64().ok()?)".to_string(),
324			KeyColumn::DateTime => "de.read_datetime().ok()?".to_string(),
325			KeyColumn::RowShapeFingerprint => "RowShapeFingerprint::new(de.read_u64().ok()?)".to_string(),
326			KeyColumn::Partition => "Partition(de.read_u128().ok()?)".to_string(),
327			KeyColumn::IndexId => "IndexId::Primary(PrimaryKeyId(de.read_u64().ok()?))".to_string(),
328			KeyColumn::ObjectId => "de.read_object_id().ok()?".to_string(),
329			KeyColumn::StorageId => "StorageId::from_object(de.read_object_id().ok()?)?".to_string(),
330			KeyColumn::OptionU8 => "match de.read_u8().ok()? { 1u8 => Some(de.read_u8().ok()?), 0u8 => { \
331				de.read_u8().ok()?; None }, _ => return None }"
332				.to_string(),
333		}
334	}
335}
336
337pub fn derive_key(input: TokenStream) -> TokenStream {
338	let tokens: Vec<TokenTree> = input.into_iter().collect();
339	let mut iter = tokens.iter().peekable();
340
341	let mut tag: Option<String> = None;
342	while let Some(TokenTree::Punct(p)) = iter.peek() {
343		if p.as_char() != '#' {
344			break;
345		}
346		iter.next();
347		match iter.next() {
348			Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
349				if let Some(found) = parse_key_attribute(g) {
350					tag = Some(found);
351				}
352			}
353			_ => return compile_error("expected an attribute after '#'"),
354		}
355	}
356
357	let tag = match tag {
358		Some(tag) => tag,
359		None => return compile_error("KeyCodec requires #[key(tag = Variant)] naming its KeyTag"),
360	};
361
362	if let Some(TokenTree::Ident(i)) = iter.peek()
363		&& *i == "pub"
364	{
365		iter.next();
366		if let Some(TokenTree::Group(g)) = iter.peek()
367			&& g.delimiter() == Delimiter::Parenthesis
368		{
369			iter.next();
370		}
371	}
372
373	match iter.next() {
374		Some(TokenTree::Ident(i)) if *i == "struct" => {}
375		_ => return compile_error("KeyCodec can only be derived for structs"),
376	}
377
378	let name = match iter.next() {
379		Some(TokenTree::Ident(i)) => i.to_string(),
380		_ => return compile_error("expected struct name"),
381	};
382
383	let body = match iter.next() {
384		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Brace => g.clone(),
385		Some(TokenTree::Punct(p)) if p.as_char() == '<' => {
386			return compile_error("KeyCodec cannot be derived for a generic struct");
387		}
388		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => {
389			return compile_error("KeyCodec requires named fields, so a tuple struct has no field order");
390		}
391		_ => return compile_error("expected struct body"),
392	};
393
394	let fields = match parse_fields(&body) {
395		Ok(fields) => fields,
396		Err(err) => return err,
397	};
398
399	expand(&name, &tag, &fields)
400}
401
402fn parse_key_attribute(group: &Group) -> Option<String> {
403	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
404	let mut iter = tokens.iter();
405
406	match iter.next() {
407		Some(TokenTree::Ident(i)) if *i == "key" => {}
408		_ => return None,
409	}
410
411	let inner = match iter.next() {
412		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => g.clone(),
413		_ => return None,
414	};
415
416	let inner_tokens: Vec<TokenTree> = inner.stream().into_iter().collect();
417	match inner_tokens.as_slice() {
418		[TokenTree::Ident(tag_ident), TokenTree::Punct(eq), TokenTree::Ident(variant)]
419			if *tag_ident == "tag" && eq.as_char() == '=' =>
420		{
421			Some(variant.to_string())
422		}
423		_ => None,
424	}
425}
426
427fn parse_fields(body: &Group) -> Result<Vec<KeyField>, TokenStream> {
428	let tokens: Vec<TokenTree> = body.stream().into_iter().collect();
429	let mut iter = tokens.iter().peekable();
430	let mut fields = Vec::new();
431
432	while iter.peek().is_some() {
433		let mut attribute: Option<FieldAttribute> = None;
434		while let Some(TokenTree::Punct(p)) = iter.peek() {
435			if p.as_char() == '#' {
436				iter.next();
437				if let Some(TokenTree::Group(g)) = iter.peek() {
438					if g.delimiter() == Delimiter::Bracket
439						&& let Some(found) = field_attribute(g)
440					{
441						attribute = Some(found);
442					}
443					iter.next();
444				}
445			} else {
446				break;
447			}
448		}
449		if iter.peek().is_none() {
450			break;
451		}
452
453		if let Some(TokenTree::Ident(i)) = iter.peek()
454			&& *i == "pub"
455		{
456			iter.next();
457			if let Some(TokenTree::Group(g)) = iter.peek()
458				&& g.delimiter() == Delimiter::Parenthesis
459			{
460				iter.next();
461			}
462		}
463
464		let field_name = match iter.next() {
465			Some(TokenTree::Ident(i)) => i.to_string(),
466			_ => return Err(compile_error("expected field name")),
467		};
468
469		match iter.next() {
470			Some(TokenTree::Punct(p)) if p.as_char() == ':' => {}
471			_ => return Err(compile_error("expected ':' after field name")),
472		}
473
474		let mut ty_tokens: Vec<TokenTree> = Vec::new();
475		loop {
476			match iter.peek() {
477				Some(TokenTree::Punct(p)) if p.as_char() == ',' => {
478					iter.next();
479					break;
480				}
481				Some(t) => {
482					ty_tokens.push((*t).clone());
483					iter.next();
484				}
485				None => break,
486			}
487		}
488
489		if ty_tokens.is_empty() {
490			return Err(compile_error(&format!("field '{}' has no type", field_name)));
491		}
492
493		let ty = render(&ty_tokens);
494		let column = match attribute {
495			Some(FieldAttribute::ReprU8) => KeyColumn::ReprU8,
496			Some(FieldAttribute::Display) => KeyColumn::DisplayStr,
497			Some(FieldAttribute::Tagged) if ty.ends_with("IndexId") => KeyColumn::TaggedIndexId,
498			Some(FieldAttribute::Varint) if ty.ends_with("u128") => KeyColumn::VarintU128,
499			Some(FieldAttribute::Tagged) | Some(FieldAttribute::Varint) => {
500				return Err(compile_error(&format!(
501					"field '{}' has type '{}', which its key attribute does not apply to",
502					field_name, ty
503				)));
504			}
505			None => match column_type(&ty_tokens) {
506				Some(column) => column,
507				None => {
508					return Err(compile_error(&format!(
509						"field '{}' has type '{}', which has no flat key column type",
510						field_name, ty
511					)));
512				}
513			},
514		};
515
516		if fields.last().is_some_and(|last: &KeyField| matches!(last.column, KeyColumn::RawTail)) {
517			return Err(compile_error(&format!(
518				"field '{}' follows a raw tail, which must be the last field because it takes the rest of the key",
519				field_name
520			)));
521		}
522
523		fields.push(KeyField {
524			name: field_name,
525			column,
526			ty,
527		});
528	}
529
530	Ok(fields)
531}
532
533#[derive(Clone, Copy)]
534enum FieldAttribute {
535	ReprU8,
536	Display,
537	Tagged,
538	Varint,
539}
540
541fn field_attribute(group: &Group) -> Option<FieldAttribute> {
542	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
543	let mut iter = tokens.iter();
544
545	match iter.next() {
546		Some(TokenTree::Ident(i)) if *i == "key" => {}
547		_ => return None,
548	}
549
550	let inner = match iter.next() {
551		Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Parenthesis => g.clone(),
552		_ => return None,
553	};
554
555	let inner_tokens: Vec<TokenTree> = inner.stream().into_iter().collect();
556	match inner_tokens.as_slice() {
557		[TokenTree::Ident(repr), TokenTree::Punct(eq), TokenTree::Ident(width)]
558			if *repr == "repr" && eq.as_char() == '=' && *width == "u8" =>
559		{
560			Some(FieldAttribute::ReprU8)
561		}
562		[TokenTree::Ident(word)] if *word == "display" => Some(FieldAttribute::Display),
563		[TokenTree::Ident(word)] if *word == "tagged" => Some(FieldAttribute::Tagged),
564		[TokenTree::Ident(word)] if *word == "varint" => Some(FieldAttribute::Varint),
565		_ => None,
566	}
567}
568
569fn column_type(tokens: &[TokenTree]) -> Option<KeyColumn> {
570	if let [TokenTree::Group(group)] = tokens {
571		return (group.delimiter() == Delimiter::Bracket && is_byte_array_16(group))
572			.then_some(KeyColumn::Blob16);
573	}
574
575	if is_option_u8(tokens) {
576		return Some(KeyColumn::OptionU8);
577	}
578
579	if is_generic_of(tokens, "Vec", "u8") {
580		return Some(KeyColumn::RawTail);
581	}
582
583	if is_generic_of(strip_path(tokens), "Asc", "RowNumber") {
584		return Some(KeyColumn::AscRowNumber);
585	}
586
587	let unqualified = strip_path(tokens);
588	if unqualified.len() != 1 {
589		return None;
590	}
591
592	let head = match unqualified.first() {
593		Some(TokenTree::Ident(i)) => i.to_string(),
594		_ => return None,
595	};
596
597	match head.as_str() {
598		"u8" => Some(KeyColumn::U8),
599		"u16" => Some(KeyColumn::U16),
600		"u32" => Some(KeyColumn::U32),
601		"u64" => Some(KeyColumn::U64),
602		"u128" => Some(KeyColumn::U128),
603		"RowNumber" => Some(KeyColumn::RowNumber),
604		"GroupId" => Some(KeyColumn::GroupId),
605		"TableId" => Some(KeyColumn::TableId),
606		"ColumnId" => Some(KeyColumn::ColumnId),
607		"ObjectId" => Some(KeyColumn::ObjectId),
608		"StorageId" => Some(KeyColumn::StorageId),
609		"QueueId" => Some(KeyColumn::QueueId),
610		"RingBufferId" => Some(KeyColumn::RingBufferId),
611		"SeriesId" => Some(KeyColumn::SeriesId),
612		"DateTime" => Some(KeyColumn::DateTime),
613		"Partition" => Some(KeyColumn::Partition),
614		"FlowId" => Some(KeyColumn::FlowId),
615		"FlowEdgeId" => Some(KeyColumn::FlowEdgeId),
616		"OperatorId" => Some(KeyColumn::OperatorId),
617		"HandlerId" => Some(KeyColumn::HandlerId),
618		"NamespaceId" => Some(KeyColumn::NamespaceId),
619		"SumTypeId" => Some(KeyColumn::SumTypeId),
620		"SequenceId" => Some(KeyColumn::SequenceId),
621		"EpochSeconds" => Some(KeyColumn::EpochSeconds),
622		"IdentityId" => Some(KeyColumn::IdentityId),
623		"AuthenticationId" | "IdentityAttributeId" | "RoleId" | "PolicyId" | "TokenId" => Some(KeyColumn::U64),
624		"ViewId" => Some(KeyColumn::ViewId),
625		"SinkId" => Some(KeyColumn::SinkId),
626		"SourceId" => Some(KeyColumn::SourceId),
627		"BindingId" => Some(KeyColumn::BindingId),
628		"DictionaryId" => Some(KeyColumn::DictionaryId),
629		"ProcedureId" => Some(KeyColumn::ProcedureId),
630		"ColumnPropertyId" => Some(KeyColumn::ColumnPropertyId),
631		"RowShapeFingerprint" => Some(KeyColumn::RowShapeFingerprint),
632		"IndexId" => Some(KeyColumn::IndexId),
633		"RelationshipId" => Some(KeyColumn::RelationshipId),
634		"MigrationId" => Some(KeyColumn::MigrationId),
635		"MigrationEventId" => Some(KeyColumn::MigrationEventId),
636		"PrimaryKeyId" => Some(KeyColumn::PrimaryKeyId),
637		"ColumnBlockId" => Some(KeyColumn::ColumnBlockId),
638		"CdcConsumerId" => Some(KeyColumn::Str),
639		"EncodedKey" => Some(KeyColumn::EscapedBytes),
640		"KeyspaceId" => Some(KeyColumn::KeyspaceId),
641		"EncodedIndexKey" => Some(KeyColumn::RawTail),
642		"SortRun" => Some(KeyColumn::SortRun),
643		_ => None,
644	}
645}
646
647fn is_generic_of(tokens: &[TokenTree], outer: &str, inner: &str) -> bool {
648	matches!(
649		tokens,
650		[TokenTree::Ident(o), TokenTree::Punct(lt), TokenTree::Ident(u), TokenTree::Punct(gt)]
651			if *o == outer && lt.as_char() == '<' && *u == inner && gt.as_char() == '>'
652	)
653}
654
655fn is_option_u8(tokens: &[TokenTree]) -> bool {
656	matches!(
657		tokens,
658		[TokenTree::Ident(o), TokenTree::Punct(lt), TokenTree::Ident(u), TokenTree::Punct(gt)]
659			if *o == "Option" && lt.as_char() == '<' && *u == "u8" && gt.as_char() == '>'
660	)
661}
662
663fn is_byte_array_16(group: &Group) -> bool {
664	let tokens: Vec<TokenTree> = group.stream().into_iter().collect();
665	match tokens.as_slice() {
666		[TokenTree::Ident(element), TokenTree::Punct(semi), TokenTree::Literal(len)] => {
667			*element == "u8" && semi.as_char() == ';' && len.to_string() == "16"
668		}
669		_ => false,
670	}
671}
672
673fn strip_path(tokens: &[TokenTree]) -> &[TokenTree] {
674	let mut at = 0;
675	if is_colon(tokens.first()) && is_colon(tokens.get(1)) {
676		at = 2;
677	}
678	while matches!(tokens.get(at), Some(TokenTree::Ident(_)))
679		&& is_colon(tokens.get(at + 1))
680		&& is_colon(tokens.get(at + 2))
681	{
682		at += 3;
683	}
684	&tokens[at..]
685}
686
687fn is_colon(token: Option<&TokenTree>) -> bool {
688	matches!(token, Some(TokenTree::Punct(p)) if p.as_char() == ':')
689}
690
691fn render(tokens: &[TokenTree]) -> String {
692	tokens.iter().map(|token| token.to_string()).collect()
693}
694
695fn expand(name: &str, tag: &str, fields: &[KeyField]) -> TokenStream {
696	let capacity: usize = 1 + fields.iter().map(|f| f.column.width()).sum::<usize>();
697
698	let mut encode_body = String::new();
699	for field in fields {
700		encode_body.push_str("\t\t");
701		encode_body.push_str(&field.column.encode_stmt(&field.name));
702		encode_body.push('\n');
703	}
704
705	let mut decode_body = String::new();
706	for field in fields {
707		decode_body.push_str(&format!("\t\t\t{}: {},\n", field.name, field.column.decode_expr(&field.ty)));
708	}
709
710	let mut out = String::new();
711	out.push_str(&format!("impl {name} {{\n"));
712	out.push_str(&format!("\tpub const TAG: KeyTag = KeyTag::{tag};\n\n"));
713	out.push_str("\tpub fn encode(&self) -> EncodedKey {\n");
714	out.push_str(&format!(
715		"\t\tlet mut serializer = ::reifydb_codec::key::serializer::KeySerializer::with_capacity({capacity});\n"
716	));
717	out.push_str("\t\tserializer.extend_u8(Self::TAG as u8);\n");
718	out.push_str(&encode_body);
719	out.push_str("\t\tserializer.to_encoded_key()\n\t}\n\n");
720	out.push_str("\tpub fn decode(key: &EncodedKey) -> Option<Self> {\n");
721	out.push_str(
722		"\t\tlet mut de = ::reifydb_codec::key::deserializer::KeyDeserializer::from_bytes(key.as_slice());\n",
723	);
724	out.push_str("\t\tlet found: KeyTag = de.read_u8().ok()?.try_into().ok()?;\n");
725	out.push_str("\t\tif found != Self::TAG {\n\t\t\treturn None;\n\t\t}\n");
726	out.push_str(&format!("\t\tlet decoded = Self {{\n{decode_body}\t\t}};\n"));
727	out.push_str("\t\tif !de.is_empty() {\n\t\t\treturn None;\n\t\t}\n");
728	out.push_str("\t\tSome(decoded)\n\t}\n}\n\n");
729
730	out.push_str(&format!("#[automatically_derived]\nimpl KeyFields for {name} {{\n"));
731	out.push_str("\tfn fields(&self) -> ::smallvec::SmallVec<[Field<'_>; 6]> {\n");
732	out.push_str("\t\t::smallvec::smallvec![\n");
733	for field in fields {
734		for expr in field.column.field_exprs(&field.name) {
735			out.push_str(&format!("\t\t\t{expr},\n"));
736		}
737	}
738	out.push_str("\t\t]\n\t}\n}\n\n");
739	out.push_str(&format!("#[automatically_derived]\nimpl {name} {{\n"));
740	out.push_str("\tpub fn cmp_fields(&self, other: &Self) -> ::core::cmp::Ordering {\n");
741	for field in fields {
742		for expr in field.column.field_exprs(&field.name) {
743			let theirs = expr.replace("self.", "other.");
744			out.push_str(&format!(
745				"\t\tmatch Ord::cmp(&{expr}, &{theirs}) {{\n\t\t\t::core::cmp::Ordering::Equal => {{}}\n\t\t\tordering => return ordering,\n\t\t}}\n"
746			));
747		}
748	}
749	out.push_str("\t\t::core::cmp::Ordering::Equal\n\t}\n}\n\n");
750	out.push_str(&expand_tests(name, fields));
751
752	out.parse().expect("derived key impl must be valid Rust")
753}
754
755#[cfg(test)]
756mod tests {
757	use super::derive_key;
758
759	fn expand(source: &str) -> String {
760		derive_key(source.parse().unwrap()).to_string()
761	}
762
763	#[test]
764	fn missing_attribute_is_rejected() {
765		let out = expand("struct RowKey { row: u64 }");
766		assert!(out.contains("compile_error"), "{out}");
767		assert!(out.contains("key"), "{out}");
768	}
769
770	#[test]
771	fn a_well_formed_struct_expands_without_an_error() {
772		let out = expand("#[key(tag = Row)] struct RowKey { table: u64, row: RowNumber }");
773		assert!(!out.contains("compile_error"), "{out}");
774		assert!(out.contains("impl RowKey { pub const TAG"), "{out}");
775		assert!(out.contains("KeyTag :: Row"), "{out}");
776	}
777
778	#[test]
779	fn fields_are_encoded_in_declaration_order() {
780		let out = expand("#[key(tag = Row)] struct RowKey { table: u64, row: RowNumber }");
781		let table = out.find("self . table").expect("table field encoded");
782		let row = out.find("self . row").expect("row field encoded");
783		assert!(table < row, "{out}");
784	}
785
786	#[test]
787	fn a_repr_u8_field_encodes_the_discriminant_and_decodes_through_try_from() {
788		// without the attribute a repr(u8) enum has no column type at all, and a decode that skipped
789		// try_from would hand back a discriminant the enum never declared
790		let out = expand("#[key(tag = SystemVersion)] struct K { #[key(repr = u8)] version: SystemVersion }");
791		assert!(!out.contains("compile_error"), "{out}");
792		assert!(out.contains("extend_u8 (self . version as u8)"), "{out}");
793		assert!(out.contains("SystemVersion :: try_from"), "{out}");
794	}
795
796	#[test]
797	fn a_repr_u8_field_still_rejects_a_type_the_attribute_was_not_put_on() {
798		let out = expand("#[key(tag = SystemVersion)] struct K { version: SystemVersion }");
799		assert!(out.contains("compile_error"), "{out}");
800	}
801
802	#[test]
803	fn an_unrecognised_field_type_is_rejected() {
804		let out = expand("#[key(tag = Row)] struct RowKey { at: String }");
805		assert!(out.contains("compile_error"), "{out}");
806		assert!(out.contains("at"), "{out}");
807		assert!(out.contains("String"), "{out}");
808	}
809
810	#[test]
811	fn a_tuple_struct_is_rejected() {
812		let out = expand("#[key(tag = Row)] struct RowKey(u64);");
813		assert!(out.contains("compile_error"), "{out}");
814	}
815
816	#[test]
817	fn a_generic_struct_is_rejected() {
818		let out = expand("#[key(tag = Row)] struct RowKey<T> { row: T }");
819		assert!(out.contains("compile_error"), "{out}");
820	}
821
822	#[test]
823	fn a_byte_array_field_maps_to_blob16() {
824		let out = expand("#[key(tag = Row)] struct RowKey { blob: [u8; 16] }");
825		assert!(!out.contains("compile_error"), "{out}");
826		assert!(out.contains("extend_raw"));
827	}
828
829	#[test]
830	fn a_byte_array_of_the_wrong_width_is_rejected() {
831		let out = expand("#[key(tag = Row)] struct RowKey { blob: [u8; 32] }");
832		assert!(out.contains("compile_error"), "{out}");
833	}
834
835	#[test]
836	fn table_id_and_column_id_encode_via_their_wrapped_u64() {
837		let out = expand("#[key(tag = Table)] struct TableKey { table: TableId, column: ColumnId }");
838		assert!(!out.contains("compile_error"), "{out}");
839		assert!(out.contains("self . table . 0"), "{out}");
840		assert!(out.contains("TableId (de . read_u64 () . ok () ?)"), "{out}");
841		assert!(out.contains("ColumnId (de . read_u64 () . ok () ?)"), "{out}");
842	}
843
844	#[test]
845	fn an_option_u8_field_encodes_as_a_fixed_two_byte_presence_and_value_pair() {
846		let out = expand("#[key(tag = Row)] struct RowKey { tag: Option<u8> }");
847		assert!(!out.contains("compile_error"), "{out}");
848		assert!(out.contains("with_capacity (3)"), "{out}: width must stay fixed at 2 bytes for None and Some");
849		assert!(out.contains("Some (v) => { serializer . extend_u8 (1u8) . extend_u8 (v) ; }"), "{out}");
850		assert!(out.contains("None => { serializer . extend_u8 (0u8) . extend_u8 (0u8) ; }"), "{out}");
851		assert!(out.contains("1u8 => Some (de . read_u8 () . ok () ?)"), "{out}");
852		assert!(out.contains("0u8 => { de . read_u8 () . ok () ?; None }"), "{out}");
853	}
854
855	#[test]
856	fn object_id_and_storage_id_encode_via_the_catalog_ext() {
857		let out = expand("#[key(tag = Row)] struct RowKey { storage: StorageId, object: ObjectId }");
858		assert!(!out.contains("compile_error"), "{out}");
859		assert!(out.contains("extend_object_id (self . storage)"), "{out}");
860		assert!(out.contains("extend_object_id (self . object)"), "{out}");
861		assert!(out.contains("StorageId :: from_object (de . read_object_id () . ok () ?) ?"), "{out}");
862		assert!(out.contains("de . read_object_id () . ok () ?"), "{out}");
863	}
864}
reifydb_macro_impl/key.rs

reifydb_macro_impl/
key.rs
