1#![allow(missing_docs)]
26
27use std::collections::HashMap;
28
29use crate::error::{Error, Result};
30use crate::ir::{
31 Annotation, AnnotationKind, Constraint, InfinityModule, Inode, InterpPart, NodeId, PatchOp,
32 ProvValue, Provenance, SchemaFieldSpec, TypeSpec,
33};
34use crate::layout::canonical_f64;
35use crate::unit::UnitTag;
36
37pub(crate) type TypeKey = (u32, Vec<(u32, u32)>);
39
40pub(crate) type BlobKey = ([u8; 32], u64, Option<Vec<u8>>);
42
43#[derive(Debug, Clone, PartialEq, Eq)]
45pub(crate) struct BlobPayload {
46 pub hash: [u8; 32],
47 pub size: u64,
48 pub payload: Option<Vec<u8>>,
49}
50
51#[derive(Debug, Clone, PartialEq)]
53pub(crate) struct PreparedType {
54 pub spec: TypeSpec,
55 pub aux: Vec<(u32, u32)>,
56}
57
58#[derive(Debug, Clone, Default)]
60pub(crate) struct Prepared {
61 pub nodes: Vec<Inode>,
62 pub root: Option<NodeId>,
63 pub bindings: Vec<(String, NodeId)>,
64 pub annotations: Vec<Annotation>,
65 pub provenance: Option<Provenance>,
66 pub strings: Vec<String>,
67 pub units: Vec<UnitTag>,
68 pub types: Vec<PreparedType>,
69 pub blobs: Vec<BlobPayload>,
70 pub hashes: Vec<[u8; 32]>,
71 pub numbers: Vec<i64>,
72 pub secrets: Vec<Vec<u8>>,
73 pub extensions: Vec<(u64, Vec<u8>)>,
74 pub string_index: HashMap<String, u32>,
75 pub unit_index: HashMap<UnitTag, u32>,
76 pub type_index: HashMap<TypeKey, u32>,
77 pub blob_index: HashMap<BlobKey, u32>,
78 pub hash_index: HashMap<[u8; 32], u32>,
79 pub number_index: HashMap<i64, u32>,
80}
81
82pub(crate) fn is_canonical_nan(value: f64) -> bool {
84 value.to_bits() == crate::layout::CANONICAL_NAN_BITS
85}
86
87pub(crate) fn prepare(module: &InfinityModule) -> Result<Prepared> {
89 module.validate_structure(u32::MAX as u64, 0)?;
90
91 let canonical_nodes: Vec<Inode> = module.nodes().iter().map(normalize).collect();
93
94 let canonical_nodes = promote_secret_locators(canonical_nodes, module);
103
104 let order = canonical_order(&canonical_nodes, module)?;
106 let mut remap = vec![NodeId::NONE; canonical_nodes.len()];
107 for (new_index, old) in order.iter().enumerate() {
108 remap[old.index()] = NodeId(new_index as u32);
109 }
110
111 let nodes: Vec<Inode> = order
113 .iter()
114 .map(|old| remap_node(&canonical_nodes[old.index()], &remap))
115 .collect::<Result<Vec<_>>>()?;
116
117 let root = module.root().map(|id| remap[id.index()]);
118 let mut bindings: Vec<(String, NodeId)> = module
119 .bindings()
120 .iter()
121 .map(|(name, id)| (name.clone(), remap[id.index()]))
122 .collect();
123 bindings.sort_by(|a, b| a.0.cmp(&b.0));
124
125 let secret_literals: Vec<String> = canonical_nodes
127 .iter()
128 .filter_map(|node| match node {
129 Inode::Secret(secret) => Some(String::from_utf8_lossy(secret.expose()).into_owned()),
130 _ => None,
131 })
132 .collect();
133
134 let annotations: Vec<Annotation> = module
135 .annotations()
136 .iter()
137 .map(|annotation| {
138 let node = remap[annotation.node.index()];
139 let kind = match &annotation.kind {
140 AnnotationKind::Intent(text) => {
141 AnnotationKind::Intent(redact_text(text, &secret_literals))
142 }
143 AnnotationKind::Constraint { source, expr } => AnnotationKind::Constraint {
144 source: redact_text(source, &secret_literals),
145 expr: remap[expr.index()],
146 },
147 AnnotationKind::AiImmutable(flag) => AnnotationKind::AiImmutable(*flag),
148 };
149 Annotation { node, kind }
150 })
151 .collect();
152
153 let mut strings = Vec::new();
155 for node in &nodes {
156 node.strings(&mut strings);
157 }
158 for node in &nodes {
160 if let Inode::Quantity {
161 unit: UnitTag::Named(name),
162 ..
163 } = node
164 {
165 strings.push(name.clone());
166 }
167 }
168 for (name, _) in &bindings {
169 strings.push(name.clone());
170 }
171 for annotation in &annotations {
172 match &annotation.kind {
173 AnnotationKind::Intent(text) => strings.push(text.clone()),
174 AnnotationKind::Constraint { source, .. } => strings.push(source.clone()),
175 AnnotationKind::AiImmutable(_) => {}
176 }
177 }
178 let provenance = module
179 .provenance()
180 .cloned()
181 .map(|prov| redact_provenance(prov, &secret_literals))
182 .transpose()?;
183 if let Some(prov) = &provenance {
184 if let Some(producer) = &prov.producer {
185 strings.push(producer.clone());
186 }
187 if let Some(model) = &prov.model {
188 strings.push(model.clone());
189 }
190 if let Some(tool) = &prov.tool {
191 strings.push(tool.clone());
192 }
193 for entry in &prov.entries {
194 strings.push(entry.key.clone());
195 if let ProvValue::Text(text) = &entry.value {
196 strings.push(text.clone());
197 }
198 }
199 }
200 strings.sort();
201 strings.dedup();
202 let string_index: HashMap<String, u32> = strings
203 .iter()
204 .enumerate()
205 .map(|(index, value)| (value.clone(), index as u32))
206 .collect();
207
208 let mut units: Vec<UnitTag> = Vec::new();
209 for node in &nodes {
210 if let Inode::Quantity { unit, .. } = node {
211 units.push(unit.clone());
212 }
213 }
214 units.sort_by(|a, b| (a.discriminant(), a.name()).cmp(&(b.discriminant(), b.name())));
215 units.dedup();
216 let unit_index: HashMap<UnitTag, u32> = units
217 .iter()
218 .enumerate()
219 .map(|(index, unit)| (unit.clone(), index as u32))
220 .collect();
221
222 let mut type_payloads: Vec<(u32, Vec<(u32, u32)>)> = Vec::new();
224 for node in &nodes {
225 if let Inode::SchemaField(spec) = node {
226 type_payloads.push(type_payload(spec, &string_index)?);
227 }
228 }
229 type_payloads.sort();
230 type_payloads.dedup();
231 let mut types = Vec::with_capacity(type_payloads.len());
232 let mut type_index = HashMap::new();
233 for (index, (code, aux)) in type_payloads.iter().enumerate() {
234 let spec = type_from_prepared(*code, aux, &strings)?;
235 types.push(PreparedType {
236 spec,
237 aux: aux.clone(),
238 });
239 type_index.insert((*code, aux.clone()), index as u32);
240 }
241
242 let mut blobs: Vec<BlobPayload> = Vec::new();
244 for node in &nodes {
245 if let Inode::Blob {
246 hash,
247 size,
248 tainted,
249 payload,
250 ..
251 } = node
252 {
253 if *tainted && payload.is_some() {
254 return Err(Error::Encode("a tainted blob must not carry a payload"));
255 }
256 blobs.push(BlobPayload {
257 hash: *hash,
258 size: *size,
259 payload: if *tainted { None } else { payload.clone() },
260 });
261 }
262 }
263 blobs.sort_by_key(|blob| (blob.hash, blob.size));
264 blobs.dedup();
265 let blob_index: HashMap<BlobKey, u32> = blobs
266 .iter()
267 .enumerate()
268 .map(|(index, blob)| ((blob.hash, blob.size, blob.payload.clone()), index as u32))
269 .collect();
270
271 let mut hashes: Vec<[u8; 32]> = Vec::new();
273 if let Some(prov) = &provenance {
274 if let Some(parent) = prov.parent {
275 hashes.push(parent);
276 }
277 for entry in &prov.entries {
278 if let ProvValue::Hash(hash) = &entry.value {
279 hashes.push(*hash);
280 }
281 }
282 }
283 hashes.sort();
284 hashes.dedup();
285 let hash_index: HashMap<[u8; 32], u32> = hashes
286 .iter()
287 .enumerate()
288 .map(|(index, hash)| (*hash, index as u32))
289 .collect();
290
291 let mut numbers: Vec<i64> = Vec::new();
293 for node in &nodes {
294 if let Inode::Constrained {
295 constraint: Constraint::MinMax { min, max },
296 ..
297 } = node
298 {
299 numbers.extend(min.iter().copied());
300 numbers.extend(max.iter().copied());
301 }
302 }
303 if let Some(prov) = &provenance {
304 for entry in &prov.entries {
305 if let ProvValue::Int(value) = &entry.value {
306 numbers.push(*value);
307 }
308 }
309 }
310 numbers.sort_unstable();
311 numbers.dedup();
312 let number_index: HashMap<i64, u32> = numbers
313 .iter()
314 .enumerate()
315 .map(|(index, value)| (*value, index as u32))
316 .collect();
317
318 let secrets: Vec<Vec<u8>> = nodes
320 .iter()
321 .filter_map(|node| match node {
322 Inode::Secret(secret) => Some(secret.expose().to_vec()),
323 _ => None,
324 })
325 .collect();
326
327 let mut extensions: Vec<(u64, Vec<u8>)> = module.extensions().to_vec();
329 extensions.sort_by_key(|(tag, _)| *tag);
330 for window in extensions.windows(2) {
331 if window[0].0 == window[1].0 {
332 return Err(Error::Encode("duplicate extension tag"));
333 }
334 }
335
336 let prepared = Prepared {
337 nodes,
338 root,
339 bindings,
340 annotations: sorted_annotations(annotations, &string_index),
341 provenance: provenance.map(sorted_provenance).transpose()?,
342 strings,
343 units,
344 types,
345 blobs,
346 hashes,
347 numbers,
348 secrets,
349 extensions,
350 string_index,
351 unit_index,
352 type_index,
353 blob_index,
354 hash_index,
355 number_index,
356 };
357 Ok(prepared)
358}
359
360fn normalize(node: &Inode) -> Inode {
361 match node {
362 Inode::Float(value) => Inode::Float(canonical_f64(*value)),
363 Inode::Table { fields } => {
364 let mut fields = fields.clone();
365 fields.sort_by(|a, b| a.0.cmp(&b.0));
366 Inode::Table { fields }
367 }
368 Inode::SchemaTable { fields, closed } => {
369 let mut fields = fields.clone();
370 fields.sort_by(|a, b| a.0.cmp(&b.0));
371 Inode::SchemaTable {
372 fields,
373 closed: *closed,
374 }
375 }
376 Inode::SchemaField(spec) => Inode::SchemaField(SchemaFieldSpec {
377 typ: normalize_type(&spec.typ),
378 ..spec.clone()
379 }),
380 Inode::Constrained {
381 target,
382 constraint: Constraint::Enum(items),
383 } => Inode::Constrained {
384 target: *target,
385 constraint: Constraint::Enum(items.clone()),
386 },
387 other => other.clone(),
388 }
389}
390
391pub const SECRET_MARKER: &str = "***SECRET***";
394
395pub(crate) fn redact_text(text: &str, secrets: &[String]) -> String {
408 let marker_tainted = secrets
409 .iter()
410 .any(|secret| !secret.is_empty() && SECRET_MARKER.contains(secret.as_str()));
411 let mut out = text.to_string();
412 for _ in 0..8 {
418 let mut matched: Vec<&String> = secrets
419 .iter()
420 .filter(|secret| !secret.is_empty() && out.contains(secret.as_str()))
421 .collect();
422 if matched.is_empty() {
423 return out;
424 }
425 if marker_tainted {
426 return SECRET_MARKER.to_string();
427 }
428 matched.sort_by(|a, b| b.len().cmp(&a.len()).then_with(|| a.cmp(b)));
431 for secret in matched {
432 out = out.replace(secret.as_str(), SECRET_MARKER);
433 }
434 }
435 SECRET_MARKER.to_string()
436}
437
438fn redact_provenance(mut prov: Provenance, secrets: &[String]) -> Result<Provenance> {
439 if secrets.is_empty() {
440 return Ok(prov);
441 }
442 prov.producer = prov.producer.map(|text| redact_text(&text, secrets));
443 prov.model = prov.model.map(|text| redact_text(&text, secrets));
444 prov.tool = prov.tool.map(|text| redact_text(&text, secrets));
445 for entry in &mut prov.entries {
446 entry.key = redact_text(&entry.key, secrets);
449 if let ProvValue::Text(text) = &entry.value {
450 entry.value = ProvValue::Text(redact_text(text, secrets));
451 }
452 }
453 Ok(prov)
454}
455
456fn normalize_type(typ: &TypeSpec) -> TypeSpec {
457 typ.clone()
458}
459
460fn promote_secret_locators(mut nodes: Vec<Inode>, module: &InfinityModule) -> Vec<Inode> {
462 if nodes.is_empty() {
463 return nodes;
464 }
465 let mut exposed = vec![false; nodes.len()];
466 let mut mark = |id: NodeId| {
467 if let Some(slot) = exposed.get_mut(id.index()) {
468 *slot = true;
469 }
470 };
471 if let Some(root) = module.root() {
472 mark(root);
473 }
474 for (_, id) in module.bindings() {
475 mark(*id);
476 }
477 for annotation in module.annotations() {
478 mark(annotation.node);
479 if let AnnotationKind::Constraint { expr, .. } = &annotation.kind {
480 mark(*expr);
481 }
482 }
483 let mut secret_only = vec![true; nodes.len()];
486 let mut referenced = vec![false; nodes.len()];
487 for node in &nodes {
488 for child in node.children() {
489 let Some(slot) = secret_only.get_mut(child.index()) else {
490 continue;
491 };
492 referenced[child.index()] = true;
493 if !node.is_secret() {
494 *slot = false;
495 }
496 }
497 }
498 for (index, node) in nodes.iter_mut().enumerate() {
499 let eligible = referenced[index] && secret_only[index] && !exposed[index];
500 if !eligible {
501 continue;
502 }
503 if let Inode::Str(text) = node {
504 *node = Inode::Secret(crate::ir::SecretBytes::from(text.clone()));
508 }
509 }
510 nodes
511}
512
513pub(crate) fn remap_node(node: &Inode, remap: &[NodeId]) -> Result<Inode> {
514 let map = |id: NodeId| -> Result<NodeId> {
515 remap
516 .get(id.index())
517 .copied()
518 .filter(|mapped| !mapped.is_none())
519 .ok_or(Error::EncodeDangling(id.0))
520 };
521 let opt = |id: Option<NodeId>| -> Result<Option<NodeId>> {
522 match id {
523 Some(id) => Ok(Some(map(id)?)),
524 None => Ok(None),
525 }
526 };
527 let out = match node {
528 Inode::Quantity { value, unit } => Inode::Quantity {
529 value: map(*value)?,
530 unit: unit.clone(),
531 },
532 Inode::Array { elements } => Inode::Array {
533 elements: elements.iter().map(|id| map(*id)).collect::<Result<_>>()?,
534 },
535 Inode::Table { fields } => Inode::Table {
536 fields: fields
537 .iter()
538 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
539 .collect::<Result<_>>()?,
540 },
541 Inode::SecretRef { locator } => Inode::SecretRef {
542 locator: map(*locator)?,
543 },
544 Inode::Ref(id) => Inode::Ref(map(*id)?),
545 Inode::Env { name, default } => Inode::Env {
546 name: name.clone(),
547 default: opt(*default)?,
548 },
549 Inode::Interpolate { parts } => Inode::Interpolate {
550 parts: parts
551 .iter()
552 .map(|part| match part {
553 InterpPart::Literal(text) => Ok(InterpPart::Literal(text.clone())),
554 InterpPart::Expr(id) => Ok(InterpPart::Expr(map(*id)?)),
555 })
556 .collect::<Result<_>>()?,
557 },
558 Inode::Binary { op, lhs, rhs } => Inode::Binary {
559 op: *op,
560 lhs: map(*lhs)?,
561 rhs: map(*rhs)?,
562 },
563 Inode::Unary { op, expr } => Inode::Unary {
564 op: *op,
565 expr: map(*expr)?,
566 },
567 Inode::Constrained { target, constraint } => Inode::Constrained {
568 target: map(*target)?,
569 constraint: constraint.clone(),
570 },
571 Inode::Patch(op) => Inode::Patch(match op {
572 PatchOp::Set { path, value } => PatchOp::Set {
573 path: path.clone(),
574 value: map(*value)?,
575 },
576 PatchOp::Remove { path } => PatchOp::Remove { path: path.clone() },
577 PatchOp::Rename { from, to } => PatchOp::Rename {
578 from: from.clone(),
579 to: to.clone(),
580 },
581 }),
582 Inode::SchemaField(spec) => Inode::SchemaField(SchemaFieldSpec {
583 typ: spec.typ.clone(),
584 min: opt(spec.min)?,
585 max: opt(spec.max)?,
586 min_len: opt(spec.min_len)?,
587 max_len: opt(spec.max_len)?,
588 pattern: opt(spec.pattern)?,
589 default: opt(spec.default)?,
590 deprecated: spec.deprecated.clone(),
591 }),
592 Inode::SchemaTable { fields, closed } => Inode::SchemaTable {
593 fields: fields
594 .iter()
595 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
596 .collect::<Result<_>>()?,
597 closed: *closed,
598 },
599 Inode::SchemaArray { element } => Inode::SchemaArray {
600 element: map(*element)?,
601 },
602 Inode::Semantic { tag, inner } => Inode::Semantic {
603 tag: *tag,
604 inner: map(*inner)?,
605 },
606 Inode::ContentAddressed { media_type, data } => Inode::ContentAddressed {
607 media_type: map(*media_type)?,
608 data: map(*data)?,
609 },
610 other => other.clone(),
611 };
612 Ok(out)
613}
614
615fn canonical_order(nodes: &[Inode], module: &InfinityModule) -> Result<Vec<NodeId>> {
617 let mut roots = Vec::new();
618 if let Some(root) = module.root() {
619 roots.push(root);
620 }
621 let mut bindings: Vec<NodeId> = module.bindings().iter().map(|(_, id)| *id).collect();
622 bindings.sort_by_key(|id| id.0);
623 roots.extend(bindings);
624 let mut subjects: Vec<NodeId> = module.annotations().iter().map(|a| a.node).collect();
625 subjects.sort_by_key(|id| id.0);
626 roots.extend(subjects);
627 let mut exprs: Vec<NodeId> = module
628 .annotations()
629 .iter()
630 .filter_map(|a| match &a.kind {
631 AnnotationKind::Constraint { expr, .. } => Some(*expr),
632 _ => None,
633 })
634 .collect();
635 exprs.sort_by_key(|id| id.0);
636 roots.extend(exprs);
637
638 let mut order = Vec::with_capacity(nodes.len());
639 let mut seen = vec![false; nodes.len()];
640 for root in roots {
641 if root.index() >= nodes.len() {
642 return Err(Error::EncodeDangling(root.0));
643 }
644 let mut stack = vec![root];
647 while let Some(id) = stack.pop() {
648 if seen[id.index()] {
649 continue;
650 }
651 seen[id.index()] = true;
652 order.push(id);
653 let children = nodes[id.index()].children();
654 for child in children.into_iter().rev() {
655 if child.index() >= nodes.len() {
656 return Err(Error::EncodeDangling(child.0));
657 }
658 if !seen[child.index()] {
659 stack.push(child);
660 }
661 }
662 }
663 }
664 Ok(order)
665}
666
667fn type_payload(
668 spec: &SchemaFieldSpec,
669 strings: &HashMap<String, u32>,
670) -> Result<(u32, Vec<(u32, u32)>)> {
671 let lookup = |text: &str| -> Result<u32> {
672 strings
673 .get(text)
674 .copied()
675 .ok_or(Error::Encode("schema type references an uninterned string"))
676 };
677 let mut aux = Vec::new();
678 match &spec.typ {
679 TypeSpec::Enum(items) => {
680 for item in items {
681 aux.push((crate::layout::AUX_TYPE_ENUM_ITEM, lookup(item)?));
682 }
683 }
684 TypeSpec::Table { closed } => {
685 aux.push((crate::layout::AUX_TYPE_TABLE, u32::from(*closed)));
686 }
687 _ => {}
688 }
689 Ok((spec.typ.code(), aux))
690}
691
692fn type_from_prepared(code: u32, aux: &[(u32, u32)], strings: &[String]) -> Result<TypeSpec> {
693 let string = |id: u32| -> Result<String> {
694 strings
695 .get(id as usize)
696 .cloned()
697 .ok_or(Error::StringRef(id))
698 };
699 let spec = match code {
700 0 => TypeSpec::Null,
701 1 => TypeSpec::Bool,
702 2 => TypeSpec::Int,
703 3 => TypeSpec::Float,
704 4 => TypeSpec::String,
705 5 => TypeSpec::Duration,
706 6 => TypeSpec::Bytes,
707 7 => {
708 let closed = aux
709 .iter()
710 .find(|(tag, _)| *tag == crate::layout::AUX_TYPE_TABLE)
711 .map(|(_, value)| *value != 0)
712 .unwrap_or(false);
713 TypeSpec::Table { closed }
714 }
715 8 => {
716 let mut items = Vec::new();
717 for (tag, value) in aux {
718 if *tag == crate::layout::AUX_TYPE_ENUM_ITEM {
719 items.push(string(*value)?);
720 }
721 }
722 TypeSpec::Enum(items)
723 }
724 9 => TypeSpec::Secret,
725 10 => TypeSpec::Any,
726 11 => TypeSpec::Blob,
727 _ => return Err(Error::Encode("unknown type discriminant")),
728 };
729 Ok(spec)
730}
731
732fn sorted_annotations(
733 mut annotations: Vec<Annotation>,
734 strings: &HashMap<String, u32>,
735) -> Vec<Annotation> {
736 annotations.sort_by(|a, b| {
737 let key = |annotation: &Annotation| -> (u32, u32, String, u32) {
738 match &annotation.kind {
739 AnnotationKind::Intent(text) => (annotation.node.0, 0, text.clone(), 0),
740 AnnotationKind::Constraint { source, expr } => {
741 (annotation.node.0, 1, source.clone(), expr.0)
742 }
743 AnnotationKind::AiImmutable(flag) => {
744 (annotation.node.0, 2, String::new(), u32::from(*flag))
745 }
746 }
747 };
748 key(a).cmp(&key(b))
749 });
750 annotations.dedup_by(|a, b| {
751 a.node == b.node
752 && match (&a.kind, &b.kind) {
753 (AnnotationKind::Intent(x), AnnotationKind::Intent(y)) => x == y,
754 (
755 AnnotationKind::Constraint {
756 source: s1,
757 expr: e1,
758 },
759 AnnotationKind::Constraint {
760 source: s2,
761 expr: e2,
762 },
763 ) => s1 == s2 && e1 == e2,
764 (AnnotationKind::AiImmutable(x), AnnotationKind::AiImmutable(y)) => x == y,
765 _ => false,
766 }
767 });
768 let _ = strings;
769 annotations
770}
771
772fn sorted_provenance(mut prov: Provenance) -> Result<Provenance> {
773 prov.entries.sort_by(|a, b| {
774 (a.key.as_str(), prov_kind_rank(&a.value)).cmp(&(b.key.as_str(), prov_kind_rank(&b.value)))
775 });
776 for window in prov.entries.windows(2) {
780 if window[0].key == window[1].key {
781 return Err(crate::error::Error::Encode(
782 "provenance keys collide after redaction; use distinct keys",
783 ));
784 }
785 }
786 Ok(prov)
787}
788fn prov_kind_rank(value: &ProvValue) -> u32 {
789 match value {
790 ProvValue::Text(_) => 0,
791 ProvValue::Int(_) => 1,
792 ProvValue::Flag(_) => 2,
793 ProvValue::Hash(_) => 3,
794 }
795}