Skip to main content

proofframe/contract/
compile.rs

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/// Arrow-specialized kernel selected once during contract compilation.
20#[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/// Runtime NaN behavior compiled for floating-point kernels.
117#[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/// Semantically checked rules ready for typed execution.
133#[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/// One schema-resolved column in the execution intermediate representation.
194#[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/// Contract execution IR compiled in Arrow schema order.
225#[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    /// Resolve and type-check a syntax tree before any record batch is scanned.
237    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    /// Compile either frozen V1 syntax or the relational V2 contract language.
286    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    /// Domain-separated digest of the Arrow schema used during compilation.
382    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    /// Domain-separated digest of the schema-resolved execution plan.
390    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}