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