Skip to main content

miden_precompiles/math/uint/
precompile.rs

1//! Precompile for fixed 256-bit uint arithmetic domains in the deferred framework.
2
3use alloc::vec::Vec;
4
5use miden_core::{
6    Felt, ZERO,
7    deferred::{
8        DeferredContext, DeferredError, Digest, Node, NodeType, Payload, Precompile,
9        PrecompileError, Tag, precompile_id,
10    },
11};
12
13use super::{Limbs, ONE_LIMBS, TWO_LIMBS, UintDomain, ZERO_LIMBS};
14
15/// Recognized uint binary operation.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum UintBinaryOp {
18    Add,
19    Sub,
20    Mul,
21}
22
23/// Structural view of a uint precompile node.
24///
25/// Operation variants expose only the structural child digests in the node payload. Value variants
26/// expose the canonical domain and limbs after the same payload checks used by evaluation.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum UintNodeRef {
29    /// Canonical uint value.
30    Value { domain: UintDomain, limbs: Limbs },
31    /// Addition over two structural child digests.
32    Add { lhs: Digest, rhs: Digest },
33    /// Subtraction over two structural child digests.
34    Sub { lhs: Digest, rhs: Digest },
35    /// Multiplication over two structural child digests.
36    Mul { lhs: Digest, rhs: Digest },
37    /// Equality assertion over two structural child digests.
38    Eq { lhs: Digest, rhs: Digest },
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42enum UintOp {
43    Value(UintDomain),
44    Binary(UintBinaryOp),
45    Eq,
46}
47
48impl UintOp {
49    fn decode(args: [Felt; 3]) -> Option<Self> {
50        match args[0].as_canonical_u64() {
51            UintPrecompile::VALUE_OP_ID if args[2] == ZERO => {
52                Some(Self::Value(domain_from_bound_ptr_arg(args[1])?))
53            },
54            UintPrecompile::ADD_OP_ID if args[1] == ZERO && args[2] == ZERO => {
55                Some(Self::Binary(UintBinaryOp::Add))
56            },
57            UintPrecompile::SUB_OP_ID if args[1] == ZERO && args[2] == ZERO => {
58                Some(Self::Binary(UintBinaryOp::Sub))
59            },
60            UintPrecompile::MUL_OP_ID if args[1] == ZERO && args[2] == ZERO => {
61                Some(Self::Binary(UintBinaryOp::Mul))
62            },
63            UintPrecompile::EQ_OP_ID if args[1] == ZERO && args[2] == ZERO => Some(Self::Eq),
64            _ => None,
65        }
66    }
67
68    const fn node_type(self) -> NodeType {
69        match self {
70            Self::Value(_) => NodeType::Data,
71            Self::Binary(_) | Self::Eq => NodeType::Join,
72        }
73    }
74}
75
76fn domain_from_bound_ptr_arg(bound_ptr: Felt) -> Option<UintDomain> {
77    let ptr = bound_ptr.as_canonical_u64();
78    if ptr > u32::MAX as u64 {
79        return None;
80    }
81    UintDomain::from_bound_ptr(ptr as u32)
82}
83
84enum UintNode {
85    Value {
86        domain: UintDomain,
87        limbs: Limbs,
88    },
89    BinaryOp {
90        op: UintBinaryOp,
91        lhs: Digest,
92        rhs: Digest,
93    },
94    Eq {
95        lhs: Digest,
96        rhs: Digest,
97    },
98}
99
100impl UintNode {
101    fn parse(op: UintOp, payload: &Payload) -> Result<Self, PrecompileError> {
102        Ok(match op {
103            UintOp::Value(domain) => {
104                let limbs = decode_limbs(payload.as_value()?)?;
105                if !domain.is_canonical(&limbs) {
106                    return Err(DeferredError::InvalidPayload.into());
107                }
108                Self::Value { domain, limbs }
109            },
110            UintOp::Binary(op) => {
111                let (lhs, rhs) = payload.as_join()?;
112                Self::BinaryOp { op, lhs, rhs }
113            },
114            UintOp::Eq => {
115                let (lhs, rhs) = payload.as_join()?;
116                Self::Eq { lhs, rhs }
117            },
118        })
119    }
120}
121
122/// Precompile for 256-bit arithmetic over fixed uint domains.
123#[derive(Clone, Copy, Debug, Default)]
124pub struct UintPrecompile;
125
126impl UintPrecompile {
127    /// Stable precompile name used to derive this precompile's tag id.
128    pub const NAME: &'static str = "uint256";
129
130    /// Operation discriminants owned by this precompile.
131    pub const VALUE_OP_ID: u64 = 0;
132    pub const ADD_OP_ID: u64 = 1;
133    pub const SUB_OP_ID: u64 = 2;
134    pub const MUL_OP_ID: u64 = 3;
135    pub const EQ_OP_ID: u64 = 4;
136
137    /// Stable precompile id derived from [`Self::NAME`].
138    pub fn id() -> Felt {
139        precompile_id(Self::NAME)
140    }
141
142    /// Builds a canonical uint `VALUE` tag for `domain`.
143    pub fn value_tag(domain: UintDomain) -> Tag {
144        let op_id = Felt::new(Self::VALUE_OP_ID).expect("uint VALUE op id must fit in a felt");
145        Tag::precompile(Self::id(), [op_id, Felt::from(domain.bound_ptr()), ZERO])
146            .expect("uint precompile id is not framework-reserved")
147    }
148
149    /// Builds a uint operation tag from `op_id`.
150    ///
151    /// Known operation ids decode to their declared shapes; unknown ids produce a tag that this
152    /// precompile rejects. Operand `VALUE` nodes carry the concrete domain.
153    pub fn op_tag(op_id: u64) -> Tag {
154        let op_id = Felt::new(op_id).expect("uint op id must fit in a felt");
155        Tag::precompile(Self::id(), [op_id, ZERO, ZERO])
156            .expect("uint precompile id is not framework-reserved")
157    }
158
159    /// Builds a uint `VALUE` node from trusted canonical limbs.
160    ///
161    /// Callers must ensure `limbs` is canonical for `domain`. Debug builds assert this
162    /// precondition; registration and evaluation validate nodes constructed from untrusted
163    /// limbs.
164    pub fn value_node(domain: UintDomain, limbs: Limbs) -> Node {
165        debug_assert!(domain.is_canonical(&limbs));
166        Node::value(Self::value_tag(domain), limbs.map(Felt::from_u32))
167            .expect("value tag is precompile-owned")
168    }
169
170    /// Decodes a canonical uint `VALUE` node for `domain`.
171    pub fn decode_value_node(node: &Node, domain: UintDomain) -> Result<Limbs, DeferredError> {
172        Self::limbs_from_value_node(node, domain)
173    }
174
175    /// Decodes a uint precompile node without evaluating its children.
176    ///
177    /// Returns `Ok(None)` when `node` belongs to another precompile. Owned operation nodes return
178    /// their structural child digests directly from the payload.
179    pub fn decode_node(node: &Node) -> Result<Option<UintNodeRef>, PrecompileError> {
180        if node.tag().id() != Self::id() {
181            return Ok(None);
182        }
183
184        let op = UintOp::decode(node.tag().args()).ok_or(PrecompileError::InvalidNode)?;
185        let parsed = UintNode::parse(op, node.payload())?;
186        Ok(Some(match parsed {
187            UintNode::Value { domain, limbs } => UintNodeRef::Value { domain, limbs },
188            UintNode::BinaryOp { op: UintBinaryOp::Add, lhs, rhs } => UintNodeRef::Add { lhs, rhs },
189            UintNode::BinaryOp { op: UintBinaryOp::Sub, lhs, rhs } => UintNodeRef::Sub { lhs, rhs },
190            UintNode::BinaryOp { op: UintBinaryOp::Mul, lhs, rhs } => UintNodeRef::Mul { lhs, rhs },
191            UintNode::Eq { lhs, rhs } => UintNodeRef::Eq { lhs, rhs },
192        }))
193    }
194
195    pub(crate) fn limbs_from_typed_value_node(
196        node: &Node,
197    ) -> Result<(UintDomain, Limbs), DeferredError> {
198        let Some(UintOp::Value(domain)) = UintOp::decode(node.tag().args()) else {
199            return Err(DeferredError::InvalidPayload);
200        };
201        let payload = node.payload_for_tag(Self::value_tag(domain))?;
202        let limbs = decode_limbs(payload.as_value()?)?;
203        if !domain.is_canonical(&limbs) {
204            return Err(DeferredError::InvalidPayload);
205        }
206        Ok((domain, limbs))
207    }
208
209    pub(crate) fn limbs_from_value_node(
210        node: &Node,
211        domain: UintDomain,
212    ) -> Result<Limbs, DeferredError> {
213        let (actual_domain, limbs) = Self::limbs_from_typed_value_node(node)?;
214        if actual_domain != domain {
215            return Err(DeferredError::InvalidPayload);
216        }
217        Ok(limbs)
218    }
219
220    fn evaluate_value_pair(
221        context: &mut DeferredContext<'_>,
222        lhs: Digest,
223        rhs: Digest,
224    ) -> Result<(UintDomain, Limbs, Limbs), PrecompileError> {
225        let (lhs, rhs) = context.evaluate_digest_pair(lhs, rhs)?;
226        let lhs = context.get_node(&lhs).ok_or(PrecompileError::MissingNode)?;
227        let rhs = context.get_node(&rhs).ok_or(PrecompileError::MissingNode)?;
228
229        let (lhs_domain, lhs) = Self::limbs_from_typed_value_node(lhs)?;
230        let (rhs_domain, rhs) = Self::limbs_from_typed_value_node(rhs)?;
231        if lhs_domain != rhs_domain {
232            return Err(DeferredError::InvalidPayload.into());
233        }
234
235        Ok((lhs_domain, lhs, rhs))
236    }
237}
238
239impl Precompile for UintPrecompile {
240    fn name(&self) -> &'static str {
241        Self::NAME
242    }
243
244    fn id(&self) -> Felt {
245        Self::id()
246    }
247
248    fn init(&self) -> Vec<Node> {
249        let mut nodes = Vec::new();
250        for domain in UintDomain::ALL {
251            for value in [ZERO_LIMBS, ONE_LIMBS, TWO_LIMBS] {
252                nodes.push(Self::value_node(domain, value));
253            }
254            if let Some(max) = domain.max() {
255                nodes.push(Self::value_node(domain, max));
256            }
257            if let Some(constants) = domain.field_constants() {
258                for value in constants {
259                    nodes.push(Self::value_node(domain, value));
260                }
261            }
262        }
263        nodes
264    }
265
266    fn decode(&self, args: [Felt; 3]) -> Option<NodeType> {
267        let op = UintOp::decode(args)?;
268        Some(op.node_type())
269    }
270
271    fn evaluate(
272        &self,
273        args: [Felt; 3],
274        payload: &Payload,
275        context: &mut DeferredContext<'_>,
276    ) -> Result<Node, PrecompileError> {
277        let op = UintOp::decode(args).ok_or(PrecompileError::InvalidNode)?;
278
279        match UintNode::parse(op, payload)? {
280            UintNode::Value { domain, limbs } => Ok(Self::value_node(domain, limbs)),
281            UintNode::BinaryOp { op, lhs, rhs } => {
282                let (domain, lhs, rhs) = Self::evaluate_value_pair(context, lhs, rhs)?;
283                let value = match op {
284                    UintBinaryOp::Add => domain.add(lhs, rhs),
285                    UintBinaryOp::Sub => domain.sub(lhs, rhs),
286                    UintBinaryOp::Mul => domain.mul(lhs, rhs),
287                };
288                Ok(Self::value_node(domain, value))
289            },
290            UintNode::Eq { lhs, rhs } => {
291                let (_, lhs, rhs) = Self::evaluate_value_pair(context, lhs, rhs)?;
292                if lhs == rhs {
293                    Ok(Node::TRUE)
294                } else {
295                    Err(PrecompileError::AssertionFailed)
296                }
297            },
298        }
299    }
300}
301
302pub(crate) fn decode_limbs(felts: &[Felt; 8]) -> Result<Limbs, DeferredError> {
303    let mut limbs = [0u32; 8];
304    for (i, felt) in felts.iter().enumerate() {
305        let v = felt.as_canonical_u64();
306        if v > u32::MAX as u64 {
307            return Err(DeferredError::InvalidPayload);
308        }
309        limbs[i] = v as u32;
310    }
311    Ok(limbs)
312}
313
314#[cfg(test)]
315mod tests {
316    use alloc::sync::Arc;
317
318    use miden_core::deferred::DeferredState;
319
320    use super::*;
321
322    fn state() -> DeferredState {
323        DeferredState::new(Arc::new(crate::registry())).expect("precompile init must succeed")
324    }
325
326    fn evaluate(state: &mut DeferredState, node: Node) -> Result<Node, PrecompileError> {
327        let digest = state.register(node)?;
328        state.require_canonical_node(digest).map(|(_, node)| node.clone())
329    }
330
331    fn assert_invalid_payload<T>(result: Result<T, PrecompileError>) {
332        let Err(error) = result else {
333            panic!("expected invalid payload");
334        };
335        assert!(
336            matches!(error.root(), PrecompileError::Other(DeferredError::InvalidPayload)),
337            "expected invalid payload, got {error:?}",
338        );
339    }
340
341    fn limbs(value: u32) -> Limbs {
342        let mut limbs = [0; 8];
343        limbs[0] = value;
344        limbs
345    }
346
347    #[test]
348    fn decode_uses_bound_ptr_value_and_op_tags() {
349        let precompile = UintPrecompile;
350        let domain = UintDomain::K1Base;
351        let bound_ptr = Felt::from(domain.bound_ptr());
352
353        assert_eq!(
354            UintPrecompile::value_tag(domain).as_word(),
355            [UintPrecompile::id(), Felt::from_u32(0), bound_ptr, ZERO],
356        );
357        assert_eq!(
358            precompile.decode(UintPrecompile::value_tag(domain).args()),
359            Some(NodeType::Data)
360        );
361
362        assert_eq!(
363            UintPrecompile::op_tag(UintPrecompile::ADD_OP_ID).as_word(),
364            [UintPrecompile::id(), Felt::from_u32(1), ZERO, ZERO],
365        );
366        assert_eq!(
367            precompile.decode(UintPrecompile::op_tag(UintPrecompile::ADD_OP_ID).args()),
368            Some(NodeType::Join)
369        );
370
371        let mut add_with_bound = UintPrecompile::op_tag(UintPrecompile::ADD_OP_ID).args();
372        add_with_bound[1] = bound_ptr;
373        assert_eq!(precompile.decode(add_with_bound), None);
374        assert_eq!(precompile.decode(UintPrecompile::op_tag(99).args()), None);
375
376        assert_eq!(precompile.decode([Felt::from_u32(0), Felt::new_unchecked(99), ZERO]), None);
377        assert_eq!(precompile.decode([Felt::from_u32(0), ZERO, ZERO]), None);
378        assert_eq!(precompile.decode([Felt::from_u32(0), bound_ptr, Felt::from_u32(1)]), None);
379        assert_eq!(
380            precompile.decode([Felt::from_u32(0), Felt::new_unchecked(u32::MAX as u64 + 1), ZERO,]),
381            None
382        );
383    }
384
385    #[test]
386    fn data_shape_does_not_bypass_one_chunk_value_semantics() {
387        let domain = UintDomain::K1Base;
388        let tag = UintPrecompile::value_tag(domain);
389        let node = Node::try_data(tag, alloc::vec![[ZERO; 8], [ZERO; 8]])
390            .expect("multi-chunk data is structurally valid");
391        let precompile = UintPrecompile;
392        assert_eq!(precompile.decode(tag.args()), Some(NodeType::Data));
393
394        let mut state = state();
395        assert_invalid_payload(state.register(node));
396    }
397
398    #[test]
399    fn decode_node_exposes_structural_uint_nodes() {
400        let domain = UintDomain::K1Base;
401        let lhs = UintPrecompile::value_node(domain, limbs(9));
402        let rhs = UintPrecompile::value_node(domain, limbs(4));
403
404        assert_eq!(
405            UintPrecompile::decode_node(&lhs).unwrap(),
406            Some(UintNodeRef::Value { domain, limbs: limbs(9) })
407        );
408        assert_eq!(UintPrecompile::decode_node(&Node::TRUE).unwrap(), None);
409
410        for (op_id, expected) in [
411            (
412                UintPrecompile::ADD_OP_ID,
413                UintNodeRef::Add { lhs: lhs.digest(), rhs: rhs.digest() },
414            ),
415            (
416                UintPrecompile::SUB_OP_ID,
417                UintNodeRef::Sub { lhs: lhs.digest(), rhs: rhs.digest() },
418            ),
419            (
420                UintPrecompile::MUL_OP_ID,
421                UintNodeRef::Mul { lhs: lhs.digest(), rhs: rhs.digest() },
422            ),
423            (
424                UintPrecompile::EQ_OP_ID,
425                UintNodeRef::Eq { lhs: lhs.digest(), rhs: rhs.digest() },
426            ),
427        ] {
428            let node = Node::join(UintPrecompile::op_tag(op_id), lhs.digest(), rhs.digest())
429                .expect("tag is uint-owned");
430            assert_eq!(UintPrecompile::decode_node(&node).unwrap(), Some(expected));
431        }
432
433        let invalid_tag = Tag::precompile(UintPrecompile::id(), [Felt::from_u32(99), ZERO, ZERO])
434            .expect("tag is precompile-owned");
435        let invalid = Node::join(invalid_tag, lhs.digest(), rhs.digest()).unwrap();
436        assert!(matches!(
437            UintPrecompile::decode_node(&invalid),
438            Err(PrecompileError::InvalidNode)
439        ));
440    }
441
442    #[test]
443    fn same_domain_binary_operation_succeeds() {
444        let mut state = state();
445        let lhs = UintPrecompile::value_node(UintDomain::U256, limbs(3));
446        let rhs = UintPrecompile::value_node(UintDomain::U256, limbs(4));
447        state.register(lhs.clone()).expect("lhs must register");
448        state.register(rhs.clone()).expect("rhs must register");
449
450        let node = Node::join(
451            UintPrecompile::op_tag(UintPrecompile::ADD_OP_ID),
452            lhs.digest(),
453            rhs.digest(),
454        )
455        .expect("tag is uint-owned");
456        let expected = UintPrecompile::value_node(UintDomain::U256, limbs(7));
457
458        assert_eq!(evaluate(&mut state, node).unwrap(), expected);
459    }
460
461    #[test]
462    fn mixed_domain_binary_operation_fails() {
463        let mut state = state();
464        let lhs = UintPrecompile::value_node(UintDomain::U256, limbs(1));
465        let rhs = UintPrecompile::value_node(UintDomain::K1Base, limbs(1));
466        state.register(lhs.clone()).expect("lhs must register");
467        state.register(rhs.clone()).expect("rhs must register");
468
469        let node = Node::join(
470            UintPrecompile::op_tag(UintPrecompile::ADD_OP_ID),
471            lhs.digest(),
472            rhs.digest(),
473        )
474        .expect("tag is uint-owned");
475
476        assert_invalid_payload(evaluate(&mut state, node));
477    }
478
479    #[test]
480    fn decode_limbs_accepts_u32_boundary_and_rejects_larger_felts() {
481        let felts = [Felt::from_u32(u32::MAX); 8];
482        assert_eq!(decode_limbs(&felts).unwrap(), [u32::MAX; 8]);
483
484        let mut felts = [Felt::from_u32(0); 8];
485        felts[3] = Felt::new_unchecked(u32::MAX as u64 + 1);
486        assert_eq!(decode_limbs(&felts), Err(DeferredError::InvalidPayload));
487    }
488}