refactor(lowering): MutationTracker (body vs signature) and drop AnnotationMap::into_any_box

Two related cleanups on `AggregateTypeLowerer`:

- The `AnnotationMap::into_any_box` trait method existed only to
  round-trip a `Box<dyn AnnotationMap>` back to a concrete
  `AstAnnotations`. The field is now `Option<AstAnnotations>`
  directly; the trait method and its two impls are gone.
- The hand-incremented `mutations: usize` counter is replaced by a
  `MutationTracker` that distinguishes body mutations (statement
  inserts) from signature rewrites. The visitor's mutation primitives
  (`push_pre_statement`, `push_post_statement`, signature-rewrite in
  `visit_pou`) call `touch()` / `signature_changed(pou_name)`
  themselves. The call site no longer has to remember to bump
  anything, and the per-POU recording means
  `AggregateTypeLowerer::post_annotate` only invalidates callers of
  POUs whose interface actually changed — not every POU in the
  touched unit.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: ghaith

compiler/plc_driver/src/pipelines/participant.rs

@@ -18,7 +18,6 @@ use plc::{
     index::{Index, PouIndexEntry},
     lowering::{calls::AggregateTypeLowerer, polymorphism::PolymorphismLowerer},
     output::FormatOption,
-    resolver::AstAnnotations,
     ConfigFormat, OnlineChange, Target,
 };
 use plc_diagnostics::diagnostics::Diagnostic;
