optimize Program.curry()

chia/types/blockchain_format/program.py

@@ -1,12 +1,12 @@
 import io
-from typing import List, Set, Tuple, Optional
+from typing import List, Set, Tuple, Optional, Any
 
 from clvm import SExp
 from clvm.casts import int_from_bytes
 from clvm.EvalError import EvalError
 from clvm.serialize import sexp_from_stream, sexp_to_stream
 from chia_rs import MEMPOOL_MODE, run_chia_program, serialized_length, run_generator
-from clvm_tools.curry import curry, uncurry
+from clvm_tools.curry import uncurry
 
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.hash import std_hash
@@ -89,8 +89,10 @@ def run(self, args) -> "Program":
         return r
 
     def curry(self, *args) -> "Program":
-        cost, r = curry(self, list(args))
-        return Program.to(r)
+        fixed_args: Any = 1
+        for arg in reversed(args):
+            fixed_args = [4, (1, arg), fixed_args]
+        return Program.to([2, (1, self), fixed_args])
 
     def uncurry(self) -> Tuple["Program", "Program"]:
         r = uncurry(self)

tests/clvm/test_program.py

@@ -2,6 +2,9 @@
 
 from chia.types.blockchain_format.program import Program
 from clvm.EvalError import EvalError
+from clvm_tools.curry import uncurry
+from clvm.operators import KEYWORD_TO_ATOM
+from clvm_tools.binutils import assemble, disassemble
 
 
 class TestProgram(TestCase):
@@ -19,3 +22,28 @@ def test_at(self):
 
         self.assertRaises(ValueError, lambda: p.at("q"))
         self.assertRaises(EvalError, lambda: p.at("ff"))
+
+
+def check_idempotency(f, *args):
+    prg = Program.to(f)
+    curried = prg.curry(*args)
+
+    r = disassemble(curried)
+    f_0, args_0 = uncurry(curried)
+
+    assert disassemble(f_0) == disassemble(f)
+    assert disassemble(args_0) == disassemble(Program.to(list(args)))
+    return r
+
+
+def test_curry_uncurry():
+    PLUS = KEYWORD_TO_ATOM["+"][0]
+    f = assemble("(+ 2 5)")
+    actual_disassembly = check_idempotency(f, 200, 30)
+    assert actual_disassembly == f"(a (q {PLUS} 2 5) (c (q . 200) (c (q . 30) 1)))"
+
+    f = assemble("(+ 2 5)")
+    args = assemble("(+ (q . 50) (q . 60))")
+    # passing "args" here wraps the arguments in a list
+    actual_disassembly = check_idempotency(f, args)
+    assert actual_disassembly == f"(a (q {PLUS} 2 5) (c (q {PLUS} (q . 50) (q . 60)) 1))"