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

reifydb_macro_impl/
key.rs
