optimize op_mul for small int

Chia-Network · Feb 2, 2024 · 8bdf69b · 8bdf69b
1 parent d29c9f6
commit 8bdf69b
Showing 1 changed file with 24 additions and 7 deletions.
diff --git a/src/more_ops.rs b/src/more_ops.rs
@@ -440,15 +440,32 @@ pub fn op_multiply(a: &mut Allocator, mut input: NodePtr, max_cost: Cost) -> Res
             continue;
         }
 
-        let (v0, l1) = int_atom(a, arg, "*")?;
-
-        total *= v0;
-        cost += MUL_COST_PER_OP;
+        a.visit_node(arg, |node| -> Result<(), EvalErr> {
+            match node {
+                NodeVisitor::Buffer(buf) => {
+                    use crate::number::number_from_u8;
+                    total *= number_from_u8(buf);
+                    let l1 = buf.len();
 
-        cost += (l0 + l1) as Cost * MUL_LINEAR_COST_PER_BYTE;
-        cost += (l0 * l1) as Cost / MUL_SQUARE_COST_PER_BYTE_DIVIDER;
+                    cost += MUL_COST_PER_OP;
+                    cost += (l0 + l1) as Cost * MUL_LINEAR_COST_PER_BYTE;
+                    cost += (l0 * l1) as Cost / MUL_SQUARE_COST_PER_BYTE_DIVIDER;
+                    l0 = limbs_for_int(&total);
+                    Ok(())
+                }
+                NodeVisitor::U32(val) => {
+                    total *= *val;
+                    let l1 = len_for_value(*val);
 
-        l0 = limbs_for_int(&total);
+                    cost += MUL_COST_PER_OP;
+                    cost += (l0 + l1) as Cost * MUL_LINEAR_COST_PER_BYTE;
+                    cost += (l0 * l1) as Cost / MUL_SQUARE_COST_PER_BYTE_DIVIDER;
+                    l0 = limbs_for_int(&total);
+                    Ok(())
+                }
+                NodeVisitor::Pair(_, _) => err(arg, "* requires int args"),
+            }
+        })?;
     }
     let total = a.new_number(total)?;
     Ok(malloc_cost(a, cost, total))