Change how preimage works to improve performance (#11)

eth_hash/backends/pycryptodome.py

@@ -1,9 +1,32 @@
 from Crypto.Hash import (
     keccak,
 )
+from eth_hash.preimage import (
+    BasePreImage,
+)
 
 
 def keccak256(prehash: bytes) -> bytes:
-    hasher = keccak.new(digest_bits=256)
-    hasher.update(prehash)
+    hasher = keccak.new(data=prehash, digest_bits=256)
     return hasher.digest()
+
+
+class preimage(BasePreImage):
+    _hash = None
+
+    def __init__(self, prehash) -> None:
+        self._hash = keccak.new(data=prehash, digest_bits=256)
+        # pycryptodome doesn't expose a `copy` mechanism for it's hash objects
+        # so we keep a record of all of the parts for when/if we need to copy
+        # them.
+        self._parts = [prehash]
+
+    def update(self, prehash) -> None:
+        self._hash.update(prehash)
+        self._parts.append(prehash)
+
+    def digest(self) -> bytes:
+        return self._hash.digest()
+
+    def copy(self) -> 'preimage':
+        return preimage(b''.join(self._parts))

eth_hash/backends/pysha3.py

@@ -1,7 +1,28 @@
+from eth_hash.preimage import (
+    BasePreImage,
+)
 from sha3 import (
     keccak_256 as _keccak_256,
 )
 
 
 def keccak256(prehash: bytes) -> bytes:
     return _keccak_256(prehash).digest()
+
+
+class preimage(BasePreImage):
+    _hash = None
+
+    def __init__(self, prehash) -> None:
+        self._hash = _keccak_256(prehash)
+
+    def update(self, prehash) -> None:
+        return self._hash.update(prehash)
+
+    def digest(self) -> bytes:
+        return self._hash.digest()
+
+    def copy(self) -> 'preimage':
+        dup = preimage(b'')
+        dup._hash = self._hash.copy()
+        return dup

eth_hash/main.py

@@ -1,40 +1,30 @@
+from typing import (
+    Union,
+)
+
+from .preimage import (
+    BasePreImage,
+)
+
+
 class Keccak256:
     def __init__(self, backend):
         self.hasher = backend.keccak256
+        self.preimage = backend.preimage
 
         assert self.hasher(b'') == b"\xc5\xd2F\x01\x86\xf7#<\x92~}\xb2\xdc\xc7\x03\xc0\xe5\x00\xb6S\xca\x82';\x7b\xfa\xd8\x04]\x85\xa4p"  # noqa: E501
 
-    def __call__(self, preimage):
+    def __call__(self, preimage: Union[bytes, bytearray]) -> bytes:
         if not isinstance(preimage, (bytes, bytearray)):
             raise TypeError(
                 "Can only compute the hash of `bytes` or `bytearray` values, not %r" % preimage
             )
 
         return self.hasher(preimage)
 
-    def new(self, part):
-        if not isinstance(part, (bytes, bytearray)):
-            raise TypeError(
-                "Can only compute the hash of `bytes` or `bytearray` values, not %r" % part
-            )
-        return PreImage(part, keccak_fn=self.hasher)
-
-
-class PreImage:
-    def __init__(self, *parts, keccak_fn):
-        self.preimage_parts = list(parts)
-        self.keccak_fn = keccak_fn
-
-    def update(self, part):
-        if not isinstance(part, (bytes, bytearray)):
+    def new(self, preimage: (Union[bytes, bytearray])) -> BasePreImage:
+        if not isinstance(preimage, (bytes, bytearray)):
             raise TypeError(
-                "Can only compute the hash of `bytes` or `bytearray` values, not %r" % part
+                "Can only compute the hash of `bytes` or `bytearray` values, not %r" % preimage
             )
-        self.preimage_parts.append(part)
-
-    def digest(self):
-        return self.keccak_fn(b''.join(self.preimage_parts))
-
-    def copy(self):
-        digest = type(self)(*self.preimage_parts, keccak_fn=self.keccak_fn)
-        return digest
+        return self.preimage(preimage)

eth_hash/preimage.py

@@ -0,0 +1,22 @@
+from abc import (
+    ABCMeta,
+    abstractmethod,
+)
+
+
+class BasePreImage(metaclass=ABCMeta):
+    @abstractmethod
+    def __init__(self, value: bytes) -> None:
+        pass
+
+    @abstractmethod
+    def update(self, value: bytes) -> None:
+        pass
+
+    @abstractmethod
+    def digest(self) -> bytes:
+        return self.keccak_fn(b''.join(self.preimage_parts))
+
+    @abstractmethod
+    def copy(self) -> 'PreImage':
+        pass