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 if let Some(prov) = &provenance {
183 if let Some(producer) = &prov.producer {
184 strings.push(producer.clone());
185 }
186 if let Some(model) = &prov.model {
187 strings.push(model.clone());
188 }
189 if let Some(tool) = &prov.tool {
190 strings.push(tool.clone());
191 }
192 for entry in &prov.entries {
193 strings.push(entry.key.clone());
194 if let ProvValue::Text(text) = &entry.value {
195 strings.push(text.clone());
196 }
197 }
198 }
199 strings.sort();
200 strings.dedup();
201 let string_index: HashMap<String, u32> = strings
202 .iter()
203 .enumerate()
204 .map(|(index, value)| (value.clone(), index as u32))
205 .collect();
206
207 let mut units: Vec<UnitTag> = Vec::new();
208 for node in &nodes {
209 if let Inode::Quantity { unit, .. } = node {
210 units.push(unit.clone());
211 }
212 }
213 units.sort_by(|a, b| (a.discriminant(), a.name()).cmp(&(b.discriminant(), b.name())));
214 units.dedup();
215 let unit_index: HashMap<UnitTag, u32> = units
216 .iter()
217 .enumerate()
218 .map(|(index, unit)| (unit.clone(), index as u32))
219 .collect();
220
221 let mut type_payloads: Vec<(u32, Vec<(u32, u32)>)> = Vec::new();
223 for node in &nodes {
224 if let Inode::SchemaField(spec) = node {
225 type_payloads.push(type_payload(spec, &string_index)?);
226 }
227 }
228 type_payloads.sort();
229 type_payloads.dedup();
230 let mut types = Vec::with_capacity(type_payloads.len());
231 let mut type_index = HashMap::new();
232 for (index, (code, aux)) in type_payloads.iter().enumerate() {
233 let spec = type_from_prepared(*code, aux, &strings)?;
234 types.push(PreparedType {
235 spec,
236 aux: aux.clone(),
237 });
238 type_index.insert((*code, aux.clone()), index as u32);
239 }
240
241 let mut blobs: Vec<BlobPayload> = Vec::new();
243 for node in &nodes {
244 if let Inode::Blob {
245 hash,
246 size,
247 tainted,
248 payload,
249 ..
250 } = node
251 {
252 if *tainted && payload.is_some() {
253 return Err(Error::Encode("a tainted blob must not carry a payload"));
254 }
255 blobs.push(BlobPayload {
256 hash: *hash,
257 size: *size,
258 payload: if *tainted { None } else { payload.clone() },
259 });
260 }
261 }
262 blobs.sort_by_key(|blob| (blob.hash, blob.size));
263 blobs.dedup();
264 let blob_index: HashMap<BlobKey, u32> = blobs
265 .iter()
266 .enumerate()
267 .map(|(index, blob)| ((blob.hash, blob.size, blob.payload.clone()), index as u32))
268 .collect();
269
270 let mut hashes: Vec<[u8; 32]> = Vec::new();
272 if let Some(prov) = &provenance {
273 if let Some(parent) = prov.parent {
274 hashes.push(parent);
275 }
276 for entry in &prov.entries {
277 if let ProvValue::Hash(hash) = &entry.value {
278 hashes.push(*hash);
279 }
280 }
281 }
282 hashes.sort();
283 hashes.dedup();
284 let hash_index: HashMap<[u8; 32], u32> = hashes
285 .iter()
286 .enumerate()
287 .map(|(index, hash)| (*hash, index as u32))
288 .collect();
289
290 let mut numbers: Vec<i64> = Vec::new();
292 for node in &nodes {
293 if let Inode::Constrained {
294 constraint: Constraint::MinMax { min, max },
295 ..
296 } = node
297 {
298 numbers.extend(min.iter().copied());
299 numbers.extend(max.iter().copied());
300 }
301 }
302 if let Some(prov) = &provenance {
303 for entry in &prov.entries {
304 if let ProvValue::Int(value) = &entry.value {
305 numbers.push(*value);
306 }
307 }
308 }
309 numbers.sort_unstable();
310 numbers.dedup();
311 let number_index: HashMap<i64, u32> = numbers
312 .iter()
313 .enumerate()
314 .map(|(index, value)| (*value, index as u32))
315 .collect();
316
317 let secrets: Vec<Vec<u8>> = nodes
319 .iter()
320 .filter_map(|node| match node {
321 Inode::Secret(secret) => Some(secret.expose().to_vec()),
322 _ => None,
323 })
324 .collect();
325
326 let mut extensions: Vec<(u64, Vec<u8>)> = module.extensions().to_vec();
328 extensions.sort_by_key(|(tag, _)| *tag);
329 for window in extensions.windows(2) {
330 if window[0].0 == window[1].0 {
331 return Err(Error::Encode("duplicate extension tag"));
332 }
333 }
334
335 let prepared = Prepared {
336 nodes,
337 root,
338 bindings,
339 annotations: sorted_annotations(annotations, &string_index),
340 provenance: provenance.map(|prov| sorted_provenance(prov, &string_index)),
341 strings,
342 units,
343 types,
344 blobs,
345 hashes,
346 numbers,
347 secrets,
348 extensions,
349 string_index,
350 unit_index,
351 type_index,
352 blob_index,
353 hash_index,
354 number_index,
355 };
356 Ok(prepared)
357}
358
359fn normalize(node: &Inode) -> Inode {
360 match node {
361 Inode::Float(value) => Inode::Float(canonical_f64(*value)),
362 Inode::Table { fields } => {
363 let mut fields = fields.clone();
364 fields.sort_by(|a, b| a.0.cmp(&b.0));
365 Inode::Table { fields }
366 }
367 Inode::SchemaTable { fields, closed } => {
368 let mut fields = fields.clone();
369 fields.sort_by(|a, b| a.0.cmp(&b.0));
370 Inode::SchemaTable {
371 fields,
372 closed: *closed,
373 }
374 }
375 Inode::SchemaField(spec) => Inode::SchemaField(SchemaFieldSpec {
376 typ: normalize_type(&spec.typ),
377 ..spec.clone()
378 }),
379 Inode::Constrained {
380 target,
381 constraint: Constraint::Enum(items),
382 } => Inode::Constrained {
383 target: *target,
384 constraint: Constraint::Enum(items.clone()),
385 },
386 other => other.clone(),
387 }
388}
389
390pub const SECRET_MARKER: &str = "***SECRET***";
393
394pub(crate) fn redact_text(text: &str, secrets: &[String]) -> String {
407 let matched: Vec<&String> = secrets
408 .iter()
409 .filter(|secret| !secret.is_empty() && text.contains(secret.as_str()))
410 .collect();
411 if matched.is_empty() {
412 return text.to_string();
413 }
414 if secrets
415 .iter()
416 .any(|secret| !secret.is_empty() && SECRET_MARKER.contains(secret.as_str()))
417 {
418 return SECRET_MARKER.to_string();
419 }
420 let mut out = text.to_string();
421 for secret in matched {
422 out = out.replace(secret.as_str(), SECRET_MARKER);
423 }
424 out
425}
426
427fn redact_provenance(mut prov: Provenance, secrets: &[String]) -> Provenance {
428 if secrets.is_empty() {
429 return prov;
430 }
431 prov.producer = prov.producer.map(|text| redact_text(&text, secrets));
432 prov.model = prov.model.map(|text| redact_text(&text, secrets));
433 prov.tool = prov.tool.map(|text| redact_text(&text, secrets));
434 for entry in &mut prov.entries {
435 entry.key = redact_text(&entry.key, secrets);
438 if let ProvValue::Text(text) = &entry.value {
439 entry.value = ProvValue::Text(redact_text(text, secrets));
440 }
441 }
442 prov
443}
444
445fn normalize_type(typ: &TypeSpec) -> TypeSpec {
446 typ.clone()
447}
448
449fn promote_secret_locators(mut nodes: Vec<Inode>, module: &InfinityModule) -> Vec<Inode> {
451 if nodes.is_empty() {
452 return nodes;
453 }
454 let mut exposed = vec![false; nodes.len()];
455 let mut mark = |id: NodeId| {
456 if let Some(slot) = exposed.get_mut(id.index()) {
457 *slot = true;
458 }
459 };
460 if let Some(root) = module.root() {
461 mark(root);
462 }
463 for (_, id) in module.bindings() {
464 mark(*id);
465 }
466 for annotation in module.annotations() {
467 mark(annotation.node);
468 if let AnnotationKind::Constraint { expr, .. } = &annotation.kind {
469 mark(*expr);
470 }
471 }
472 let mut secret_only = vec![true; nodes.len()];
475 let mut referenced = vec![false; nodes.len()];
476 for node in &nodes {
477 for child in node.children() {
478 let Some(slot) = secret_only.get_mut(child.index()) else {
479 continue;
480 };
481 referenced[child.index()] = true;
482 if !node.is_secret() {
483 *slot = false;
484 }
485 }
486 }
487 for (index, node) in nodes.iter_mut().enumerate() {
488 let eligible = referenced[index] && secret_only[index] && !exposed[index];
489 if !eligible {
490 continue;
491 }
492 if let Inode::Str(text) = node {
493 *node = Inode::Secret(crate::ir::SecretBytes::from(text.clone()));
497 }
498 }
499 nodes
500}
501
502pub(crate) fn remap_node(node: &Inode, remap: &[NodeId]) -> Result<Inode> {
503 let map = |id: NodeId| -> Result<NodeId> {
504 remap
505 .get(id.index())
506 .copied()
507 .filter(|mapped| !mapped.is_none())
508 .ok_or(Error::EncodeDangling(id.0))
509 };
510 let opt = |id: Option<NodeId>| -> Result<Option<NodeId>> {
511 match id {
512 Some(id) => Ok(Some(map(id)?)),
513 None => Ok(None),
514 }
515 };
516 let out = match node {
517 Inode::Quantity { value, unit } => Inode::Quantity {
518 value: map(*value)?,
519 unit: unit.clone(),
520 },
521 Inode::Array { elements } => Inode::Array {
522 elements: elements.iter().map(|id| map(*id)).collect::<Result<_>>()?,
523 },
524 Inode::Table { fields } => Inode::Table {
525 fields: fields
526 .iter()
527 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
528 .collect::<Result<_>>()?,
529 },
530 Inode::SecretRef { locator } => Inode::SecretRef {
531 locator: map(*locator)?,
532 },
533 Inode::Ref(id) => Inode::Ref(map(*id)?),
534 Inode::Env { name, default } => Inode::Env {
535 name: name.clone(),
536 default: opt(*default)?,
537 },
538 Inode::Interpolate { parts } => Inode::Interpolate {
539 parts: parts
540 .iter()
541 .map(|part| match part {
542 InterpPart::Literal(text) => Ok(InterpPart::Literal(text.clone())),
543 InterpPart::Expr(id) => Ok(InterpPart::Expr(map(*id)?)),
544 })
545 .collect::<Result<_>>()?,
546 },
547 Inode::Binary { op, lhs, rhs } => Inode::Binary {
548 op: *op,
549 lhs: map(*lhs)?,
550 rhs: map(*rhs)?,
551 },
552 Inode::Unary { op, expr } => Inode::Unary {
553 op: *op,
554 expr: map(*expr)?,
555 },
556 Inode::Constrained { target, constraint } => Inode::Constrained {
557 target: map(*target)?,
558 constraint: constraint.clone(),
559 },
560 Inode::Patch(op) => Inode::Patch(match op {
561 PatchOp::Set { path, value } => PatchOp::Set {
562 path: path.clone(),
563 value: map(*value)?,
564 },
565 PatchOp::Remove { path } => PatchOp::Remove { path: path.clone() },
566 PatchOp::Rename { from, to } => PatchOp::Rename {
567 from: from.clone(),
568 to: to.clone(),
569 },
570 }),
571 Inode::SchemaField(spec) => Inode::SchemaField(SchemaFieldSpec {
572 typ: spec.typ.clone(),
573 min: opt(spec.min)?,
574 max: opt(spec.max)?,
575 min_len: opt(spec.min_len)?,
576 max_len: opt(spec.max_len)?,
577 pattern: opt(spec.pattern)?,
578 default: opt(spec.default)?,
579 deprecated: spec.deprecated.clone(),
580 }),
581 Inode::SchemaTable { fields, closed } => Inode::SchemaTable {
582 fields: fields
583 .iter()
584 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
585 .collect::<Result<_>>()?,
586 closed: *closed,
587 },
588 Inode::SchemaArray { element } => Inode::SchemaArray {
589 element: map(*element)?,
590 },
591 Inode::Semantic { tag, inner } => Inode::Semantic {
592 tag: *tag,
593 inner: map(*inner)?,
594 },
595 Inode::ContentAddressed { media_type, data } => Inode::ContentAddressed {
596 media_type: map(*media_type)?,
597 data: map(*data)?,
598 },
599 other => other.clone(),
600 };
601 Ok(out)
602}
603
604fn canonical_order(nodes: &[Inode], module: &InfinityModule) -> Result<Vec<NodeId>> {
606 let mut roots = Vec::new();
607 if let Some(root) = module.root() {
608 roots.push(root);
609 }
610 let mut bindings: Vec<NodeId> = module.bindings().iter().map(|(_, id)| *id).collect();
611 bindings.sort_by_key(|id| id.0);
612 roots.extend(bindings);
613 let mut subjects: Vec<NodeId> = module.annotations().iter().map(|a| a.node).collect();
614 subjects.sort_by_key(|id| id.0);
615 roots.extend(subjects);
616 let mut exprs: Vec<NodeId> = module
617 .annotations()
618 .iter()
619 .filter_map(|a| match &a.kind {
620 AnnotationKind::Constraint { expr, .. } => Some(*expr),
621 _ => None,
622 })
623 .collect();
624 exprs.sort_by_key(|id| id.0);
625 roots.extend(exprs);
626
627 let mut order = Vec::with_capacity(nodes.len());
628 let mut seen = vec![false; nodes.len()];
629 for root in roots {
630 if root.index() >= nodes.len() {
631 return Err(Error::EncodeDangling(root.0));
632 }
633 let mut stack = vec![root];
636 while let Some(id) = stack.pop() {
637 if seen[id.index()] {
638 continue;
639 }
640 seen[id.index()] = true;
641 order.push(id);
642 let children = nodes[id.index()].children();
643 for child in children.into_iter().rev() {
644 if child.index() >= nodes.len() {
645 return Err(Error::EncodeDangling(child.0));
646 }
647 if !seen[child.index()] {
648 stack.push(child);
649 }
650 }
651 }
652 }
653 Ok(order)
654}
655
656fn type_payload(
657 spec: &SchemaFieldSpec,
658 strings: &HashMap<String, u32>,
659) -> Result<(u32, Vec<(u32, u32)>)> {
660 let lookup = |text: &str| -> Result<u32> {
661 strings
662 .get(text)
663 .copied()
664 .ok_or(Error::Encode("schema type references an uninterned string"))
665 };
666 let mut aux = Vec::new();
667 match &spec.typ {
668 TypeSpec::Enum(items) => {
669 for item in items {
670 aux.push((crate::layout::AUX_TYPE_ENUM_ITEM, lookup(item)?));
671 }
672 }
673 TypeSpec::Table { closed } => {
674 aux.push((crate::layout::AUX_TYPE_TABLE, u32::from(*closed)));
675 }
676 _ => {}
677 }
678 Ok((spec.typ.code(), aux))
679}
680
681fn type_from_prepared(code: u32, aux: &[(u32, u32)], strings: &[String]) -> Result<TypeSpec> {
682 let string = |id: u32| -> Result<String> {
683 strings
684 .get(id as usize)
685 .cloned()
686 .ok_or(Error::StringRef(id))
687 };
688 let spec = match code {
689 0 => TypeSpec::Null,
690 1 => TypeSpec::Bool,
691 2 => TypeSpec::Int,
692 3 => TypeSpec::Float,
693 4 => TypeSpec::String,
694 5 => TypeSpec::Duration,
695 6 => TypeSpec::Bytes,
696 7 => {
697 let closed = aux
698 .iter()
699 .find(|(tag, _)| *tag == crate::layout::AUX_TYPE_TABLE)
700 .map(|(_, value)| *value != 0)
701 .unwrap_or(false);
702 TypeSpec::Table { closed }
703 }
704 8 => {
705 let mut items = Vec::new();
706 for (tag, value) in aux {
707 if *tag == crate::layout::AUX_TYPE_ENUM_ITEM {
708 items.push(string(*value)?);
709 }
710 }
711 TypeSpec::Enum(items)
712 }
713 9 => TypeSpec::Secret,
714 10 => TypeSpec::Any,
715 11 => TypeSpec::Blob,
716 _ => return Err(Error::Encode("unknown type discriminant")),
717 };
718 Ok(spec)
719}
720
721fn sorted_annotations(
722 mut annotations: Vec<Annotation>,
723 strings: &HashMap<String, u32>,
724) -> Vec<Annotation> {
725 annotations.sort_by(|a, b| {
726 let key = |annotation: &Annotation| -> (u32, u32, String, u32) {
727 match &annotation.kind {
728 AnnotationKind::Intent(text) => (annotation.node.0, 0, text.clone(), 0),
729 AnnotationKind::Constraint { source, expr } => {
730 (annotation.node.0, 1, source.clone(), expr.0)
731 }
732 AnnotationKind::AiImmutable(flag) => {
733 (annotation.node.0, 2, String::new(), u32::from(*flag))
734 }
735 }
736 };
737 key(a).cmp(&key(b))
738 });
739 annotations.dedup_by(|a, b| {
740 a.node == b.node
741 && match (&a.kind, &b.kind) {
742 (AnnotationKind::Intent(x), AnnotationKind::Intent(y)) => x == y,
743 (
744 AnnotationKind::Constraint {
745 source: s1,
746 expr: e1,
747 },
748 AnnotationKind::Constraint {
749 source: s2,
750 expr: e2,
751 },
752 ) => s1 == s2 && e1 == e2,
753 (AnnotationKind::AiImmutable(x), AnnotationKind::AiImmutable(y)) => x == y,
754 _ => false,
755 }
756 });
757 let _ = strings;
758 annotations
759}
760
761fn sorted_provenance(mut prov: Provenance, strings: &HashMap<String, u32>) -> Provenance {
762 let _ = strings;
763 prov.entries.sort_by(|a, b| {
764 (a.key.as_str(), prov_kind_rank(&a.value)).cmp(&(b.key.as_str(), prov_kind_rank(&b.value)))
765 });
766 prov.entries.dedup_by(|a, b| a.key == b.key);
770 prov
771}
772fn prov_kind_rank(value: &ProvValue) -> u32 {
773 match value {
774 ProvValue::Text(_) => 0,
775 ProvValue::Int(_) => 1,
776 ProvValue::Flag(_) => 2,
777 ProvValue::Hash(_) => 3,
778 }
779}