@@ -365,18 +364,24 @@ impl PipelineParticipantMut for AggregateTypeLowerer {
 
         let AnnotatedProject { mut units, index, annotations, diagnostics } = annotated_project;
         self.index = Some(index);
-        self.annotation = Some(Box::new(annotations));
+        self.annotation = Some(annotations);
 
         // Walk each unit; bookkeeper records which were mutated and which
         // POU names had their signature rewritten (aggregate return →
         // VAR_IN_OUT). Constants enter the index because the new
         // VAR_IN_OUT parameter's type can be e.g. `STRING[K+1]`.
+        //
+        // `signature_changed_pous()` returns only POUs whose interface was
+        // actually rewritten — body-only mutations (e.g. inserted alloca
+        // prelude for an aggregate-returning call site) flip the dirty
+        // flag but must not invalidate callers of every other POU in the
+        // unit.
         let mut book = super::bookkeeping::LoweringBookkeeper::new();
         for (idx, annotated_unit) in units.iter_mut().enumerate() {
             if self.visit_unit_tracked(&mut annotated_unit.unit) {
                 book.mark_unit(idx);
-                for pou in &annotated_unit.unit.pous {
-                    book.signature_changed(&pou.name);
+                for name in self.signature_changed_pous() {
+                    book.signature_changed(&name);
                 }
             }
         }
@@ -388,10 +393,6 @@ impl PipelineParticipantMut for AggregateTypeLowerer {
         // borrowed them for the walk).
         let index = self.index.take().expect("index returned by visit");
         let annotations = self.annotation.take().expect("annotations returned by visit");
-        let annotations = *annotations
-            .into_any_box()
-            .downcast::<AstAnnotations>()
-            .expect("post_annotate participant always receives AstAnnotations");
 
         let project = AnnotatedProject { units, index, annotations, diagnostics };
         book.apply_to_annotated(project, &reverse_deps, self.id_provider.clone())

src/lowering.rs

@@ -1,4 +1,5 @@
 pub mod calls;
 pub mod helper;
+pub mod mutation;
 pub mod polymorphism;
 pub mod property;

src/lowering/calls.rs

@@ -62,28 +62,27 @@ use plc_source::source_location::SourceLocation;
 
 use crate::{
     index::Index,
-    lowering::helper::create_member_reference_with_location,
-    resolver::{AnnotationMap, StatementAnnotation},
+    lowering::{helper::create_member_reference_with_location, mutation::MutationTracker},
+    resolver::{AnnotationMap, AstAnnotations, StatementAnnotation},
     typesystem::{DataType, DataTypeInformation},
 };
 
 // Performs lowering for aggregate types defined in functions
 #[derive(Default)]
 pub struct AggregateTypeLowerer {
     pub index: Option<Index>,
-    pub annotation: Option<Box<dyn AnnotationMap>>,
+    pub annotation: Option<AstAnnotations>,
     pub id_provider: IdProvider,
     /// New statements to be added during visit, that should happen before the call. This should always be drained when read
     pre_stmts: Vec<Vec<AstNode>>,
     /// New statements to be added during visit, that should happen after the call. This should always be drained when read
     post_stmts: Vec<Vec<AstNode>>,
     counter: AtomicI32,
-    /// Incremented every time the lowerer mutates the AST in a way that
-    /// affects the global index (POU signature changes) or the annotations
-    /// (new pre/post statements inserted). Compared before / after a unit
-    /// visit so the participant can decide whether to re-index / re-annotate
-    /// that unit.
-    mutations: usize,
+    /// Set whenever the lowerer mutates the AST in a way that affects the
+    /// global index (POU signature changes) or the annotations (new pre/post
+    /// statements inserted). Reset before each unit walk so the participant
+    /// can decide per unit whether to re-index / re-annotate.
+    mutations: MutationTracker,
 }
 
 impl AggregateTypeLowerer {
@@ -102,10 +101,21 @@ impl AggregateTypeLowerer {
     /// Visits a single unit and returns whether the lowerer actually mutated
     /// it. The caller uses the result to scope re-indexing / re-annotation
     /// to only the units that need it.
+    ///
+    /// Pair with [`signature_changed_pous`](Self::signature_changed_pous) to
+    /// learn *which* POUs in this unit had their signature rewritten — body
+    /// edits flip the dirty flag but do not invalidate callers of every POU
+    /// in the unit.
     pub fn visit_unit_tracked(&mut self, unit: &mut CompilationUnit) -> bool {
-        let before = self.mutations;
+        self.mutations.reset();
         self.visit_compilation_unit(unit);
-        self.mutations > before
+        self.mutations.is_dirty()
+    }
+
+    /// Names of POUs whose public signatures were rewritten during the most
+    /// recent [`visit_unit_tracked`](Self::visit_unit_tracked) call.
+    pub fn signature_changed_pous(&self) -> Vec<String> {
+        self.mutations.signature_changed_pous()
     }
 
     fn steal_and_walk_list(&mut self, list: &mut Vec<AstNode>) {
@@ -127,7 +137,7 @@ impl AggregateTypeLowerer {
     fn push_pre_statement(&mut self, stmt: AstNode) {
         if let Some(stmts) = self.pre_stmts.last_mut() {
             stmts.push(stmt);
-            self.mutations += 1;
+            self.mutations.touch();
         } else {
             unreachable!("Statement lists should exist at this point");
         }
@@ -140,7 +150,7 @@ impl AggregateTypeLowerer {
     fn push_post_statement(&mut self, stmt: AstNode) {
         if let Some(stmts) = self.post_stmts.last_mut() {
             stmts.push(stmt);
-            self.mutations += 1;
+            self.mutations.touch();
         } else {
             unreachable!("Statement lists should exist at this point");
         }
@@ -192,7 +202,7 @@ impl AstVisitorMut for AggregateTypeLowerer {
         // If the return type is aggregate, remove it from the signature and add a matching variable
         // in a VAR_IN_OUT block
         if index.get_effective_type_or_void_by_name(&return_type_name).is_aggregate_type() {
-            self.mutations += 1;
+            self.mutations.signature_changed(&pou.name);
             let original_return = pou.return_type.take().unwrap();
             let location = original_return.get_location();
             //Create a new return type for the pou
@@ -374,8 +384,8 @@ impl AstVisitorMut for AggregateTypeLowerer {
             .iter()
             .enumerate()
             .any(|(param_index, param)|
-                is_output_assignment_and_type_cast_needed(param, annotation.as_ref(), index, &qualified_name, param_index, function_entry.is_method())
-                || is_output_assignment_and_has_direct_access(param, annotation.as_ref(), index, &qualified_name, param_index, function_entry.is_method()))
+                is_output_assignment_and_type_cast_needed(param, annotation, index, &qualified_name, param_index, function_entry.is_method())
+                || is_output_assignment_and_has_direct_access(param, annotation, index, &qualified_name, param_index, function_entry.is_method()))
             // Stateful structs (such as function blocks) have their own output assignment mechanism in codegen,
             // and are not handled by this lowerer
             && (function_entry.is_function() || function_entry.is_method())
@@ -392,7 +402,7 @@ impl AstVisitorMut for AggregateTypeLowerer {
                     (&self.counter, self.id_provider.clone()),
                     param,
                     (&qualified_name, &return_name, function_entry.is_method()),
-                    annotation.as_ref(),
+                    annotation,
                     index,
                     param_index,
                     (&mut pre_statements, &mut expressions, &mut post_statements),
@@ -803,6 +813,7 @@ mod tests {
 
     use crate::index::indexer;
     use crate::lowering::calls::AggregateTypeLowerer;
+    use crate::resolver::AstAnnotations;
     use crate::test_utils::tests::{
         annotate_and_lower_with_ids, annotate_with_ids, index as test_index, index_and_lower,
         index_unit_with_id, index_with_ids,
@@ -972,7 +983,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1013,7 +1024,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[2]);
     }
@@ -1051,7 +1062,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1091,7 +1102,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1131,7 +1142,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1171,7 +1182,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1214,7 +1225,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1252,7 +1263,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1290,7 +1301,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1330,7 +1341,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1371,7 +1382,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[1]);
     }
@@ -1399,7 +1410,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[0]);
     }
@@ -1627,7 +1638,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[2].statements[0]);
     }
@@ -1664,7 +1675,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
         assert_debug_snapshot!(unit.implementations[2].statements[0]);
     }
@@ -1708,7 +1719,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
 
         let implementations = &unit.implementations;
@@ -1765,7 +1776,7 @@ mod tests {
         lowerer.visit_compilation_unit(&mut unit);
         lowerer.index.replace(index_unit_with_id(&unit, id_provider.clone()));
         let annotations = annotate_with_ids(&unit, lowerer.index.as_mut().unwrap(), id_provider.clone());
-        lowerer.annotation.replace(Box::new(annotations));
+        lowerer.annotation.replace(AstAnnotations::new(annotations, id_provider.clone().next_id()));
         lowerer.visit_compilation_unit(&mut unit);
 
         let implementations = &unit.implementations;