Rewrite sha256tree nonrecursively.

src/types/blockchain_format/program.py

@@ -13,6 +13,8 @@ from clvm_tools.curry import curry, uncurry
 from src.types.blockchain_format.sized_bytes import bytes32
 from src.util.hash import std_hash
 
+from .tree_hash import sha256_treehash
+
 
 def run_program(
     program,
@@ -55,27 +57,12 @@ class Program(SExp):
     def __str__(self) -> str:
         return bytes(self).hex()
 
-    def _tree_hash(self, precalculated: Set[bytes32]) -> bytes32:
-        """
-        Hash values in `precalculated` are presumed to have been hashed already.
-        """
-        if self.listp():
-            left = self.to(self.first())._tree_hash(precalculated)
-            right = self.to(self.rest())._tree_hash(precalculated)
-            s = b"\2" + left + right
-        else:
-            atom = self.as_atom()
-            if atom in precalculated:
-                return bytes32(atom)
-            s = b"\1" + atom
-        return bytes32(std_hash(s))
-
     def get_tree_hash(self, *args: List[bytes32]) -> bytes32:
         """
         Any values in `args` that appear in the tree
         are presumed to have been hashed already.
         """
-        return self._tree_hash(set(args))
+        return sha256_treehash(self, set(args))
 
     def run_with_cost(self, args) -> Tuple[int, "Program"]:
         prog_args = Program.to(args)

src/types/blockchain_format/tree_hash.py

@@ -0,0 +1,58 @@
+"""
+This is an implementation of `sha256_treehash`, used to calculate
+puzzle hashes in clvm.
+
+This implementation goes to great pains to be non-recursive so we don't
+have to worry about blowing out the python stack.
+"""
+
+from typing import Optional, Set
+
+from clvm import CLVMObject
+
+from src.types.blockchain_format.sized_bytes import bytes32
+from src.util.hash import std_hash
+
+
+def sha256_treehash(sexp: CLVMObject, precalculated: Optional[Set[bytes32]] = None) -> bytes32:
+    """
+    Hash values in `precalculated` are presumed to have been hashed already.
+    """
+
+    if precalculated is None:
+        precalculated = set()
+
+    def handle_sexp(sexp_stack, op_stack, precalculated: Set[bytes32]) -> None:
+        sexp = sexp_stack.pop()
+        if sexp.pair:
+            p0, p1 = sexp.pair
+            sexp_stack.append(p0)
+            sexp_stack.append(p1)
+            op_stack.append(handle_pair)
+            op_stack.append(handle_sexp)
+            op_stack.append(roll)
+            op_stack.append(handle_sexp)
+        else:
+            if sexp.atom in precalculated:
+                r = sexp.atom
+            else:
+                r = std_hash(b"\1" + sexp.atom)
+            sexp_stack.append(r)
+
+    def handle_pair(sexp_stack, op_stack, precalculated) -> None:
+        p0 = sexp_stack.pop()
+        p1 = sexp_stack.pop()
+        sexp_stack.append(std_hash(b"\2" + p0 + p1))
+
+    def roll(sexp_stack, op_stack, precalculated) -> None:
+        p0 = sexp_stack.pop()
+        p1 = sexp_stack.pop()
+        sexp_stack.append(p0)
+        sexp_stack.append(p1)
+
+    sexp_stack = [sexp]
+    op_stack = [handle_sexp]
+    while len(op_stack) > 0:
+        op = op_stack.pop()
+        op(sexp_stack, op_stack, precalculated)
+    return bytes32(sexp_stack[0])