optimize uint_atom() for small atoms

Chia-Network · Feb 5, 2024 · ed2f13e · ed2f13e
1 parent cf68113
commit ed2f13e
Showing 1 changed file with 32 additions and 31 deletions.
diff --git 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<const SIZE: usize>(
     args: NodePtr,
     op_name: &str,
 ) -> Result<u64, EvalErr> {
-    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)]