1use crate::nary::{NaryRulePattern, PatternVar};
24
25pub const NARY_MAX_BODY_ATOMS: usize = 8;
27pub const NARY_MAX_JOIN_VARS: usize = 8;
28pub const NARY_MAX_ATOM_ARITY: usize = 8;
29pub const NARY_MAX_HEAD_ARITY: usize = 8;
35
36const JOIN_FLAG: u32 = 1 << 31;
37
38#[derive(Debug, Clone, PartialEq, Eq)]
40pub enum NaryLayoutError {
41 EmptyBatch,
42 EmptyBody {
43 pattern: usize,
44 },
45 TooManyBodyAtoms {
46 pattern: usize,
47 atoms: usize,
48 },
49 AtomArityOutOfRange {
50 pattern: usize,
51 atom: usize,
52 arity: usize,
53 },
54 JoinIndexOutOfRange {
55 pattern: usize,
56 atom: usize,
57 join: u8,
58 },
59 HeadIndexOutOfRange {
60 pattern: usize,
61 atom: usize,
62 head: u8,
63 },
64 HeadArityOutOfRange {
65 pattern: usize,
66 head_arity: usize,
67 },
68}
69
70impl std::fmt::Display for NaryLayoutError {
71 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72 match self {
73 Self::EmptyBatch => write!(f, "pattern batch is empty"),
74 Self::EmptyBody { pattern } => {
75 write!(f, "pattern {pattern} has an empty body")
76 }
77 Self::TooManyBodyAtoms { pattern, atoms } => write!(
78 f,
79 "pattern {pattern} has {atoms} body atoms; device bound is \
80 {NARY_MAX_BODY_ATOMS}"
81 ),
82 Self::AtomArityOutOfRange {
83 pattern,
84 atom,
85 arity,
86 } => write!(
87 f,
88 "pattern {pattern} atom {atom} arity {arity} outside \
89 1..={NARY_MAX_ATOM_ARITY}"
90 ),
91 Self::JoinIndexOutOfRange {
92 pattern,
93 atom,
94 join,
95 } => write!(
96 f,
97 "pattern {pattern} atom {atom} join index {join} >= device \
98 bound {NARY_MAX_JOIN_VARS}"
99 ),
100 Self::HeadIndexOutOfRange {
101 pattern,
102 atom,
103 head,
104 } => write!(
105 f,
106 "pattern {pattern} atom {atom} head index {head} >= the \
107 pattern's head arity"
108 ),
109 Self::HeadArityOutOfRange {
110 pattern,
111 head_arity,
112 } => write!(
113 f,
114 "pattern {pattern} head arity {head_arity} exceeds the \
115 device bound {NARY_MAX_HEAD_ARITY}"
116 ),
117 }
118 }
119}
120
121impl std::error::Error for NaryLayoutError {}
122
123#[derive(Debug, Clone, PartialEq, Eq)]
131pub struct NaryPatternBatchLayout {
132 pub head_arity: Vec<u32>,
133 pub join_count: Vec<u32>,
134 pub body_offset: Vec<u32>,
135 pub body_len: Vec<u32>,
136 pub atom_candidate_slot: Vec<u32>,
137 pub atom_arity: Vec<u32>,
138 pub atom_binding_offset: Vec<u32>,
139 pub binding_codes: Vec<u32>,
140}
141
142pub fn binding_code(var: PatternVar) -> u32 {
144 match var {
145 PatternVar::Head(i) => u32::from(i),
146 PatternVar::Join(j) => JOIN_FLAG | u32::from(j),
147 }
148}
149
150pub fn decode_binding(code: u32) -> PatternVar {
152 if code & JOIN_FLAG != 0 {
153 PatternVar::Join((code & 0xFF) as u8)
154 } else {
155 PatternVar::Head((code & 0xFF) as u8)
156 }
157}
158
159pub fn flatten_patterns(
165 patterns: &[NaryRulePattern],
166) -> Result<NaryPatternBatchLayout, NaryLayoutError> {
167 if patterns.is_empty() {
168 return Err(NaryLayoutError::EmptyBatch);
169 }
170 let mut layout = NaryPatternBatchLayout {
171 head_arity: Vec::with_capacity(patterns.len()),
172 join_count: Vec::with_capacity(patterns.len()),
173 body_offset: Vec::with_capacity(patterns.len()),
174 body_len: Vec::with_capacity(patterns.len()),
175 atom_candidate_slot: Vec::new(),
176 atom_arity: Vec::new(),
177 atom_binding_offset: Vec::new(),
178 binding_codes: Vec::new(),
179 };
180 for (p, pattern) in patterns.iter().enumerate() {
181 if pattern.body.is_empty() {
182 return Err(NaryLayoutError::EmptyBody { pattern: p });
183 }
184 if pattern.body.len() > NARY_MAX_BODY_ATOMS {
185 return Err(NaryLayoutError::TooManyBodyAtoms {
186 pattern: p,
187 atoms: pattern.body.len(),
188 });
189 }
190 if usize::from(pattern.head_arity) > NARY_MAX_HEAD_ARITY {
195 return Err(NaryLayoutError::HeadArityOutOfRange {
196 pattern: p,
197 head_arity: usize::from(pattern.head_arity),
198 });
199 }
200 let mut joins = 0u32;
201 layout.head_arity.push(u32::from(pattern.head_arity));
202 layout.body_offset.push(
203 u32::try_from(layout.atom_candidate_slot.len())
204 .expect("atom count exceeds u32; the batch bounds make this unreachable"),
205 );
206 layout.body_len.push(
207 u32::try_from(pattern.body.len())
208 .expect("body length is bounded by NARY_MAX_BODY_ATOMS"),
209 );
210 for (a, atom) in pattern.body.iter().enumerate() {
211 let arity = atom.bindings.len();
212 if arity == 0 || arity > NARY_MAX_ATOM_ARITY {
213 return Err(NaryLayoutError::AtomArityOutOfRange {
214 pattern: p,
215 atom: a,
216 arity,
217 });
218 }
219 layout.atom_candidate_slot.push(atom.candidate_slot);
220 layout.atom_arity.push(arity as u32);
221 layout.atom_binding_offset.push(
222 u32::try_from(layout.binding_codes.len())
223 .expect("binding-code count exceeds u32; unreachable at sane scales"),
224 );
225 for binding in &atom.bindings {
226 match *binding {
227 PatternVar::Head(i) => {
228 if u32::from(i) >= u32::from(pattern.head_arity) {
229 return Err(NaryLayoutError::HeadIndexOutOfRange {
230 pattern: p,
231 atom: a,
232 head: i,
233 });
234 }
235 }
236 PatternVar::Join(j) => {
237 if usize::from(j) >= NARY_MAX_JOIN_VARS {
238 return Err(NaryLayoutError::JoinIndexOutOfRange {
239 pattern: p,
240 atom: a,
241 join: j,
242 });
243 }
244 joins = joins.max(u32::from(j) + 1);
245 }
246 }
247 layout.binding_codes.push(binding_code(*binding));
248 }
249 }
250 layout.join_count.push(joins);
251 }
252 Ok(layout)
253}
254
255#[derive(Debug, Clone, PartialEq, Eq)]
264pub struct FlatRelation {
265 pub arity: u32,
266 pub row_count: u32,
267 pub values: Vec<u64>,
269}
270
271impl FlatRelation {
272 pub fn from_rows(rows: &[Vec<u64>], arity: u32) -> Self {
274 let mut values = vec![0u64; rows.len() * arity as usize];
275 for (row_index, row) in rows.iter().enumerate() {
276 assert_eq!(row.len(), arity as usize, "row arity mismatch");
277 for (position, value) in row.iter().enumerate() {
278 values[position * rows.len() + row_index] = *value;
279 }
280 }
281 Self {
282 arity,
283 row_count: u32::try_from(rows.len()).expect("relation row count exceeds u32"),
284 values,
285 }
286 }
287}
288
289pub fn score_pattern_flat(
298 layout: &NaryPatternBatchLayout,
299 pattern_index: usize,
300 candidates: &[FlatRelation],
301 example: &[u64],
302) -> bool {
303 let body_offset = layout.body_offset[pattern_index] as usize;
304 let body_len = layout.body_len[pattern_index] as usize;
305 debug_assert_eq!(example.len(), layout.head_arity[pattern_index] as usize);
306
307 let mut joins = [0u64; NARY_MAX_JOIN_VARS];
308 let mut bound: u32 = 0;
309 let mut row_cursor = [0u32; NARY_MAX_BODY_ATOMS];
310 let mut depth_mask = [0u32; NARY_MAX_BODY_ATOMS];
311
312 let mut depth = 0usize;
313 loop {
314 if depth == body_len {
315 return true;
316 }
317 let atom = body_offset + depth;
318 let relation = &candidates[layout.atom_candidate_slot[atom] as usize];
319 let arity = layout.atom_arity[atom] as usize;
320 let binding_offset = layout.atom_binding_offset[atom] as usize;
321 debug_assert_eq!(relation.arity as usize, arity);
322
323 let rows = relation.row_count as usize;
324 let mut descended = false;
325 while row_cursor[depth] < relation.row_count {
326 let row = row_cursor[depth] as usize;
327 let mut mask: u32 = 0;
331 let mut matched = true;
332 for position in 0..arity {
333 let value = relation.values[position * rows + row];
335 let code = layout.binding_codes[binding_offset + position];
336 if code & JOIN_FLAG != 0 {
337 let j = (code & 0xFF) as usize;
338 let bit = 1u32 << j;
339 if bound & bit != 0 {
340 if joins[j] != value {
341 matched = false;
342 break;
343 }
344 } else {
345 joins[j] = value;
346 bound |= bit;
347 mask |= bit;
348 }
349 } else {
350 let head = (code & 0xFF) as usize;
351 if example[head] != value {
352 matched = false;
353 break;
354 }
355 }
356 }
357 if matched {
358 depth_mask[depth] = mask;
359 depth += 1;
360 if depth < body_len {
361 row_cursor[depth] = 0;
362 }
363 descended = true;
364 break;
365 }
366 bound &= !mask;
367 row_cursor[depth] += 1;
368 }
369 if descended {
370 continue;
371 }
372 if depth == 0 {
374 return false;
375 }
376 depth -= 1;
377 bound &= !depth_mask[depth];
378 row_cursor[depth] += 1;
379 }
380}
381
382pub fn score_pattern_flat_coverage(
384 layout: &NaryPatternBatchLayout,
385 pattern_index: usize,
386 candidates: &[FlatRelation],
387 positives: &[Vec<u64>],
388 negatives: &[Vec<u64>],
389) -> (u32, u32) {
390 let count = |examples: &[Vec<u64>]| -> u32 {
391 examples
392 .iter()
393 .filter(|example| score_pattern_flat(layout, pattern_index, candidates, example))
394 .count() as u32
395 };
396 (count(positives), count(negatives))
397}
398
399#[cfg(test)]
400mod tests {
401 use super::*;
402 use crate::nary::{canonical_binary_pattern, BodyAtomPattern};
403 use crate::nary_reference::{score_pattern_reference, HostRelation, ReferenceCoverage};
404 use crate::types::Topology;
405
406 fn flat_pairs(rows: &[(u64, u64)]) -> FlatRelation {
407 FlatRelation::from_rows(
408 &rows.iter().map(|(a, b)| vec![*a, *b]).collect::<Vec<_>>(),
409 2,
410 )
411 }
412
413 fn host_pairs(rows: &[(u64, u64)]) -> HostRelation {
414 HostRelation {
415 rows: rows.iter().map(|(a, b)| vec![*a, *b]).collect(),
416 }
417 }
418
419 #[test]
420 fn binding_code_round_trips() {
421 for var in [
422 PatternVar::Head(0),
423 PatternVar::Head(7),
424 PatternVar::Join(0),
425 PatternVar::Join(7),
426 ] {
427 assert_eq!(decode_binding(binding_code(var)), var);
428 }
429 }
430
431 #[test]
435 fn flat_relation_is_column_major() {
436 let relation = FlatRelation::from_rows(&[vec![1, 2], vec![3, 4], vec![5, 6]], 2);
437 assert_eq!(relation.values, vec![1, 3, 5, 2, 4, 6]);
438 let rows = relation.row_count as usize;
439 assert_eq!(relation.values[1], 3); assert_eq!(relation.values[rows + 2], 6); }
442
443 #[test]
446 fn flat_walk_matches_reference_on_kernel_fixture() {
447 let flat_candidates = vec![
448 flat_pairs(&[(1, 2), (2, 3)]),
449 flat_pairs(&[(2, 4), (3, 5), (4, 6)]),
450 ];
451 let host_candidates = vec![
452 host_pairs(&[(1, 2), (2, 3)]),
453 host_pairs(&[(2, 4), (3, 5), (4, 6)]),
454 ];
455 let positives = vec![vec![1u64, 4u64], vec![2, 5]];
456 let negatives = vec![vec![7u64, 8u64]];
457
458 let mut patterns = Vec::new();
459 for topology in Topology::ALL {
460 for left in 0..2u32 {
461 for right in 0..2u32 {
462 patterns.push(canonical_binary_pattern(topology, left, right));
463 }
464 }
465 }
466 let layout = flatten_patterns(&patterns).expect("bounded batch flattens");
467 for (index, pattern) in patterns.iter().enumerate() {
468 let reference =
469 score_pattern_reference(pattern, &host_candidates, &positives, &negatives);
470 let (pos, neg) = score_pattern_flat_coverage(
471 &layout,
472 index,
473 &flat_candidates,
474 &positives,
475 &negatives,
476 );
477 assert_eq!(
478 ReferenceCoverage {
479 positives_covered: pos,
480 negatives_covered: neg
481 },
482 reference,
483 "pattern {index} diverges from the recursive reference"
484 );
485 }
486 }
487
488 #[test]
491 fn flat_walk_matches_reference_on_ternary_fixture() {
492 use PatternVar::{Head, Join};
493 let pattern = NaryRulePattern {
494 head_arity: 3,
495 body: vec![
496 BodyAtomPattern {
497 candidate_slot: 0,
498 bindings: vec![Head(0), Head(1), Join(0)],
499 },
500 BodyAtomPattern {
501 candidate_slot: 1,
502 bindings: vec![Join(0), Head(2)],
503 },
504 ],
505 };
506 let ternary_rows = vec![vec![1u64, 2, 9], vec![4, 5, 8]];
507 let binary_rows = [(9u64, 3u64), (8, 7)];
508 let flat_candidates = vec![
509 FlatRelation::from_rows(&ternary_rows, 3),
510 flat_pairs(&binary_rows),
511 ];
512 let host_candidates = vec![
513 HostRelation {
514 rows: ternary_rows.clone(),
515 },
516 host_pairs(&binary_rows),
517 ];
518 let positives = vec![vec![1u64, 2, 3], vec![4, 5, 6], vec![1, 2, 7]];
519
520 let layout = flatten_patterns(std::slice::from_ref(&pattern)).unwrap();
521 let reference = score_pattern_reference(&pattern, &host_candidates, &positives, &[]);
522 let (pos, neg) = score_pattern_flat_coverage(&layout, 0, &flat_candidates, &positives, &[]);
523 assert_eq!(pos, reference.positives_covered);
524 assert_eq!(neg, reference.negatives_covered);
525 assert_eq!(pos, 1);
526 }
527
528 #[test]
531 fn flat_walk_backtracks_across_atoms() {
532 use PatternVar::{Head, Join};
533 let pattern = NaryRulePattern {
534 head_arity: 2,
535 body: vec![
536 BodyAtomPattern {
537 candidate_slot: 0,
538 bindings: vec![Head(0), Join(0)],
539 },
540 BodyAtomPattern {
541 candidate_slot: 1,
542 bindings: vec![Join(0), Head(1)],
543 },
544 ],
545 };
546 let candidates = vec![flat_pairs(&[(1, 8), (1, 9)]), flat_pairs(&[(9, 2)])];
547 let layout = flatten_patterns(std::slice::from_ref(&pattern)).unwrap();
548 assert!(score_pattern_flat(&layout, 0, &candidates, &[1, 2]));
549 assert!(!score_pattern_flat(&layout, 0, &candidates, &[1, 3]));
550 }
551
552 #[test]
555 fn undo_masks_do_not_leak_bindings_within_a_row() {
556 use PatternVar::Join;
557 let pattern = NaryRulePattern {
560 head_arity: 1,
561 body: vec![BodyAtomPattern {
562 candidate_slot: 0,
563 bindings: vec![Join(0), Join(0)],
564 }],
565 };
566 let candidates = vec![flat_pairs(&[(3, 4), (5, 5)])];
567 let layout = flatten_patterns(std::slice::from_ref(&pattern)).unwrap();
568 assert!(score_pattern_flat(&layout, 0, &candidates, &[0]));
571 }
572
573 #[test]
574 fn device_bounds_are_refused_typed() {
575 use PatternVar::{Head, Join};
576 assert_eq!(flatten_patterns(&[]), Err(NaryLayoutError::EmptyBatch));
577
578 let empty_body = NaryRulePattern {
579 head_arity: 2,
580 body: vec![],
581 };
582 assert_eq!(
583 flatten_patterns(std::slice::from_ref(&empty_body)),
584 Err(NaryLayoutError::EmptyBody { pattern: 0 })
585 );
586
587 let atom = BodyAtomPattern {
588 candidate_slot: 0,
589 bindings: vec![Head(0), Head(1)],
590 };
591 let too_many_atoms = NaryRulePattern {
592 head_arity: 2,
593 body: vec![atom.clone(); NARY_MAX_BODY_ATOMS + 1],
594 };
595 assert!(matches!(
596 flatten_patterns(std::slice::from_ref(&too_many_atoms)),
597 Err(NaryLayoutError::TooManyBodyAtoms { pattern: 0, .. })
598 ));
599
600 let wide_head = NaryRulePattern {
603 head_arity: (NARY_MAX_HEAD_ARITY + 1) as u8,
604 body: vec![BodyAtomPattern {
605 candidate_slot: 0,
606 bindings: vec![Head(0), Head(1)],
607 }],
608 };
609 assert!(matches!(
610 flatten_patterns(std::slice::from_ref(&wide_head)),
611 Err(NaryLayoutError::HeadArityOutOfRange { pattern: 0, .. })
612 ));
613
614 let wide_atom = NaryRulePattern {
615 head_arity: 2,
616 body: vec![BodyAtomPattern {
617 candidate_slot: 0,
618 bindings: vec![Head(0); NARY_MAX_ATOM_ARITY + 1],
619 }],
620 };
621 assert!(matches!(
622 flatten_patterns(std::slice::from_ref(&wide_atom)),
623 Err(NaryLayoutError::AtomArityOutOfRange { .. })
624 ));
625
626 let join_out_of_range = NaryRulePattern {
627 head_arity: 2,
628 body: vec![BodyAtomPattern {
629 candidate_slot: 0,
630 bindings: vec![Head(0), Join(NARY_MAX_JOIN_VARS as u8)],
631 }],
632 };
633 assert!(matches!(
634 flatten_patterns(std::slice::from_ref(&join_out_of_range)),
635 Err(NaryLayoutError::JoinIndexOutOfRange { .. })
636 ));
637
638 let head_out_of_range = NaryRulePattern {
639 head_arity: 2,
640 body: vec![BodyAtomPattern {
641 candidate_slot: 0,
642 bindings: vec![Head(2), Head(0)],
643 }],
644 };
645 assert!(matches!(
646 flatten_patterns(std::slice::from_ref(&head_out_of_range)),
647 Err(NaryLayoutError::HeadIndexOutOfRange { .. })
648 ));
649 }
650
651 #[test]
653 fn layout_offsets_tile_the_flat_arrays() {
654 let patterns = vec![
655 canonical_binary_pattern(Topology::Chain, 0, 1),
656 canonical_binary_pattern(Topology::Star, 1, 0),
657 ];
658 let layout = flatten_patterns(&patterns).unwrap();
659 assert_eq!(layout.body_offset, vec![0, 2]);
660 assert_eq!(layout.body_len, vec![2, 2]);
661 assert_eq!(layout.atom_candidate_slot.len(), 4);
662 assert_eq!(layout.atom_binding_offset, vec![0, 2, 4, 6]);
663 assert_eq!(layout.binding_codes.len(), 8);
664 assert_eq!(layout.join_count.len(), 2);
665 }
666}