1use std::cmp::Ordering;
2use std::collections::HashSet;
3use std::sync::Arc;
4
5use ahash::RandomState;
6use arrow::datatypes::{DataType, FieldRef, Schema, TimeUnit};
7use regex::Regex;
8
9use super::ast::{append_path, column_path};
10use super::bounds::parse_bound;
11use super::{
12 CompareOpAst, ComparePlan, CompositeNullPolicyAst, ContractAst, ContractAstV2,
13 ContractDocument, ContractVersion, CountRangeAst, DatasetPlan, NaNPolicyAst, NullPolicyAst,
14 OperandPlan, ParameterizedTypeAst, PrimitiveTypeAst, RowPlan, RowPlanKind, RuleAst, RuleAstV2,
15 ScalarValuePlan, TimeUnitAst, TypeAst, TypedBound,
16};
17use crate::{ErrorCode, ProofFrameError};
18
19#[derive(Debug, Clone, Eq, PartialEq)]
21pub enum KernelKind {
22 Boolean,
23 I8,
24 I16,
25 I32,
26 I64,
27 U8,
28 U16,
29 U32,
30 U64,
31 F32,
32 F64,
33 Date32,
34 Date64,
35 Decimal128 { precision: u8, scale: i8 },
36 Timestamp(TimeUnit),
37 Utf8,
38 LargeUtf8,
39 Utf8View,
40 Binary,
41 LargeBinary,
42 BinaryView,
43 Nested,
44 NullOnly,
45}
46
47impl KernelKind {
48 fn from_data_type(data_type: &DataType) -> Self {
49 match data_type {
50 DataType::Boolean => Self::Boolean,
51 DataType::Int8 => Self::I8,
52 DataType::Int16 => Self::I16,
53 DataType::Int32 => Self::I32,
54 DataType::Int64 => Self::I64,
55 DataType::UInt8 => Self::U8,
56 DataType::UInt16 => Self::U16,
57 DataType::UInt32 => Self::U32,
58 DataType::UInt64 => Self::U64,
59 DataType::Float32 => Self::F32,
60 DataType::Float64 => Self::F64,
61 DataType::Date32 => Self::Date32,
62 DataType::Date64 => Self::Date64,
63 DataType::Decimal128(precision, scale) => Self::Decimal128 {
64 precision: *precision,
65 scale: *scale,
66 },
67 DataType::Timestamp(unit, _) => Self::Timestamp(*unit),
68 DataType::Utf8 => Self::Utf8,
69 DataType::LargeUtf8 => Self::LargeUtf8,
70 DataType::Utf8View => Self::Utf8View,
71 DataType::Binary => Self::Binary,
72 DataType::LargeBinary => Self::LargeBinary,
73 DataType::BinaryView => Self::BinaryView,
74 DataType::List(_)
75 | DataType::LargeList(_)
76 | DataType::FixedSizeList(_, _)
77 | DataType::Struct(_)
78 | DataType::Map(_, _) => Self::Nested,
79 _ => Self::NullOnly,
80 }
81 }
82
83 pub(crate) fn from_data_type_for_plan(data_type: &DataType) -> Self {
84 Self::from_data_type(data_type)
85 }
86
87 fn supports_bounds(&self) -> bool {
88 matches!(
89 self,
90 Self::I8
91 | Self::I16
92 | Self::I32
93 | Self::I64
94 | Self::U8
95 | Self::U16
96 | Self::U32
97 | Self::U64
98 | Self::F32
99 | Self::F64
100 | Self::Date32
101 | Self::Date64
102 | Self::Decimal128 { .. }
103 | Self::Timestamp(_)
104 )
105 }
106
107 fn supports_text_rules(&self) -> bool {
108 matches!(self, Self::Utf8 | Self::LargeUtf8 | Self::Utf8View)
109 }
110
111 fn supports_unique(&self) -> bool {
112 !matches!(self, Self::NullOnly)
113 }
114}
115
116#[derive(Debug, Clone, Copy, Eq, PartialEq)]
118pub enum NaNPolicy {
119 Reject,
120 Allow,
121}
122
123impl From<Option<NaNPolicyAst>> for NaNPolicy {
124 fn from(value: Option<NaNPolicyAst>) -> Self {
125 match value.unwrap_or_default() {
126 NaNPolicyAst::Reject => Self::Reject,
127 NaNPolicyAst::Allow => Self::Allow,
128 }
129 }
130}
131
132#[derive(Debug, Clone)]
134pub struct CompiledRules {
135 pub(crate) required: bool,
136 pub(crate) not_null: bool,
137 pub(crate) unique: bool,
138 pub(crate) min: Option<TypedBound>,
139 pub(crate) max: Option<TypedBound>,
140 pub(crate) nan: NaNPolicy,
141 pub(crate) validate_nan: bool,
142 pub(crate) pattern: Option<Regex>,
143 pub(crate) allowed: Option<Arc<HashSet<Box<str>, RandomState>>>,
144}
145
146impl CompiledRules {
147 #[must_use]
148 pub const fn required(&self) -> bool {
149 self.required
150 }
151
152 #[must_use]
153 pub const fn not_null(&self) -> bool {
154 self.not_null
155 }
156
157 #[must_use]
158 pub const fn unique(&self) -> bool {
159 self.unique
160 }
161
162 #[must_use]
163 pub fn min(&self) -> Option<&TypedBound> {
164 self.min.as_ref()
165 }
166
167 #[must_use]
168 pub fn max(&self) -> Option<&TypedBound> {
169 self.max.as_ref()
170 }
171
172 #[must_use]
173 pub const fn nan(&self) -> NaNPolicy {
174 self.nan
175 }
176
177 #[must_use]
178 pub const fn validates_nan(&self) -> bool {
179 self.validate_nan
180 }
181
182 #[must_use]
183 pub fn pattern(&self) -> Option<&Regex> {
184 self.pattern.as_ref()
185 }
186
187 #[must_use]
188 pub fn allowed(&self) -> Option<&HashSet<Box<str>, RandomState>> {
189 self.allowed.as_deref()
190 }
191}
192
193#[derive(Debug, Clone)]
195pub struct ColumnPlan {
196 column_index: usize,
197 field: FieldRef,
198 kernel: KernelKind,
199 rules: CompiledRules,
200}
201
202impl ColumnPlan {
203 #[must_use]
204 pub const fn column_index(&self) -> usize {
205 self.column_index
206 }
207
208 #[must_use]
209 pub fn field(&self) -> &FieldRef {
210 &self.field
211 }
212
213 #[must_use]
214 pub const fn kernel(&self) -> &KernelKind {
215 &self.kernel
216 }
217
218 #[must_use]
219 pub const fn rules(&self) -> &CompiledRules {
220 &self.rules
221 }
222}
223
224#[derive(Debug, Clone)]
226pub struct CompiledContract {
227 schema: Schema,
228 columns: Vec<ColumnPlan>,
229 row_plans: Vec<RowPlan>,
230 dataset_plan: DatasetPlan,
231 version: ContractVersion,
232 max_findings: usize,
233}
234
235impl CompiledContract {
236 pub fn compile(contract: &ContractAst, schema: &Schema) -> Result<Self, ProofFrameError> {
238 for (name, rules) in &contract.columns {
239 if schema.index_of(name).is_ok() {
240 continue;
241 }
242 if rules.required {
243 return Err(ProofFrameError::contract(
244 ErrorCode::MissingColumn,
245 format!("Required contract column `{name}` is absent from the Arrow schema"),
246 Some(column_path(name)),
247 ));
248 } else if has_value_rules(rules) {
249 return Err(ProofFrameError::contract(
250 ErrorCode::MissingColumn,
251 format!("Contract column `{name}` is absent from the Arrow schema"),
252 Some(column_path(name)),
253 ));
254 }
255 }
256
257 let mut columns = Vec::with_capacity(contract.columns.len());
258 for (column_index, field) in schema.fields().iter().enumerate() {
259 let Some(source_rules) = contract.columns.get(field.name()) else {
260 continue;
261 };
262 if !has_runtime_rules(source_rules) {
263 continue;
264 }
265 let kernel = KernelKind::from_data_type(field.data_type());
266 let rules = compile_rules(field.name(), field.data_type(), &kernel, source_rules)?;
267 columns.push(ColumnPlan {
268 column_index,
269 field: field.clone(),
270 kernel,
271 rules,
272 });
273 }
274
275 Ok(Self {
276 schema: schema.clone(),
277 columns,
278 row_plans: Vec::new(),
279 dataset_plan: DatasetPlan::default(),
280 version: ContractVersion::V1,
281 max_findings: contract.max_findings,
282 })
283 }
284
285 pub fn compile_document(
287 document: &ContractDocument,
288 schema: &Schema,
289 ) -> Result<Self, ProofFrameError> {
290 match document {
291 ContractDocument::V1(contract) => Self::compile(contract, schema),
292 ContractDocument::V2(contract) => Self::compile_v2(contract, schema),
293 }
294 }
295
296 fn compile_v2(contract: &ContractAstV2, schema: &Schema) -> Result<Self, ProofFrameError> {
297 for (name, rules) in &contract.columns {
298 let Ok(column_index) = schema.index_of(name) else {
299 if rules.required || has_v2_value_rules(rules) {
300 return Err(ProofFrameError::contract(
301 ErrorCode::MissingColumn,
302 format!("Contract column `{name}` is absent from the Arrow schema"),
303 Some(column_path(name)),
304 ));
305 }
306 continue;
307 };
308 if let Some(expected) = rules.expected_type.as_ref() {
309 let actual = schema.field(column_index).data_type();
310 if !expected_type_matches(expected, actual) {
311 return Err(ProofFrameError::contract(
312 ErrorCode::ContractTypeMismatch,
313 format!("Column `{name}` requires `{expected:?}`, found `{actual}`"),
314 Some(append_path(&column_path(name), "type")),
315 ));
316 }
317 }
318 }
319
320 let mut columns = Vec::with_capacity(contract.columns.len());
321 for (column_index, field) in schema.fields().iter().enumerate() {
322 let Some(source) = contract.columns.get(field.name()) else {
323 continue;
324 };
325 let source_rules = v2_rules_as_v1(source);
326 if !has_runtime_rules(&source_rules) {
327 continue;
328 }
329 let kernel = KernelKind::from_data_type(field.data_type());
330 let rules = compile_rules(field.name(), field.data_type(), &kernel, &source_rules)?;
331 columns.push(ColumnPlan {
332 column_index,
333 field: field.clone(),
334 kernel,
335 rules,
336 });
337 }
338
339 let row_plans = contract
340 .row_rules
341 .iter()
342 .enumerate()
343 .map(|(index, rule)| RowPlan::compile(rule, schema, index))
344 .collect::<Result<Vec<_>, _>>()?;
345 let dataset_plan = DatasetPlan::compile(contract, schema)?;
346 Ok(Self {
347 schema: schema.clone(),
348 columns,
349 row_plans,
350 dataset_plan,
351 version: ContractVersion::V2,
352 max_findings: contract.max_findings,
353 })
354 }
355
356 #[must_use]
357 pub fn columns(&self) -> &[ColumnPlan] {
358 &self.columns
359 }
360
361 #[must_use]
362 pub const fn schema(&self) -> &Schema {
363 &self.schema
364 }
365
366 #[must_use]
367 pub fn row_plans(&self) -> &[RowPlan] {
368 &self.row_plans
369 }
370
371 #[must_use]
372 pub const fn dataset_plan(&self) -> &DatasetPlan {
373 &self.dataset_plan
374 }
375
376 #[must_use]
377 pub const fn max_findings(&self) -> usize {
378 self.max_findings
379 }
380
381 pub fn schema_digest(&self) -> Result<String, ProofFrameError> {
383 let mut hasher = blake3::Hasher::new();
384 hasher.update(b"proofframe:schema:v1\0");
385 hasher.update(&crate::encoding::canonical_schema_digest(&self.schema)?);
386 Ok(tagged_digest("pf-schema-v1:", hasher.finalize()))
387 }
388
389 pub fn compiled_plan_digest(&self) -> Result<String, ProofFrameError> {
391 if self.version == ContractVersion::V2 {
392 return self.compiled_plan_digest_v2();
393 }
394 let mut hasher = blake3::Hasher::new();
395 hasher.update(b"proofframe:compiled-plan:v1\0");
396 hasher.update(&crate::encoding::canonical_schema_digest(&self.schema)?);
397 hasher.update(&(self.max_findings as u64).to_le_bytes());
398 for column in &self.columns {
399 hasher.update(&(column.column_index as u64).to_le_bytes());
400 hash_part(&mut hasher, column.field.name().as_bytes());
401 hash_kernel(&mut hasher, &column.kernel);
402 let rules = &column.rules;
403 hasher.update(&[
404 u8::from(rules.required),
405 u8::from(rules.not_null),
406 u8::from(rules.unique),
407 u8::from(rules.validate_nan),
408 match rules.nan {
409 NaNPolicy::Reject => 0,
410 NaNPolicy::Allow => 1,
411 },
412 ]);
413 hash_optional_bound(&mut hasher, rules.min.as_ref());
414 hash_optional_bound(&mut hasher, rules.max.as_ref());
415 hash_optional_part(
416 &mut hasher,
417 rules
418 .pattern
419 .as_ref()
420 .map(|value| value.as_str().as_bytes()),
421 );
422 if let Some(allowed) = rules.allowed.as_deref() {
423 hasher.update(&[1]);
424 let mut values = allowed.iter().map(AsRef::as_ref).collect::<Vec<&str>>();
425 values.sort_unstable();
426 hasher.update(&(values.len() as u64).to_le_bytes());
427 for value in values {
428 hash_part(&mut hasher, value.as_bytes());
429 }
430 } else {
431 hasher.update(&[0]);
432 }
433 }
434 Ok(tagged_digest("pf-plan-v1:", hasher.finalize()))
435 }
436
437 fn compiled_plan_digest_v2(&self) -> Result<String, ProofFrameError> {
438 let mut hasher = blake3::Hasher::new();
439 hasher.update(b"proofframe:compiled-plan:v2\0");
440 hasher.update(&crate::encoding::canonical_schema_digest(&self.schema)?);
441 hasher.update(&(self.max_findings as u64).to_le_bytes());
442 for column in &self.columns {
443 hasher.update(&(column.column_index as u64).to_le_bytes());
444 hash_part(&mut hasher, column.field.name().as_bytes());
445 hash_kernel(&mut hasher, &column.kernel);
446 hash_compiled_rules(&mut hasher, &column.rules);
447 }
448 hasher.update(&(self.row_plans.len() as u64).to_le_bytes());
449 for row in &self.row_plans {
450 hash_row_plan(&mut hasher, row);
451 }
452 hash_dataset_plan(&mut hasher, &self.dataset_plan);
453 Ok(tagged_digest("pf-plan-v2:", hasher.finalize()))
454 }
455}
456
457fn hash_row_plan(hasher: &mut blake3::Hasher, plan: &RowPlan) {
458 hash_part(hasher, plan.name().as_bytes());
459 match plan.kind() {
460 RowPlanKind::Compare(compare) => {
461 hasher.update(&[0]);
462 hash_compare_plan(hasher, compare);
463 }
464 RowPlanKind::Conditional {
465 predicate,
466 assertion_column,
467 assertion_field,
468 assertion,
469 } => {
470 hasher.update(&[1]);
471 hash_compare_plan(hasher, predicate);
472 hasher.update(&(*assertion_column as u64).to_le_bytes());
473 hash_part(hasher, assertion_field.name().as_bytes());
474 hash_compiled_rules(hasher, assertion);
475 }
476 }
477}
478
479fn hash_compare_plan(hasher: &mut blake3::Hasher, plan: &ComparePlan) {
480 hash_operand_plan(hasher, plan.left());
481 hasher.update(&[compare_op_tag(plan.op()), null_policy_tag(plan.nulls())]);
482 hash_operand_plan(hasher, plan.right());
483}
484
485fn hash_operand_plan(hasher: &mut blake3::Hasher, operand: &OperandPlan) {
486 match operand {
487 OperandPlan::Column {
488 column_index,
489 field,
490 kernel,
491 } => {
492 hasher.update(&[0]);
493 hasher.update(&(*column_index as u64).to_le_bytes());
494 hash_part(hasher, field.name().as_bytes());
495 hash_kernel(hasher, kernel);
496 }
497 OperandPlan::Literal(value) => {
498 hasher.update(&[1]);
499 match value {
500 ScalarValuePlan::Boolean(value) => {
501 hasher.update(&[0, u8::from(*value)]);
502 }
503 ScalarValuePlan::I64(value) => {
504 hasher.update(&[1]);
505 hasher.update(&value.to_le_bytes());
506 }
507 ScalarValuePlan::U64(value) => {
508 hasher.update(&[2]);
509 hasher.update(&value.to_le_bytes());
510 }
511 ScalarValuePlan::F64(value) => {
512 hasher.update(&[3]);
513 hasher.update(&value.to_bits().to_le_bytes());
514 }
515 ScalarValuePlan::Text(value) => {
516 hasher.update(&[4]);
517 hash_part(hasher, value.as_bytes());
518 }
519 }
520 }
521 }
522}
523
524fn hash_dataset_plan(hasher: &mut blake3::Hasher, plan: &DatasetPlan) {
525 hash_optional_count_range(hasher, plan.row_count());
526 hasher.update(&(plan.null_ratios().len() as u64).to_le_bytes());
527 for ratio in plan.null_ratios() {
528 hasher.update(&(ratio.column_index() as u64).to_le_bytes());
529 hash_part(hasher, ratio.column().as_bytes());
530 hash_optional_f64(hasher, ratio.range().min);
531 hash_optional_f64(hasher, ratio.range().max);
532 hash_kernel(hasher, ratio.kernel());
533 }
534 hasher.update(&(plan.distinct_counts().len() as u64).to_le_bytes());
535 for count in plan.distinct_counts() {
536 hasher.update(&(count.column_index() as u64).to_le_bytes());
537 hash_part(hasher, count.column().as_bytes());
538 hash_count_range(hasher, count.range());
539 hash_kernel(hasher, count.kernel());
540 }
541 hasher.update(&(plan.distinct_ratios().len() as u64).to_le_bytes());
542 for ratio in plan.distinct_ratios() {
543 hasher.update(&(ratio.column_index() as u64).to_le_bytes());
544 hash_part(hasher, ratio.column().as_bytes());
545 hash_optional_f64(hasher, ratio.range().min);
546 hash_optional_f64(hasher, ratio.range().max);
547 hash_kernel(hasher, ratio.kernel());
548 }
549 hasher.update(&(plan.composite_unique().len() as u64).to_le_bytes());
550 for composite in plan.composite_unique() {
551 hash_part(hasher, composite.name().as_bytes());
552 hasher.update(&(composite.columns().len() as u64).to_le_bytes());
553 for column in composite.columns() {
554 hasher.update(&(*column as u64).to_le_bytes());
555 }
556 hasher.update(&[match composite.nulls() {
557 CompositeNullPolicyAst::Equal => 0,
558 CompositeNullPolicyAst::Reject => 1,
559 }]);
560 }
561}
562
563fn hash_optional_count_range(hasher: &mut blake3::Hasher, range: Option<&CountRangeAst>) {
564 if let Some(range) = range {
565 hasher.update(&[1]);
566 hash_count_range(hasher, range);
567 } else {
568 hasher.update(&[0]);
569 }
570}
571
572fn hash_count_range(hasher: &mut blake3::Hasher, range: &CountRangeAst) {
573 hash_optional_u64(hasher, range.exact);
574 hash_optional_u64(hasher, range.min);
575 hash_optional_u64(hasher, range.max);
576}
577
578fn hash_optional_u64(hasher: &mut blake3::Hasher, value: Option<u64>) {
579 match value {
580 Some(value) => {
581 hasher.update(&[1]);
582 hasher.update(&value.to_le_bytes());
583 }
584 None => {
585 hasher.update(&[0]);
586 }
587 }
588}
589
590fn hash_optional_f64(hasher: &mut blake3::Hasher, value: Option<f64>) {
591 match value {
592 Some(value) => {
593 hasher.update(&[1]);
594 hasher.update(&value.to_bits().to_le_bytes());
595 }
596 None => {
597 hasher.update(&[0]);
598 }
599 }
600}
601
602const fn compare_op_tag(op: CompareOpAst) -> u8 {
603 match op {
604 CompareOpAst::Eq => 0,
605 CompareOpAst::Ne => 1,
606 CompareOpAst::Lt => 2,
607 CompareOpAst::Lte => 3,
608 CompareOpAst::Gt => 4,
609 CompareOpAst::Gte => 5,
610 }
611}
612
613const fn null_policy_tag(policy: NullPolicyAst) -> u8 {
614 match policy {
615 NullPolicyAst::Skip => 0,
616 NullPolicyAst::Fail => 1,
617 NullPolicyAst::Equal => 2,
618 }
619}
620
621fn hash_compiled_rules(hasher: &mut blake3::Hasher, rules: &CompiledRules) {
622 hasher.update(&[
623 u8::from(rules.required),
624 u8::from(rules.not_null),
625 u8::from(rules.unique),
626 u8::from(rules.validate_nan),
627 match rules.nan {
628 NaNPolicy::Reject => 0,
629 NaNPolicy::Allow => 1,
630 },
631 ]);
632 hash_optional_bound(hasher, rules.min.as_ref());
633 hash_optional_bound(hasher, rules.max.as_ref());
634 hash_optional_part(
635 hasher,
636 rules
637 .pattern
638 .as_ref()
639 .map(|value| value.as_str().as_bytes()),
640 );
641 if let Some(allowed) = rules.allowed.as_deref() {
642 hasher.update(&[1]);
643 let mut values = allowed.iter().map(AsRef::as_ref).collect::<Vec<&str>>();
644 values.sort_unstable();
645 hasher.update(&(values.len() as u64).to_le_bytes());
646 for value in values {
647 hash_part(hasher, value.as_bytes());
648 }
649 } else {
650 hasher.update(&[0]);
651 }
652}
653
654fn has_v2_value_rules(rules: &RuleAstV2) -> bool {
655 rules.not_null
656 || rules.unique
657 || rules.expected_type.is_some()
658 || rules.min.is_some()
659 || rules.max.is_some()
660 || rules.nan.is_some()
661 || rules.pattern.is_some()
662 || rules.allowed.is_some()
663}
664
665fn v2_rules_as_v1(rules: &RuleAstV2) -> RuleAst {
666 RuleAst {
667 required: rules.required,
668 not_null: rules.not_null,
669 unique: rules.unique,
670 min: rules.min.clone(),
671 max: rules.max.clone(),
672 nan: rules.nan,
673 pattern: rules.pattern.clone(),
674 allowed: rules.allowed.clone(),
675 }
676}
677
678fn expected_type_matches(expected: &TypeAst, actual: &DataType) -> bool {
679 match expected {
680 TypeAst::Primitive(expected) => matches_primitive_type(*expected, actual),
681 TypeAst::Parameterized(ParameterizedTypeAst::Decimal128 { precision, scale }) => {
682 matches!(actual, DataType::Decimal128(actual_precision, actual_scale) if actual_precision == precision && actual_scale == scale)
683 }
684 TypeAst::Parameterized(ParameterizedTypeAst::Timestamp { unit, timezone }) => {
685 matches!(actual, DataType::Timestamp(actual_unit, actual_timezone)
686 if time_unit_matches(*unit, *actual_unit)
687 && actual_timezone.as_deref() == timezone.as_deref())
688 }
689 }
690}
691
692fn matches_primitive_type(expected: PrimitiveTypeAst, actual: &DataType) -> bool {
693 matches!(
694 (expected, actual),
695 (PrimitiveTypeAst::Boolean, DataType::Boolean)
696 | (PrimitiveTypeAst::Int8, DataType::Int8)
697 | (PrimitiveTypeAst::Int16, DataType::Int16)
698 | (PrimitiveTypeAst::Int32, DataType::Int32)
699 | (PrimitiveTypeAst::Int64, DataType::Int64)
700 | (PrimitiveTypeAst::Uint8, DataType::UInt8)
701 | (PrimitiveTypeAst::Uint16, DataType::UInt16)
702 | (PrimitiveTypeAst::Uint32, DataType::UInt32)
703 | (PrimitiveTypeAst::Uint64, DataType::UInt64)
704 | (PrimitiveTypeAst::Float32, DataType::Float32)
705 | (PrimitiveTypeAst::Float64, DataType::Float64)
706 | (PrimitiveTypeAst::Date32, DataType::Date32)
707 | (PrimitiveTypeAst::Date64, DataType::Date64)
708 | (PrimitiveTypeAst::Utf8, DataType::Utf8)
709 | (PrimitiveTypeAst::LargeUtf8, DataType::LargeUtf8)
710 | (PrimitiveTypeAst::Utf8View, DataType::Utf8View)
711 | (PrimitiveTypeAst::Binary, DataType::Binary)
712 | (PrimitiveTypeAst::LargeBinary, DataType::LargeBinary)
713 | (PrimitiveTypeAst::BinaryView, DataType::BinaryView)
714 )
715}
716
717fn time_unit_matches(expected: TimeUnitAst, actual: TimeUnit) -> bool {
718 matches!(
719 (expected, actual),
720 (TimeUnitAst::S, TimeUnit::Second)
721 | (TimeUnitAst::Ms, TimeUnit::Millisecond)
722 | (TimeUnitAst::Us, TimeUnit::Microsecond)
723 | (TimeUnitAst::Ns, TimeUnit::Nanosecond)
724 )
725}
726
727fn tagged_digest(prefix: &str, digest: blake3::Hash) -> String {
728 let hex = digest.to_hex();
729 let mut output = String::with_capacity(prefix.len() + hex.len());
730 output.push_str(prefix);
731 output.push_str(hex.as_str());
732 output
733}
734
735fn hash_part(hasher: &mut blake3::Hasher, bytes: &[u8]) {
736 hasher.update(&(bytes.len() as u64).to_le_bytes());
737 hasher.update(bytes);
738}
739
740fn hash_optional_part(hasher: &mut blake3::Hasher, bytes: Option<&[u8]>) {
741 match bytes {
742 Some(bytes) => {
743 hasher.update(&[1]);
744 hash_part(hasher, bytes);
745 }
746 None => {
747 hasher.update(&[0]);
748 }
749 }
750}
751
752fn hash_kernel(hasher: &mut blake3::Hasher, kernel: &KernelKind) {
753 hasher.update(&[kernel_tag(kernel)]);
754 hash_kernel_payload(hasher, kernel);
755}
756
757fn scalar_kernel_tag(kernel: &KernelKind) -> Option<u8> {
758 Some(match kernel {
759 KernelKind::Boolean => 0,
760 KernelKind::I8 => 1,
761 KernelKind::I16 => 2,
762 KernelKind::I32 => 3,
763 KernelKind::I64 => 4,
764 KernelKind::U8 => 5,
765 KernelKind::U16 => 6,
766 KernelKind::U32 => 7,
767 KernelKind::U64 => 8,
768 KernelKind::F32 => 9,
769 KernelKind::F64 => 10,
770 _ => return None,
771 })
772}
773
774fn temporal_kernel_tag(kernel: &KernelKind) -> Option<u8> {
775 match kernel {
776 KernelKind::Date32 => Some(11),
777 KernelKind::Date64 => Some(12),
778 KernelKind::Decimal128 { .. } => Some(13),
779 KernelKind::Timestamp(_) => Some(14),
780 _ => None,
781 }
782}
783
784fn variable_kernel_tag(kernel: &KernelKind) -> u8 {
785 match kernel {
786 KernelKind::Utf8 => 15,
787 KernelKind::LargeUtf8 => 16,
788 KernelKind::Binary => 17,
789 KernelKind::LargeBinary => 18,
790 KernelKind::Nested => 19,
791 KernelKind::NullOnly => 20,
792 KernelKind::Utf8View => 21,
793 KernelKind::BinaryView => 22,
794 _ => unreachable!("scalar and temporal kernels are handled first"),
795 }
796}
797
798fn kernel_tag(kernel: &KernelKind) -> u8 {
799 if let Some(tag) = scalar_kernel_tag(kernel) {
800 return tag;
801 }
802 if let Some(tag) = temporal_kernel_tag(kernel) {
803 return tag;
804 }
805 variable_kernel_tag(kernel)
806}
807
808fn hash_kernel_payload(hasher: &mut blake3::Hasher, kernel: &KernelKind) {
809 match kernel {
810 KernelKind::Decimal128 { precision, scale } => {
811 hasher.update(&[*precision, *scale as u8]);
812 }
813 KernelKind::Timestamp(unit) => hash_part(hasher, format!("{unit:?}").as_bytes()),
814 _ => {}
815 }
816}
817
818fn hash_optional_bound(hasher: &mut blake3::Hasher, bound: Option<&TypedBound>) {
819 let Some(bound) = bound else {
820 hasher.update(&[0]);
821 return;
822 };
823 hasher.update(&[1]);
824 match bound {
825 TypedBound::I64(value) => {
826 hasher.update(&[0]);
827 hasher.update(&value.to_le_bytes());
828 }
829 TypedBound::U64(value) => {
830 hasher.update(&[1]);
831 hasher.update(&value.to_le_bytes());
832 }
833 TypedBound::F32(value) => {
834 hasher.update(&[2]);
835 hasher.update(&value.to_bits().to_le_bytes());
836 }
837 TypedBound::F64(value) => {
838 hasher.update(&[3]);
839 hasher.update(&value.to_bits().to_le_bytes());
840 }
841 TypedBound::Decimal128 { value, scale } => {
842 hasher.update(&[4, *scale as u8]);
843 hasher.update(&value.to_le_bytes());
844 }
845 TypedBound::Timestamp { value, unit } => {
846 hasher.update(&[5]);
847 hasher.update(&value.to_le_bytes());
848 hash_part(hasher, format!("{unit:?}").as_bytes());
849 }
850 }
851}
852
853fn has_value_rules(rules: &RuleAst) -> bool {
854 rules.not_null
855 || rules.unique
856 || rules.min.is_some()
857 || rules.max.is_some()
858 || rules.nan.is_some()
859 || rules.pattern.is_some()
860 || rules.allowed.is_some()
861}
862
863fn has_runtime_rules(rules: &RuleAst) -> bool {
864 has_value_rules(rules)
865}
866
867pub(super) fn compile_rules(
868 column: &str,
869 data_type: &DataType,
870 kernel: &KernelKind,
871 source: &RuleAst,
872) -> Result<CompiledRules, ProofFrameError> {
873 let base_path = column_path(column);
874 if (source.min.is_some() || source.max.is_some()) && !kernel.supports_bounds() {
875 let field = if source.min.is_some() { "min" } else { "max" };
876 return Err(type_mismatch(
877 column,
878 data_type,
879 field,
880 "numeric, decimal, date, or timestamp",
881 ));
882 }
883 if source.pattern.is_some() && !kernel.supports_text_rules() {
884 return Err(type_mismatch(column, data_type, "pattern", "UTF-8"));
885 }
886 if source.allowed.is_some() && !kernel.supports_text_rules() {
887 return Err(type_mismatch(column, data_type, "allowed", "UTF-8"));
888 }
889 if source.nan.is_some() && !matches!(kernel, KernelKind::F32 | KernelKind::F64) {
890 return Err(type_mismatch(column, data_type, "nan", "floating point"));
891 }
892 if source.unique && !kernel.supports_unique() {
893 return Err(type_mismatch(
894 column,
895 data_type,
896 "unique",
897 "canonically encodable",
898 ));
899 }
900
901 let min_path = append_path(&base_path, "min");
902 let max_path = append_path(&base_path, "max");
903 let min = source
904 .min
905 .as_ref()
906 .map(|bound| parse_bound(bound, data_type, &min_path))
907 .transpose()?;
908 let max = source
909 .max
910 .as_ref()
911 .map(|bound| parse_bound(bound, data_type, &max_path))
912 .transpose()?;
913 if min
914 .as_ref()
915 .zip(max.as_ref())
916 .is_some_and(|(minimum, maximum)| minimum.compare(maximum) == Some(Ordering::Greater))
917 {
918 return Err(ProofFrameError::contract(
919 ErrorCode::ContractInvalidBound,
920 format!("Contract minimum exceeds maximum for column `{column}`"),
921 Some(base_path.clone()),
922 ));
923 }
924
925 let pattern = source
926 .pattern
927 .as_ref()
928 .map(|pattern| {
929 Regex::new(pattern).map_err(|error| {
930 ProofFrameError::contract(
931 ErrorCode::ContractInvalidBound,
932 format!("Invalid regular expression for column `{column}`: {error}"),
933 Some(append_path(&base_path, "pattern")),
934 )
935 })
936 })
937 .transpose()?;
938 let allowed = source.allowed.as_ref().map(|values| {
939 let mut compiled = HashSet::with_capacity_and_hasher(values.len(), RandomState::new());
940 compiled.extend(values.iter().map(|value| value.clone().into_boxed_str()));
941 Arc::new(compiled)
942 });
943
944 Ok(CompiledRules {
945 required: source.required,
946 not_null: source.not_null,
947 unique: source.unique,
948 min,
949 max,
950 nan: source.nan.into(),
951 validate_nan: source.nan.is_some() || source.min.is_some() || source.max.is_some(),
952 pattern,
953 allowed,
954 })
955}
956
957fn type_mismatch(
958 column: &str,
959 data_type: &DataType,
960 rule: &str,
961 expected: &str,
962) -> ProofFrameError {
963 let path = append_path(&column_path(column), rule);
964 ProofFrameError::contract(
965 ErrorCode::ContractTypeMismatch,
966 format!(
967 "Rule `{rule}` on column `{column}` requires {expected} values, found `{data_type}`"
968 ),
969 Some(path),
970 )
971}