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
394fn redact_text(text: &str, secrets: &[String]) -> String {
401 let mut out = text.to_string();
402 for secret in secrets {
403 if secret.is_empty() || !out.contains(secret.as_str()) {
404 continue;
405 }
406 out = out.replace(secret.as_str(), SECRET_MARKER);
407 }
408 out
409}
410
411fn redact_provenance(mut prov: Provenance, secrets: &[String]) -> Provenance {
412 if secrets.is_empty() {
413 return prov;
414 }
415 prov.producer = prov.producer.map(|text| redact_text(&text, secrets));
416 prov.model = prov.model.map(|text| redact_text(&text, secrets));
417 prov.tool = prov.tool.map(|text| redact_text(&text, secrets));
418 for entry in &mut prov.entries {
419 if let ProvValue::Text(text) = &entry.value {
420 entry.value = ProvValue::Text(redact_text(text, secrets));
421 }
422 }
423 prov
424}
425
426fn normalize_type(typ: &TypeSpec) -> TypeSpec {
427 typ.clone()
428}
429
430fn promote_secret_locators(mut nodes: Vec<Inode>, module: &InfinityModule) -> Vec<Inode> {
432 if nodes.is_empty() {
433 return nodes;
434 }
435 let mut exposed = vec![false; nodes.len()];
436 let mut mark = |id: NodeId| {
437 if let Some(slot) = exposed.get_mut(id.index()) {
438 *slot = true;
439 }
440 };
441 if let Some(root) = module.root() {
442 mark(root);
443 }
444 for (_, id) in module.bindings() {
445 mark(*id);
446 }
447 for annotation in module.annotations() {
448 mark(annotation.node);
449 if let AnnotationKind::Constraint { expr, .. } = &annotation.kind {
450 mark(*expr);
451 }
452 }
453 let mut secret_only = vec![true; nodes.len()];
456 let mut referenced = vec![false; nodes.len()];
457 for node in &nodes {
458 for child in node.children() {
459 let Some(slot) = secret_only.get_mut(child.index()) else {
460 continue;
461 };
462 referenced[child.index()] = true;
463 if !node.is_secret() {
464 *slot = false;
465 }
466 }
467 }
468 for (index, node) in nodes.iter_mut().enumerate() {
469 let eligible = referenced[index] && secret_only[index] && !exposed[index];
470 if !eligible {
471 continue;
472 }
473 if let Inode::Str(text) = node {
474 *node = Inode::Secret(crate::ir::SecretBytes::from(text.clone()));
478 }
479 }
480 nodes
481}
482
483pub(crate) fn remap_node(node: &Inode, remap: &[NodeId]) -> Result<Inode> {
484 let map = |id: NodeId| -> Result<NodeId> {
485 remap
486 .get(id.index())
487 .copied()
488 .filter(|mapped| !mapped.is_none())
489 .ok_or(Error::EncodeDangling(id.0))
490 };
491 let opt = |id: Option<NodeId>| -> Result<Option<NodeId>> {
492 match id {
493 Some(id) => Ok(Some(map(id)?)),
494 None => Ok(None),
495 }
496 };
497 let out = match node {
498 Inode::Quantity { value, unit } => Inode::Quantity {
499 value: map(*value)?,
500 unit: unit.clone(),
501 },
502 Inode::Array { elements } => Inode::Array {
503 elements: elements.iter().map(|id| map(*id)).collect::<Result<_>>()?,
504 },
505 Inode::Table { fields } => Inode::Table {
506 fields: fields
507 .iter()
508 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
509 .collect::<Result<_>>()?,
510 },
511 Inode::SecretRef { locator } => Inode::SecretRef {
512 locator: map(*locator)?,
513 },
514 Inode::Ref(id) => Inode::Ref(map(*id)?),
515 Inode::Env { name, default } => Inode::Env {
516 name: name.clone(),
517 default: opt(*default)?,
518 },
519 Inode::Interpolate { parts } => Inode::Interpolate {
520 parts: parts
521 .iter()
522 .map(|part| match part {
523 InterpPart::Literal(text) => Ok(InterpPart::Literal(text.clone())),
524 InterpPart::Expr(id) => Ok(InterpPart::Expr(map(*id)?)),
525 })
526 .collect::<Result<_>>()?,
527 },
528 Inode::Binary { op, lhs, rhs } => Inode::Binary {
529 op: *op,
530 lhs: map(*lhs)?,
531 rhs: map(*rhs)?,
532 },
533 Inode::Unary { op, expr } => Inode::Unary {
534 op: *op,
535 expr: map(*expr)?,
536 },
537 Inode::Constrained { target, constraint } => Inode::Constrained {
538 target: map(*target)?,
539 constraint: constraint.clone(),
540 },
541 Inode::Patch(op) => Inode::Patch(match op {
542 PatchOp::Set { path, value } => PatchOp::Set {
543 path: path.clone(),
544 value: map(*value)?,
545 },
546 PatchOp::Remove { path } => PatchOp::Remove { path: path.clone() },
547 PatchOp::Rename { from, to } => PatchOp::Rename {
548 from: from.clone(),
549 to: to.clone(),
550 },
551 }),
552 Inode::SchemaField(spec) => Inode::SchemaField(SchemaFieldSpec {
553 typ: spec.typ.clone(),
554 min: opt(spec.min)?,
555 max: opt(spec.max)?,
556 min_len: opt(spec.min_len)?,
557 max_len: opt(spec.max_len)?,
558 pattern: opt(spec.pattern)?,
559 default: opt(spec.default)?,
560 deprecated: spec.deprecated.clone(),
561 }),
562 Inode::SchemaTable { fields, closed } => Inode::SchemaTable {
563 fields: fields
564 .iter()
565 .map(|(name, id)| Ok((name.clone(), map(*id)?)))
566 .collect::<Result<_>>()?,
567 closed: *closed,
568 },
569 Inode::SchemaArray { element } => Inode::SchemaArray {
570 element: map(*element)?,
571 },
572 Inode::Semantic { tag, inner } => Inode::Semantic {
573 tag: *tag,
574 inner: map(*inner)?,
575 },
576 Inode::ContentAddressed { media_type, data } => Inode::ContentAddressed {
577 media_type: map(*media_type)?,
578 data: map(*data)?,
579 },
580 other => other.clone(),
581 };
582 Ok(out)
583}
584
585fn canonical_order(nodes: &[Inode], module: &InfinityModule) -> Result<Vec<NodeId>> {
587 let mut roots = Vec::new();
588 if let Some(root) = module.root() {
589 roots.push(root);
590 }
591 let mut bindings: Vec<NodeId> = module.bindings().iter().map(|(_, id)| *id).collect();
592 bindings.sort_by_key(|id| id.0);
593 roots.extend(bindings);
594 let mut subjects: Vec<NodeId> = module.annotations().iter().map(|a| a.node).collect();
595 subjects.sort_by_key(|id| id.0);
596 roots.extend(subjects);
597 let mut exprs: Vec<NodeId> = module
598 .annotations()
599 .iter()
600 .filter_map(|a| match &a.kind {
601 AnnotationKind::Constraint { expr, .. } => Some(*expr),
602 _ => None,
603 })
604 .collect();
605 exprs.sort_by_key(|id| id.0);
606 roots.extend(exprs);
607
608 let mut order = Vec::with_capacity(nodes.len());
609 let mut seen = vec![false; nodes.len()];
610 for root in roots {
611 if root.index() >= nodes.len() {
612 return Err(Error::EncodeDangling(root.0));
613 }
614 let mut stack = vec![root];
617 while let Some(id) = stack.pop() {
618 if seen[id.index()] {
619 continue;
620 }
621 seen[id.index()] = true;
622 order.push(id);
623 let children = nodes[id.index()].children();
624 for child in children.into_iter().rev() {
625 if child.index() >= nodes.len() {
626 return Err(Error::EncodeDangling(child.0));
627 }
628 if !seen[child.index()] {
629 stack.push(child);
630 }
631 }
632 }
633 }
634 Ok(order)
635}
636
637fn type_payload(
638 spec: &SchemaFieldSpec,
639 strings: &HashMap<String, u32>,
640) -> Result<(u32, Vec<(u32, u32)>)> {
641 let lookup = |text: &str| -> Result<u32> {
642 strings
643 .get(text)
644 .copied()
645 .ok_or(Error::Encode("schema type references an uninterned string"))
646 };
647 let mut aux = Vec::new();
648 match &spec.typ {
649 TypeSpec::Enum(items) => {
650 for item in items {
651 aux.push((crate::layout::AUX_TYPE_ENUM_ITEM, lookup(item)?));
652 }
653 }
654 TypeSpec::Table { closed } => {
655 aux.push((crate::layout::AUX_TYPE_TABLE, u32::from(*closed)));
656 }
657 _ => {}
658 }
659 Ok((spec.typ.code(), aux))
660}
661
662fn type_from_prepared(code: u32, aux: &[(u32, u32)], strings: &[String]) -> Result<TypeSpec> {
663 let string = |id: u32| -> Result<String> {
664 strings
665 .get(id as usize)
666 .cloned()
667 .ok_or(Error::StringRef(id))
668 };
669 let spec = match code {
670 0 => TypeSpec::Null,
671 1 => TypeSpec::Bool,
672 2 => TypeSpec::Int,
673 3 => TypeSpec::Float,
674 4 => TypeSpec::String,
675 5 => TypeSpec::Duration,
676 6 => TypeSpec::Bytes,
677 7 => {
678 let closed = aux
679 .iter()
680 .find(|(tag, _)| *tag == crate::layout::AUX_TYPE_TABLE)
681 .map(|(_, value)| *value != 0)
682 .unwrap_or(false);
683 TypeSpec::Table { closed }
684 }
685 8 => {
686 let mut items = Vec::new();
687 for (tag, value) in aux {
688 if *tag == crate::layout::AUX_TYPE_ENUM_ITEM {
689 items.push(string(*value)?);
690 }
691 }
692 TypeSpec::Enum(items)
693 }
694 9 => TypeSpec::Secret,
695 10 => TypeSpec::Any,
696 11 => TypeSpec::Blob,
697 _ => return Err(Error::Encode("unknown type discriminant")),
698 };
699 Ok(spec)
700}
701
702fn sorted_annotations(
703 mut annotations: Vec<Annotation>,
704 strings: &HashMap<String, u32>,
705) -> Vec<Annotation> {
706 annotations.sort_by(|a, b| {
707 let key = |annotation: &Annotation| -> (u32, u32, String, u32) {
708 match &annotation.kind {
709 AnnotationKind::Intent(text) => (annotation.node.0, 0, text.clone(), 0),
710 AnnotationKind::Constraint { source, expr } => {
711 (annotation.node.0, 1, source.clone(), expr.0)
712 }
713 AnnotationKind::AiImmutable(flag) => {
714 (annotation.node.0, 2, String::new(), u32::from(*flag))
715 }
716 }
717 };
718 key(a).cmp(&key(b))
719 });
720 annotations.dedup_by(|a, b| {
721 a.node == b.node
722 && match (&a.kind, &b.kind) {
723 (AnnotationKind::Intent(x), AnnotationKind::Intent(y)) => x == y,
724 (
725 AnnotationKind::Constraint {
726 source: s1,
727 expr: e1,
728 },
729 AnnotationKind::Constraint {
730 source: s2,
731 expr: e2,
732 },
733 ) => s1 == s2 && e1 == e2,
734 (AnnotationKind::AiImmutable(x), AnnotationKind::AiImmutable(y)) => x == y,
735 _ => false,
736 }
737 });
738 let _ = strings;
739 annotations
740}
741
742fn sorted_provenance(mut prov: Provenance, strings: &HashMap<String, u32>) -> Provenance {
743 let _ = strings;
744 prov.entries.sort_by(|a, b| {
745 (a.key.as_str(), prov_kind_rank(&a.value)).cmp(&(b.key.as_str(), prov_kind_rank(&b.value)))
746 });
747 prov.entries.dedup_by(|a, b| a == b);
748 prov
749}
750fn prov_kind_rank(value: &ProvValue) -> u32 {
751 match value {
752 ProvValue::Text(_) => 0,
753 ProvValue::Int(_) => 1,
754 ProvValue::Flag(_) => 2,
755 ProvValue::Hash(_) => 3,
756 }
757}