1 files changed, 21 insertions, 4 deletions

diff --git a/Base58.hs b/Base58.hs
index de35b01..3c1a113 100644
--- a/Base58.hs
+++ b/Base58.hs
@@ -1,6 +1,12 @@
+{-# LANGUAGE CPP #-}
 module Base58 where
 
+#if !defined(VERSION_cryptonite)
 import qualified Crypto.Hash.SHA256 as SHA256
+#else
+import Crypto.Hash
+import Data.ByteArray (convert)
+#endif
 import qualified Data.ByteString      as S
 import Data.Maybe
 import Data.List
@@ -28,7 +34,12 @@ base58_decode str = do
     
     let (rcksum,rpayload) = splitAt 4 $ rbytes
         a_payload = reverse rpayload
+#if !defined(VERSION_cryptonite)
         hash_result = S.take 4 . SHA256.hash . SHA256.hash . S.pack $ a_payload
+#else
+        hash_result = S.take 4 . convert $ digest
+          where digest = hash (S.pack a_payload) :: Digest SHA256
+#endif
         expected_hash = S.pack $ reverse rcksum
         (network_id,payload) = splitAt 1 a_payload
     
@@ -37,11 +48,17 @@ base58_decode str = do
     return (network_id,payload)
 
 base58_encode :: S.ByteString -> String
-base58_encode hash = replicate zcount '1' ++ map (base58chars !!) (reverse rdigits)
+base58_encode hsh = replicate zcount '1' ++ map (base58chars !!) (reverse rdigits)
   where
-    zcount = S.length . S.takeWhile (==0) $ hash
-    cksum = S.take 4 . SHA256.hash . SHA256.hash $ hash
-    n = foldl' (\a b->a*256+b) 0 . map asInteger $ concatMap S.unpack [hash, cksum]
+    zcount = S.length . S.takeWhile (==0) $ hsh
+#if !defined(VERSION_cryptonite)
+    cksum = S.take 4 . SHA256.hash . SHA256.hash $ hsh
+#else
+    cksum = S.take 4 (convert digest2 :: S.ByteString)
+      where digest2 = hash ( convert digest1 :: S.ByteString) :: Digest SHA256
+            digest1 = hash hsh :: Digest SHA256
+#endif
+    n = foldl' (\a b->a*256+b) 0 . map asInteger $ concatMap S.unpack [hsh, cksum]
     asInteger x = fromIntegral x :: Integer
     rdigits = unfoldr getdigit n
      where