Phased and Confused: typestate enforcement for post-mono Body invariants

src/mk_graph/output/d2.rs

@@ -23,7 +23,7 @@ impl SmirJson {
         for item in self.items {
             match item.mono_item_kind {
                 MonoItemKind::MonoItemFn { name, body, .. } => {
-                    render_d2_function(&name, body.as_ref(), &ctx, &mut output);
+                    render_d2_function(&name, body.as_ref().map(|b| b.inner()), &ctx, &mut output);
                 }
                 MonoItemKind::MonoItemGlobalAsm { asm } => {
                     render_d2_asm(&asm, &mut output);

src/mk_graph/output/dot.rs

@@ -68,6 +68,7 @@ impl SmirJson {
             for item in self.items {
                 match item.mono_item_kind {
                     MonoItemKind::MonoItemFn { name, body, id: _ } => {
+                        let body = body.as_ref().map(|b| b.inner());
                         let mut c = graph.cluster();
                         c.set_label(&name_lines(&name));
                         c.set_style(Style::Filled);
@@ -81,7 +82,7 @@ impl SmirJson {
                         let mut local_node = c.node_auto();
                         let mut vector: Vec<String> = vec![];
                         vector.push(String::from("LOCALS"));
-                        for (index, decl) in body.clone().unwrap().local_decls() {
+                        for (index, decl) in body.cloned().unwrap().local_decls() {
                             let ty_with_layout = ctx.render_type_with_layout(decl.ty);
                             vector.push(format!("{index} = {}", ty_with_layout));
                         }
@@ -287,7 +288,7 @@ impl SmirJson {
                                 }
                             };
 
-                        if let Some(ref body) = body {
+                        if let Some(body) = body {
                             add_call_edges(&mut graph, 0, &body.blocks);
                         }
                     }

src/printer/collect.rs

@@ -24,6 +24,7 @@ use stable_mir::CrateDef;
 
 use super::items::{get_foreign_module_details, mk_item};
 use super::mir_visitor::{maybe_add_to_link_map, BodyAnalyzer, UnevalConstInfo};
+use super::phased::Raw;
 use super::schema::{
     AllocInfo, AllocMap, CollectedCrate, DerivedInfo, Item, LinkMap, SmirJson, SmirJsonDebugInfo,
     SpanMap,
@@ -53,7 +54,7 @@ fn warn_missing_body(mono_item: &MonoItem) {
     }
 }
 
-fn collect_items(tcx: TyCtxt<'_>) -> HashMap<String, (MonoItem, Item)> {
+fn collect_items(tcx: TyCtxt<'_>) -> HashMap<String, (MonoItem, Item<Raw>)> {
     // get initial set of mono_items
     let items = mono_collect(tcx);
     items
@@ -73,7 +74,7 @@ fn enqueue_unevaluated_consts(
     tcx: TyCtxt<'_>,
     discovered: Vec<UnevalConstInfo>,
     known_names: &mut HashSet<String>,
-    pending: &mut HashMap<String, (MonoItem, Item)>,
+    pending: &mut HashMap<String, (MonoItem, Item<Raw>)>,
     unevaluated_consts: &mut HashMap<stable_mir::ty::ConstDef, String>,
 ) {
     for info in discovered {
@@ -98,7 +99,7 @@ fn enqueue_unevaluated_consts(
 /// this phase, then dropped; only the `Item` survives into `CollectedCrate`.
 fn collect_and_analyze_items(
     tcx: TyCtxt<'_>,
-    initial_items: HashMap<String, (MonoItem, Item)>,
+    initial_items: HashMap<String, (MonoItem, Item<Raw>)>,
 ) -> (CollectedCrate, DerivedInfo) {
     let mut calls_map: LinkMap = HashMap::new();
     let mut visited_allocs = AllocMap::new();
@@ -107,8 +108,8 @@ fn collect_and_analyze_items(
     let mut unevaluated_consts: HashMap<stable_mir::ty::ConstDef, String> = HashMap::new();
 
     let mut known_names: HashSet<String> = initial_items.keys().cloned().collect();
-    let mut pending: HashMap<String, (MonoItem, Item)> = initial_items;
-    let mut all_items: Vec<Item> = Vec::new();
+    let mut pending: HashMap<String, (MonoItem, Item<Raw>)> = initial_items;
+    let mut all_items: Vec<Item<Raw>> = Vec::new();
 
     while let Some((_name, (mono_item, item))) = take_any(&mut pending) {
         maybe_add_to_link_map(tcx, &mono_item, &mut calls_map);
@@ -142,9 +143,13 @@ fn collect_and_analyze_items(
         all_items.push(item);
     }
 
+    // Phase 2 -> 3 boundary: validate that all bodies are free of
+    // ConstantKind::Unevaluated before handing them to phase 3.
+    let validated_items: Vec<_> = all_items.into_iter().map(|i| i.validate()).collect();
+
     (
         CollectedCrate {
-            items: all_items,
+            items: validated_items,
             unevaluated_consts,
         },
         DerivedInfo {

src/printer/items.rs

@@ -22,21 +22,23 @@ use stable_mir::{CrateDef, CrateItem};
 
 use crate::compat::bridge::mono_instance;
 
+use super::phased::{Phased, Raw};
 use super::schema::{BodyDetails, ForeignItem, ForeignModule, GenericData, Item, ItemDetails};
 
 #[derive(Serialize, Clone)]
-pub enum MonoItemKind {
+#[serde(bound = "")]
+pub enum MonoItemKind<S> {
     MonoItemFn {
         name: String,
         id: stable_mir::DefId,
-        body: Option<Body>,
+        body: Option<Phased<Body, S>>,
     },
     MonoItemStatic {
         name: String,
         id: stable_mir::DefId,
         allocation: Option<Allocation>,
         #[serde(skip)]
-        body: Option<Body>,
+        body: Option<Phased<Body, S>>,
     },
     MonoItemGlobalAsm {
         asm: String,
@@ -73,7 +75,7 @@ fn get_item_details(
     }
 }
 
-pub(super) fn mk_item(tcx: TyCtxt<'_>, item: MonoItem, sym_name: String) -> (MonoItem, Item) {
+pub(super) fn mk_item(tcx: TyCtxt<'_>, item: MonoItem, sym_name: String) -> (MonoItem, Item<Raw>) {
     match item {
         MonoItem::Fn(inst) => {
             let id = inst.def.def_id();
@@ -89,7 +91,7 @@ pub(super) fn mk_item(tcx: TyCtxt<'_>, item: MonoItem, sym_name: String) -> (Mon
                     MonoItemKind::MonoItemFn {
                         name: name.clone(),
                         id,
-                        body,
+                        body: body.map(Phased::new),
                     },
                     details,
                 ),
@@ -118,7 +120,7 @@ pub(super) fn mk_item(tcx: TyCtxt<'_>, item: MonoItem, sym_name: String) -> (Mon
                         name: static_def.name(),
                         id: static_def.def_id(),
                         allocation: alloc,
-                        body,
+                        body: body.map(Phased::new),
                     },
                     get_item_details(tcx, internal_id, None, None),
                 ),

src/printer/mod.rs

@@ -9,8 +9,9 @@
 //! | Module | Responsibility |
 //! |--------|----------------|
 //! | [`schema`] | Data model types ([`SmirJson`], [`Item`], [`AllocInfo`], etc.) and type aliases; [`Item`] deliberately excludes `MonoItem` for structural phase separation |
-//! | [`collect`] | Three-phase pipeline: collect items, analyze bodies, assemble final output; phase boundary is enforced structurally via the `(MonoItem, Item)` split |
-//! | [`items`] | Constructing `(MonoItem, Item)` pairs and extracting debug-level details |
+//! | [`collect`] | Three-phase pipeline: collect items, analyze bodies, assemble final output; phase boundary is enforced structurally via the `(MonoItem, Item)` split and [`phased::Phased`] body wrapper |
+//! | [`items`] | Constructing `(MonoItem, Item<Raw>)` pairs and extracting debug-level details |
+//! | [`phased`] | [`Phased<T, S>`](phased::Phased) typestate wrapper encoding the post-monomorphization body invariant |
 //! | [`mir_visitor`] | `BodyAnalyzer`: single-pass MIR body traversal collecting calls, allocs, types, spans |
 //! | [`ty_visitor`] | `TyCollector`: recursively collects reachable types with layout info (some special kinds are traversed but not stored) |
 //! | [`link_map`] | Function resolution map: type + instance kind to symbol name |
@@ -49,6 +50,7 @@ mod collect;
 mod items;
 mod link_map;
 mod mir_visitor;
+mod phased;
 mod schema;
 mod ty_visitor;
 mod types;
@@ -57,6 +59,7 @@ mod util;
 // Re-exports preserving the public API
 pub use collect::collect_smir;
 pub use items::MonoItemKind;
+pub use phased::Mono;
 pub use schema::{AllocInfo, FnSymType, Item, LinkMapKey, SmirJson, TypeMetadata};
 pub(crate) use util::hash;
 

src/printer/phased.rs

@@ -0,0 +1,103 @@
+//! Typestate wrapper encoding the post-monomorphization body invariant.
+//!
+//! After the collection pipeline's fixpoint loop resolves all unevaluated
+//! constants, bodies stored in [`Item`](super::schema::Item) should contain
+//! no `ConstantKind::Unevaluated`. [`Phased<T, S>`] encodes this guarantee
+//! in the type system: phase 1+2 code works with `Phased<Body, Raw>`, and
+//! the phase 2 -> 3 boundary validates each body (via [`UnevaluatedChecker`])
+//! before producing `Phased<Body, Mono>`.
+//!
+//! [`Phased`] serializes transparently (delegates to the inner `T`), so
+//! JSON output is unchanged.
+
+use std::marker::PhantomData;
+
+use crate::compat::serde;
+use crate::compat::stable_mir;
+
+use serde::{Serialize, Serializer};
+use stable_mir::mir::visit::MirVisitor;
+use stable_mir::mir::Body;
+
+/// Pre-validation phase marker. Bodies may still contain `ConstantKind::Unevaluated`.
+pub struct Raw;
+
+/// Post-validation phase marker. Bodies are guaranteed free of `ConstantKind::Unevaluated`.
+pub struct Mono;
+
+/// A value of type `T` tagged with a phase marker `S`.
+///
+/// Construction of `Phased<T, Raw>` is unrestricted; `Phased<T, Mono>` can
+/// only be obtained through [`Phased::<Body, Raw>::validate`], which walks
+/// the body and panics if any unevaluated constant is found.
+pub struct Phased<T, S>(T, PhantomData<S>);
+
+impl<T, S> Phased<T, S> {
+    /// Borrow the inner value.
+    pub fn inner(&self) -> &T {
+        &self.0
+    }
+
+    /// Consume the wrapper, returning the inner value.
+    pub fn into_inner(self) -> T {
+        self.0
+    }
+}
+
+impl<T: Clone, S> Clone for Phased<T, S> {
+    fn clone(&self) -> Self {
+        Phased(self.0.clone(), PhantomData)
+    }
+}
+
+impl<T: Serialize, S> Serialize for Phased<T, S> {
+    fn serialize<Ser>(&self, serializer: Ser) -> Result<Ser::Ok, Ser::Error>
+    where
+        Ser: Serializer,
+    {
+        self.0.serialize(serializer)
+    }
+}
+
+impl Phased<Body, Raw> {
+    /// Wrap a body in the `Raw` phase. Unrestricted.
+    pub fn new(body: Body) -> Self {
+        Phased(body, PhantomData)
+    }
+
+    /// Validate that the body contains no `ConstantKind::Unevaluated`, then
+    /// promote to the `Mono` phase. Panics if the invariant is violated.
+    pub fn validate(self) -> Phased<Body, Mono> {
+        UnevaluatedChecker.visit_body(&self.0);
+        Phased(self.0, PhantomData)
+    }
+}
+
+/// MIR visitor that panics if any `ConstantKind::Unevaluated` is encountered.
+///
+/// Used by [`Phased::<Body, Raw>::validate`] to enforce the post-mono invariant
+/// at the phase 2 -> 3 boundary.
+struct UnevaluatedChecker;
+
+impl MirVisitor for UnevaluatedChecker {
+    fn visit_mir_const(
+        &mut self,
+        constant: &stable_mir::ty::MirConst,
+        loc: stable_mir::mir::visit::Location,
+    ) {
+        if let stable_mir::ty::ConstantKind::Unevaluated(_) = constant.kind() {
+            panic!(
+                "Phased<Body, Raw>::validate: body contains ConstantKind::Unevaluated \
+                 at location {loc:?}.\n\n\
+                 The fixpoint loop in collect_and_analyze_items should have resolved \
+                 all unevaluated constants before phase 3, but this one survived. \
+                 The input program likely triggers a code path we haven't encountered yet.\n\n\
+                 Action: please file an issue at \
+                 https://github.com/runtimeverification/stable-mir-json/issues \
+                 with the .rs file that triggered this panic so we can add it as a \
+                 test case and extend the fixpoint loop to handle it."
+            );
+        }
+        self.super_mir_const(constant, loc);
+    }
+}