diff --git a/src/op_utils.rs b/src/op_utils.rs index a2074e9c..afca6667 100644 --- a/src/op_utils.rs +++ b/src/op_utils.rs @@ -1,4 +1,4 @@ -use crate::allocator::{Allocator, NodePtr, SExp}; +use crate::allocator::{Allocator, NodePtr, NodeVisitor, SExp}; use crate::cost::Cost; use crate::err_utils::err; use crate::number::Number; @@ -279,37 +279,38 @@ pub fn uint_atom( args: NodePtr, op_name: &str, ) -> Result { - let bytes = match a.sexp(args) { - SExp::Atom => a.atom(args), - _ => { - return err(args, &format!("{op_name} requires int arg")); + a.visit_node(args, |node| { + match node { + NodeVisitor::Buffer(bytes) => { + if bytes.is_empty() { + return Ok(0); + } + + if (bytes[0] & 0x80) != 0 { + return err(args, &format!("{op_name} requires positive int arg")); + } + + // strip leading zeros + let mut buf: &[u8] = bytes; + while !buf.is_empty() && buf[0] == 0 { + buf = &buf[1..]; + } + + if buf.len() > SIZE { + return err(args, &format!("{op_name} requires u{} arg", SIZE * 8)); + } + + let mut ret = 0; + for b in buf { + ret <<= 8; + ret |= *b as u64; + } + Ok(ret) + } + NodeVisitor::U32(val) => Ok(*val as u64), + NodeVisitor::Pair(_, _) => err(args, &format!("{op_name} requires int arg")), } - }; - - if bytes.is_empty() { - return Ok(0); - } - - if (bytes[0] & 0x80) != 0 { - return err(args, &format!("{op_name} requires positive int arg")); - } - - // strip leading zeros - let mut buf: &[u8] = bytes; - while !buf.is_empty() && buf[0] == 0 { - buf = &buf[1..]; - } - - if buf.len() > SIZE { - return err(args, &format!("{op_name} requires u{} arg", SIZE * 8)); - } - - let mut ret = 0; - for b in buf { - ret <<= 8; - ret |= *b as u64; - } - Ok(ret) + }) } #[cfg(test)]