Add MASM lint mode for advice taint analysis

codegen/masm/src/emit/binary.rs

@@ -870,6 +870,16 @@ impl OpEmitter<'_> {
         self.push(ty);
     }
 
+    pub fn exp_u32_exponent(&mut self, span: SourceSpan) {
+        let rhs = self.pop().expect("operand stack is empty");
+        let lhs = self.pop().expect("operand stack is empty");
+        let ty = lhs.ty();
+        assert_eq!(ty, Type::Felt, "expected exp.u32 base to be felt");
+        assert_eq!(rhs.ty(), Type::Felt, "expected exp.u32 exponent to be felt");
+        self.emit(masm::Instruction::ExpBitLength(32), span);
+        self.push(ty);
+    }
+
     #[allow(unused)]
     pub fn exp_imm(&mut self, imm: Immediate, span: SourceSpan) {
         let lhs = self.pop().expect("operand stack is empty");

codegen/masm/src/emit/unary.rs

@@ -306,6 +306,7 @@ impl OpEmitter<'_> {
                 self.felt_to_int(dst_bits, span);
             }
             // u32
+            (Type::U32, Type::Felt) => (),
             (Type::U32, Type::I64 | Type::U64 | Type::I128) => self.zext_int32(dst_bits, span),
             (Type::U32, Type::I32) => self.assert_i32(span),
             (Type::U32, Type::U16 | Type::U8 | Type::I1) => {

codegen/masm/src/lower/lowering.rs

@@ -682,7 +682,12 @@ impl HirLowering for arith::MulOverflowing {
 
 impl HirLowering for arith::Exp {
     fn emit(&self, emitter: &mut BlockEmitter<'_>) -> Result<(), Report> {
-        emitter.inst_emitter(self.as_operation()).exp(self.span());
+        let mut emitter = emitter.inst_emitter(self.as_operation());
+        if *self.get_exponent_must_be_u32() {
+            emitter.exp_u32_exponent(self.span());
+        } else {
+            emitter.exp(self.span());
+        }
         Ok(())
     }
 }

dialects/arith/src/builders.rs

@@ -359,11 +359,24 @@ pub trait ArithOpBuilder<'f, B: ?Sized + Builder> {
 
     /// Exponentiation
     fn exp(&mut self, lhs: ValueRef, rhs: ValueRef, span: SourceSpan) -> Result<ValueRef, Report> {
-        let op_builder = self.builder_mut().create::<crate::ops::Exp, _>(span);
+        let op_builder = self.builder_mut().create::<crate::ops::Exp, (ValueRef, ValueRef)>(span);
         let op = op_builder(lhs, rhs)?;
         Ok(op.borrow().result().as_value_ref())
     }
 
+    /// Exponentiation whose exponent operand must be u32-range constrained.
+    fn exp_u32_exponent(
+        &mut self,
+        lhs: ValueRef,
+        rhs: ValueRef,
+        span: SourceSpan,
+    ) -> Result<ValueRef, Report> {
+        let op_builder =
+            self.builder_mut().create::<crate::ops::Exp, (ValueRef, ValueRef, bool)>(span);
+        let op = op_builder(lhs, rhs, true)?;
+        Ok(op.borrow().result().as_value_ref())
+    }
+
     /// Compute 2^n
     fn pow2(&mut self, n: ValueRef, span: SourceSpan) -> Result<ValueRef, Report> {
         let op_builder = self.builder_mut().create::<crate::ops::Pow2, _>(span);

dialects/arith/src/ops/binary.rs

@@ -2,7 +2,7 @@ use alloc::rc::Rc;
 
 use midenc_hir::{
     derive::{EffectOpInterface, OpParser, OpPrinter, operation},
-    dialects::builtin::attributes::OverflowAttr,
+    dialects::builtin::attributes::{BoolAttr, OverflowAttr},
     effects::*,
     traits::*,
     *,
@@ -184,19 +184,39 @@ infer_return_ty_for_binary_op!(MulOverflowing, overflowed: Type::I1);
 #[operation(
     dialect = ArithDialect,
     traits(BinaryOp, SameTypeOperands, SameOperandsAndResultType),
-    implements(InferTypeOpInterface, MemoryEffectOpInterface, OpPrinter)
+    implements(
+        InferTypeOpInterface,
+        MemoryEffectOpInterface,
+        OperandRangeRequirementOpInterface,
+        OpPrinter
+    )
 )]
 pub struct Exp {
     #[operand]
     lhs: IntFelt,
     #[operand]
     rhs: IntFelt,
+    #[attr]
+    #[default]
+    exponent_must_be_u32: BoolAttr,
     #[result]
     result: IntFelt,
 }
 
 infer_return_ty_for_binary_op!(Exp);
 
+impl OperandRangeRequirementOpInterface for Exp {
+    fn operand_range_requirement(&self, operand_index: usize) -> OperandRangeRequirement {
+        if operand_index == 1 && *self.get_exponent_must_be_u32() {
+            OperandRangeRequirement::Required(ValueRangeConstraint::Type(Type::U32))
+        } else {
+            default_operand_range_requirement(self.as_operation(), operand_index)
+                .map(OperandRangeRequirement::Required)
+                .unwrap_or(OperandRangeRequirement::None)
+        }
+    }
+}
+
 /// Unsigned integer division, traps on division by zero
 #[derive(EffectOpInterface, OpPrinter, OpParser)]
 #[operation(

dialects/hir/src/analyses/advice_taint/lattice.rs

@@ -156,6 +156,7 @@ impl LatticeLike for AdviceTaintValue {
 }
 
 const MAX_CALL_CONTEXT_DEPTH: usize = 4;
+const MAX_CALL_CONTEXTS: usize = 32;
 
 type CallContext = SmallVec<[CallContextFrame; MAX_CALL_CONTEXT_DEPTH]>;
 pub(super) type AdviceTaintSparseLattice = Lattice<ContextualAdviceTaintValue>;
@@ -225,6 +226,24 @@ impl ContextualAdviceTaintValue {
         origins.into_iter()
     }
 
+    pub(super) fn call_context_spans_containing_origin(
+        &self,
+        origin: AdviceTaintOrigin,
+    ) -> Vec<SourceSpan> {
+        let mut spans = Vec::new();
+        for (context, taint) in self.contexts.iter() {
+            if !taint.contains_origin(origin) {
+                continue;
+            }
+            for frame in context {
+                if !spans.contains(&frame.span) {
+                    spans.push(frame.span);
+                }
+            }
+        }
+        spans
+    }
+
     pub fn mark_reported(&self) -> Self {
         Self {
             contexts: self
@@ -300,10 +319,24 @@ impl ContextualAdviceTaintValue {
         context: CallContext,
         taint: AdviceTaintValue,
     ) {
+        if context.is_empty() && !taint.is_clean() && !contexts.is_empty() {
+            let collapsed =
+                contexts.values().fold(taint, |acc, taint| LatticeLike::join(&acc, taint));
+            contexts.clear();
+            contexts.insert(CallContext::new(), collapsed);
+            return;
+        }
+        if let Some(empty) = contexts.get_mut(&CallContext::new())
+            && !empty.is_clean()
+        {
+            *empty = LatticeLike::join(empty, &taint);
+            return;
+        }
         contexts
             .entry(context)
             .and_modify(|current| *current = LatticeLike::join(current, &taint))
             .or_insert(taint);
+        normalize_contexts(contexts);
     }
 }
 
@@ -342,6 +375,28 @@ fn push_call_context(context: &CallContext, frame: CallContextFrame) -> CallCont
     pushed
 }
 
+fn normalize_contexts(contexts: &mut BTreeMap<CallContext, AdviceTaintValue>) {
+    if contexts.values().all(AdviceTaintValue::is_clean) {
+        contexts.clear();
+        contexts.insert(CallContext::new(), AdviceTaintValue::clean());
+        return;
+    }
+
+    if contexts.len() > 1 {
+        contexts.retain(|_, taint| !taint.is_clean());
+    }
+
+    if contexts.len() <= MAX_CALL_CONTEXTS {
+        return;
+    }
+
+    let collapsed = contexts
+        .values()
+        .fold(AdviceTaintValue::clean(), |acc, taint| LatticeLike::join(&acc, taint));
+    contexts.clear();
+    contexts.insert(CallContext::new(), collapsed);
+}
+
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 enum OriginState {
     Unreported,
@@ -385,3 +440,78 @@ pub(super) fn value_taint(value: ValueRef, solver: &DataFlowSolver) -> Contextua
         .map(|state| state.value().clone())
         .unwrap_or_default()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn contextual_taint_collapses_when_call_context_cap_is_exceeded() {
+        let raw = ContextualAdviceTaintValue::raw(SourceSpan::UNKNOWN);
+        let mut joined = ContextualAdviceTaintValue::clean();
+
+        for id in 0..=MAX_CALL_CONTEXTS {
+            let frame = CallContextFrame {
+                id,
+                span: SourceSpan::UNKNOWN,
+            };
+            joined = LatticeLike::join(&joined, &raw.enter_call(frame));
+        }
+
+        assert_eq!(joined.contexts.len(), 1);
+        assert!(joined.contexts.contains_key(&CallContext::new()));
+        assert!(joined.has_unreported_origin());
+    }
+
+    #[test]
+    fn contextual_taint_keeps_distinct_call_contexts_below_cap() {
+        let raw = ContextualAdviceTaintValue::raw(SourceSpan::UNKNOWN);
+        let mut joined = ContextualAdviceTaintValue::clean();
+
+        for id in 0..MAX_CALL_CONTEXTS {
+            let frame = CallContextFrame {
+                id,
+                span: SourceSpan::UNKNOWN,
+            };
+            joined = LatticeLike::join(&joined, &raw.enter_call(frame));
+        }
+
+        assert_eq!(joined.contexts.len(), MAX_CALL_CONTEXTS);
+        assert!(!joined.contexts.contains_key(&CallContext::new()));
+    }
+
+    #[test]
+    fn collapsed_contextual_taint_absorbs_precise_contexts() {
+        let raw = ContextualAdviceTaintValue::raw(SourceSpan::UNKNOWN);
+        let mut precise = ContextualAdviceTaintValue::clean();
+
+        for id in 0..MAX_CALL_CONTEXTS {
+            let frame = CallContextFrame {
+                id,
+                span: SourceSpan::UNKNOWN,
+            };
+            precise = LatticeLike::join(&precise, &raw.enter_call(frame));
+        }
+
+        let collapsed = ContextualAdviceTaintValue::raw(SourceSpan::UNKNOWN);
+        let joined = LatticeLike::join(&collapsed, &precise);
+
+        assert_eq!(joined, collapsed);
+    }
+
+    #[test]
+    fn clean_contextual_taint_stays_canonical_across_call_contexts() {
+        let clean = ContextualAdviceTaintValue::clean();
+        let mut joined = ContextualAdviceTaintValue::clean();
+
+        for id in 0..MAX_CALL_CONTEXTS {
+            let frame = CallContextFrame {
+                id,
+                span: SourceSpan::UNKNOWN,
+            };
+            joined = LatticeLike::join(&joined, &clean.enter_call(frame));
+        }
+
+        assert_eq!(joined, ContextualAdviceTaintValue::clean());
+    }
+}