{-# LANGUAGE ParallelListComp #-}

module Language.Bitcoin.Script.Descriptors.Checksum (
  validDescriptorChecksum,
  descriptorChecksum,
) where

import Data.Bifunctor (first)
import Data.Bits (Bits (shiftL, shiftR, testBit, xor, (.&.)))
import Data.Char (ord)
import Data.Foldable (foldl')
import Data.IntMap (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as Vector

{- | Test whether the textual representation of an output descriptor has the
 given checksum.
-}
validDescriptorChecksum :: Text -> Text -> Bool
validDescriptorChecksum :: Text -> Text -> Bool
validDescriptorChecksum Text
desc Text
checksum =
  case forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Charset -> Char -> Maybe Word
charsetFind Charset
checksumCharset) (Text -> String
Text.unpack Text
checksum) of
    Maybe [Word]
Nothing -> Bool
False
    Just [Word]
checkSymbols ->
      Word
1 forall a. Eq a => a -> a -> Bool
== [Word] -> Word
polymodChecksum (Text -> [Word]
expandChecksum Text
desc forall a. Semigroup a => a -> a -> a
<> [Word]
checkSymbols)

{- | Compute the checksum of the textual representation of an output descriptor
 if possible.
-}
descriptorChecksum :: Text -> Maybe Text
descriptorChecksum :: Text -> Maybe Text
descriptorChecksum Text
desc = String -> Text
Text.pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA [Maybe Char]
checksumChars
 where
  checksumChars :: [Maybe Char]
checksumChars = [Charset
checksumCharset Charset -> Word -> Maybe Char
`charsetGet` Int -> Word
charsetIndex Int
i | Int
i <- [Int
0 .. Int
7]]
  charsetIndex :: Int -> Word
charsetIndex Int
i = (Word
checksum forall a. Bits a => a -> Int -> a
`shiftR` (Int
5 forall a. Num a => a -> a -> a
* (Int
7 forall a. Num a => a -> a -> a
- Int
i))) forall a. Bits a => a -> a -> a
.&. Word
31
  symbols :: [Word]
symbols = Text -> [Word]
expandChecksum Text
desc forall a. Semigroup a => a -> a -> a
<> forall a. Int -> a -> [a]
replicate Int
8 Word
0
  checksum :: Word
checksum = Word
1 forall a. Bits a => a -> a -> a
`xor` [Word] -> Word
polymodChecksum [Word]
symbols

expandChecksum :: Text -> [Word]
expandChecksum :: Text -> [Word]
expandChecksum =
  forall a. [a] -> [a]
reverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. Num a => ([a], [a]) -> [a]
end
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
      ( \([Word]
gs, [Word]
s) Word
v -> case (Word
v forall a. Bits a => a -> Int -> a
`shiftR` Int
5 forall a. a -> [a] -> [a]
: [Word]
gs, Word
v forall a. Bits a => a -> a -> a
.&. Word
31 forall a. a -> [a] -> [a]
: [Word]
s) of
          ([Word
g2, Word
g1, Word
g0], [Word]
s') -> ([], Word
9 forall a. Num a => a -> a -> a
* Word
g0 forall a. Num a => a -> a -> a
+ Word
3 forall a. Num a => a -> a -> a
* Word
g1 forall a. Num a => a -> a -> a
+ Word
g2 forall a. a -> [a] -> [a]
: [Word]
s')
          ([Word], [Word])
x -> ([Word], [Word])
x
      )
      forall a. Monoid a => a
mempty
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a -> a
fromMaybe []
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Charset -> Char -> Maybe Word
charsetFind Charset
inputCharset)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
Text.unpack
 where
  end :: ([a], [a]) -> [a]
end ([a
g0], [a]
s) = a
g0 forall a. a -> [a] -> [a]
: [a]
s
  end ([a
g1, a
g0], [a]
s) = a
3 forall a. Num a => a -> a -> a
* a
g0 forall a. Num a => a -> a -> a
+ a
g1 forall a. a -> [a] -> [a]
: [a]
s
  end ([a]
_, [a]
s) = [a]
s

polymodChecksum :: [Word] -> Word
polymodChecksum :: [Word] -> Word
polymodChecksum =
  forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
    ( \Word
chk Word
value ->
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          forall a. Bits a => a -> a -> a
xor
          ((Word
chk forall a. Bits a => a -> a -> a
.&. Word
0x7ffffffff) forall a. Bits a => a -> Int -> a
`shiftL` Int
5 forall a. Bits a => a -> a -> a
`xor` Word
value)
          [if Word
chk forall a. Bits a => a -> Int -> Bool
`testBit` Int
i then Word
g else Word
0 | Int
i <- [Int
35 ..] | Word
g <- [Word]
generator]
    )
    Word
1
 where
  generator :: [Word]
generator =
    [ Word
0xf5dee51989
    , Word
0xa9fdca3312
    , Word
0x1bab10e32d
    , Word
0x3706b1677a
    , Word
0x644d626ffd
    ]

data Charset = Charset
  { Charset -> IntMap Word
charToIndex :: IntMap Word
  , Charset -> Vector Char
indexToChar :: Vector Char
  }

charsetFromString :: String -> Charset
charsetFromString :: String -> Charset
charsetFromString String
s =
  let xs :: [(Char, Word)]
xs = [(Char
c, Word
i) | Char
c <- String
s | Word
i <- [Word
0 ..]]
   in Charset
        { charToIndex :: IntMap Word
charToIndex = forall a. [(Int, a)] -> IntMap a
IntMap.fromList forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Char -> Int
ord forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Char, Word)]
xs
        , indexToChar :: Vector Char
indexToChar = forall a. Unbox a => [a] -> Vector a
Vector.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Char, Word)]
xs
        }

charsetFind :: Charset -> Char -> Maybe Word
charsetFind :: Charset -> Char -> Maybe Word
charsetFind Charset
set Char
c = forall a. Int -> IntMap a -> Maybe a
IntMap.lookup (Char -> Int
ord Char
c) forall a b. (a -> b) -> a -> b
$ Charset -> IntMap Word
charToIndex Charset
set

charsetGet :: Charset -> Word -> Maybe Char
charsetGet :: Charset -> Word -> Maybe Char
charsetGet Charset
set Word
i = Charset -> Vector Char
indexToChar Charset
set forall a. Unbox a => Vector a -> Int -> Maybe a
Vector.!? forall a. Num a => Integer -> a
fromInteger (forall a. Integral a => a -> Integer
toInteger Word
i)

inputCharset :: Charset
inputCharset :: Charset
inputCharset = String -> Charset
charsetFromString String
"0123456789()[],'/*abcdefgh@:$%{}IJKLMNOPQRSTUVWXYZ&+-.;<=>?!^_|~ijklmnopqrstuvwxyzABCDEFGH`#\"\\ "

checksumCharset :: Charset
checksumCharset :: Charset
checksumCharset = String -> Charset
charsetFromString String
"qpzry9x8gf2tvdw0s3jn54khce6mua7l"