{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Module: Data.ByteString.Base58
-- Copyright: (c) 2024 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- base58 encoding and decoding of strict bytestrings.

module Data.ByteString.Base58 (
    encode
  , decode
  ) where

import Control.Monad (guard)
import qualified Data.Bits as B
import Data.Bits ((.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BU

fi :: (Integral a, Num b) => a -> b
fi :: forall a b. (Integral a, Num b) => a -> b
fi = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE fi #-}

-- | Encode a base256 'ByteString' as base58.
--
--   >>> encode "hello world"
--   "StV1DL6CwTryKyV"
encode :: BS.ByteString -> BS.ByteString
encode :: ByteString -> ByteString
encode ByteString
bs = ByteString
ls ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Integer -> ByteString
unroll_base58 (ByteString -> Integer
roll_base256 ByteString
bs) where
  ls :: ByteString
ls = ByteString -> ByteString
leading_ones ByteString
bs

-- | Decode a base58 'ByteString' to base256.
--
--   Invalid inputs will produce 'Nothing'.
--
--   >>> decode "StV1DL6CwTryKyV"
--   Just "hello world"
--   >>> decode "StV1DL0CwTryKyV" -- s/6/0
--   Nothing
decode :: BS.ByteString -> Maybe BS.ByteString
decode :: ByteString -> Maybe ByteString
decode ByteString
bs = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ByteString -> Bool
verify_base58 ByteString
bs)
  let ls :: ByteString
ls = ByteString -> ByteString
leading_zeros ByteString
bs
  ByteString -> Maybe ByteString
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
ls ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Integer -> ByteString
unroll_base256 (ByteString -> Integer
roll_base58 ByteString
bs)

verify_base58 :: BS.ByteString -> Bool
verify_base58 :: ByteString -> Bool
verify_base58 ByteString
bs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
bs of
  Maybe (Word8, ByteString)
Nothing -> Bool
True
  Just (Word8
h, ByteString
t)
    | Word8 -> ByteString -> Bool
BS.elem Word8
h ByteString
base58_charset -> ByteString -> Bool
verify_base58 ByteString
t
    | Bool
otherwise -> Bool
False

base58_charset :: BS.ByteString
base58_charset :: ByteString
base58_charset = ByteString
"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"

-- produce leading ones from leading zeros
leading_ones :: BS.ByteString -> BS.ByteString
leading_ones :: ByteString -> ByteString
leading_ones = ByteString -> ByteString -> ByteString
go ByteString
forall a. Monoid a => a
mempty where
  go :: ByteString -> ByteString -> ByteString
go ByteString
acc ByteString
bs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
bs of
    Maybe (Word8, ByteString)
Nothing -> ByteString
acc
    Just (Word8
h, ByteString
t)
      | Word8
h Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0 -> ByteString -> ByteString -> ByteString
go (Word8 -> ByteString -> ByteString
BS.cons Word8
0x31 ByteString
acc) ByteString
t
      | Bool
otherwise -> ByteString
acc

-- produce leading zeros from leading ones
leading_zeros :: BS.ByteString -> BS.ByteString
leading_zeros :: ByteString -> ByteString
leading_zeros = ByteString -> ByteString -> ByteString
go ByteString
forall a. Monoid a => a
mempty where
  go :: ByteString -> ByteString -> ByteString
go ByteString
acc ByteString
bs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
bs of
    Maybe (Word8, ByteString)
Nothing -> ByteString
acc
    Just (Word8
h, ByteString
t)
      | Word8
h Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x31 -> ByteString -> ByteString -> ByteString
go (Word8 -> ByteString -> ByteString
BS.cons Word8
0x00 ByteString
acc) ByteString
t
      | Bool
otherwise -> ByteString
acc

-- to base256
unroll_base256 :: Integer -> BS.ByteString
unroll_base256 :: Integer -> ByteString
unroll_base256 = ByteString -> ByteString
BS.reverse (ByteString -> ByteString)
-> (Integer -> ByteString) -> Integer -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Maybe (Word8, Integer)) -> Integer -> ByteString
forall a. (a -> Maybe (Word8, a)) -> a -> ByteString
BS.unfoldr Integer -> Maybe (Word8, Integer)
forall {b} {a}. (Integral b, Num a) => b -> Maybe (a, b)
coalg where
  coalg :: b -> Maybe (a, b)
coalg b
a
    | b
a b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0 = Maybe (a, b)
forall a. Maybe a
Nothing
    | Bool
otherwise = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just ((a, b) -> Maybe (a, b)) -> (a, b) -> Maybe (a, b)
forall a b. (a -> b) -> a -> b
$
        let (b
b, b
c) = b -> b -> (b, b)
forall a. Integral a => a -> a -> (a, a)
quotRem b
a b
256
        in  (b -> a
forall a b. (Integral a, Num b) => a -> b
fi b
c, b
b)

-- from base256
roll_base256 :: BS.ByteString -> Integer
roll_base256 :: ByteString -> Integer
roll_base256 = (Integer -> Word8 -> Integer) -> Integer -> ByteString -> Integer
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' Integer -> Word8 -> Integer
forall {a} {a}. (Bits a, Integral a, Num a) => a -> a -> a
alg Integer
0 where
  alg :: a -> a -> a
alg !a
a !a
b = a
a a -> Int -> a
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
8 a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fi a
b

-- to base58
unroll_base58 :: Integer -> BS.ByteString
unroll_base58 :: Integer -> ByteString
unroll_base58 = ByteString -> ByteString
BS.reverse (ByteString -> ByteString)
-> (Integer -> ByteString) -> Integer -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Maybe (Word8, Integer)) -> Integer -> ByteString
forall a. (a -> Maybe (Word8, a)) -> a -> ByteString
BS.unfoldr Integer -> Maybe (Word8, Integer)
forall {b}. Integral b => b -> Maybe (Word8, b)
coalg where
  coalg :: b -> Maybe (Word8, b)
coalg b
a
    | b
a b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0 = Maybe (Word8, b)
forall a. Maybe a
Nothing
    | Bool
otherwise = (Word8, b) -> Maybe (Word8, b)
forall a. a -> Maybe a
Just ((Word8, b) -> Maybe (Word8, b)) -> (Word8, b) -> Maybe (Word8, b)
forall a b. (a -> b) -> a -> b
$
        let (b
b, b
c) = b -> b -> (b, b)
forall a. Integral a => a -> a -> (a, a)
quotRem b
a b
58
        in  (ByteString -> Int -> Word8
BU.unsafeIndex ByteString
base58_charset (b -> Int
forall a b. (Integral a, Num b) => a -> b
fi b
c), b
b)

-- from base58
roll_base58 :: BS.ByteString -> Integer
roll_base58 :: ByteString -> Integer
roll_base58 ByteString
bs = (Integer -> Word8 -> Integer) -> Integer -> ByteString -> Integer
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' Integer -> Word8 -> Integer
forall {a}. Num a => a -> Word8 -> a
alg Integer
0 ByteString
bs where
  alg :: a -> Word8 -> a
alg !a
b !Word8
a = case Word8 -> ByteString -> Maybe Int
BS.elemIndex Word8
a ByteString
base58_charset of
    Just Int
w -> a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
58 a -> a -> a
forall a. Num a => a -> a -> a
+ Int -> a
forall a b. (Integral a, Num b) => a -> b
fi Int
w
    Maybe Int
Nothing ->
      [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"ppad-base58 (roll_base58): not a base58-encoded bytestring"