packages feed

bitcoin-compact-filters-0.1.0.2: src/Bitcoin/CompactFilter.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_GHC -Wno-type-defaults #-}

module Bitcoin.CompactFilter (
    BlockFilter,
    blockFilter,
    BlockFilterHeader,
    blockFilterHeader,
    filterHeaderToHex,
    filterHeaderFromHex,
    genesisHeader,
    filterContents,
    encodeFilter,
    isMember,
) where

import Control.Monad (replicateM, (>=>))
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State.Strict (StateT, evalStateT)
import qualified Control.Monad.Trans.State.Strict as St
import Data.Bits (shiftL, shiftR, testBit)
import Data.Bitstream (Bitstream, Right)
import qualified Data.Bitstream as BiS
import Data.Bool (bool)
import Data.ByteArray.Hash (
    SipHash (..),
    SipKey (..),
    sipHash,
 )
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Foldable (foldl')
import Data.List (sort)
import Data.Serialize (
    Get,
    Serialize (..),
    decode,
    encode,
    getWord16le,
    getWord32le,
    getWord64le,
    getWord8,
    putByteString,
    putWord16le,
    putWord32le,
    putWord64le,
    putWord8,
    runGet,
 )
import Data.Text (Text)
import Data.Word (Word64, Word8)
import Haskoin.Block (
    Block,
    BlockHash,
    blockHeader,
    blockTxns,
    headerHash,
 )
import Haskoin.Crypto (Hash256, doubleSHA256)
import Haskoin.Script (Script (..), ScriptOp (..))
import Haskoin.Transaction (scriptOutput, txOut)
import Haskoin.Util (decodeHex, encodeHex)

-- | SIP hash parameter
paramP :: Int
paramP = 19

-- | SIP hash parameter
paramM :: Word64
paramM = 784931

-- | Hashes of scripts in the block
newtype BlockFilter = BlockFilter
    { -- | Get the list of hashes in increasing order
      blockFilter :: [Word64]
    }
    deriving (Eq, Show)

-- | Number of elements in the filter
blockFilterSize :: BlockFilter -> CompactSize
blockFilterSize = CompactSize . length . blockFilter

newtype BlockFilterHeader = BlockFilterHeader {getBFHeader :: Hash256}
    deriving (Eq, Show, Ord, Serialize)

filterHeaderToHex :: BlockFilterHeader -> Text
filterHeaderToHex = encodeHex . BS.reverse . encode . getBFHeader

filterHeaderFromHex :: Text -> Either String BlockFilterHeader
filterHeaderFromHex = maybe (Left "Invalid hex") Right . decodeHex >=> decode . BS.reverse

genesisHeader :: BlockFilterHeader
Right genesisHeader = BlockFilterHeader <$> decode (BS.replicate 32 0x0)

-- | Calculate the header for the block filter
blockFilterHeader ::
    -- | previous header
    BlockFilterHeader ->
    -- | current filter
    BlockFilter ->
    BlockFilterHeader
blockFilterHeader prev bf =
    BlockFilterHeader . doubleSHA256 $
        (encode . doubleSHA256) (encode bf) <> encode (getBFHeader prev)

instance Serialize BlockFilter where
    put bf = put (blockFilterSize bf) >> putByteString (constructGCS paramP $ blockFilter bf)
    get = get >>= fmap BlockFilter . getGolombRiceSet paramP . unCompactSize

-- | Calculate the list of scripts which belong in the BIP158 block filter
filterContents ::
    -- | previous output scripts spent in this block
    [ByteString] ->
    Block ->
    [ByteString]
filterContents prev b = filter scriptFilter prev <> these
  where
    these = filter scriptFilter . fmap scriptOutput $ blockTxns b >>= txOut
    scriptFilter scr = not (BS.null scr) && contentFilter scr

    contentFilter bs = case decode bs of
        Right (Script (OP_RETURN : _)) -> False
        _ -> True

-- | Construct a BIP158 filter from a block
encodeFilter ::
    -- | output scripts spent in this block
    [ByteString] ->
    Block ->
    BlockFilter
encodeFilter os b = BlockFilter s
  where
    h = headerHash $ blockHeader b
    bs = toSet $ filterContents os b
    s = hashedSetConstruct (sipKey h) paramM (length bs) bs

-- | Test membership.  The test succeeds if /any/ of the scripts matches the block filter.
isMember ::
    BlockHash ->
    -- | Scripts we want to match against the filter
    [ByteString] ->
    BlockFilter ->
    Bool
isMember h bs (BlockFilter bf) = orderedScan hs bf
  where
    k = sipKey h
    hs = hashedSetConstruct k paramM (length bf) bs

orderedScan :: [Word64] -> [Word64] -> Bool
orderedScan xs@(x : xs') hs@(h : hs')
    | x > h = orderedScan xs hs'
    | x < h = orderedScan xs' hs
    | otherwise = True
orderedScan _ _ = False

type GetBits = StateT [Bool] Get

getGolombRiceSet :: Int -> Int -> Get [Word64]
getGolombRiceSet p n = fmap unDiffs . evalBitstream $ replicateM n getEncoded
  where
    getEncoded = do
        q <- unaryPart
        r <- fromBits <$> getBits p
        return $ q * 2 ^ p + r

unaryPart :: GetBits Word64
unaryPart = go 0
  where
    go q = getBit >>= bool (return q) (go $ q + 1)
    getBit = head <$> getBits 1

evalBitstream :: GetBits a -> Get a
evalBitstream = (`evalStateT` mempty)

getBits :: Int -> GetBits [Bool]
getBits n = do
    bs <- St.get
    let l = length bs
        (q, r) = (n - l) `quotRem` 8

        combine m t = bs <> mconcat m <> t

        lastByte
            | r > 0 = getBs >>= takeSome r
            | otherwise = mempty <$ St.put mempty

        getBs = byteBits <$> lift getWord8
        takeSome m xs = take m xs <$ St.put (drop m xs)

    if n < l
        then takeSome n bs
        else combine <$> replicateM q getBs <*> lastByte

sipKey :: BlockHash -> SipKey
sipKey h = SipKey k1 k2
  where
    Right (k1, k2) = runGet word64Pair $ encode h
    word64Pair = (,) <$> getWord64le <*> getWord64le

hashToRange :: Word64 -> SipKey -> ByteString -> Word64
hashToRange f k bs = v
  where
    SipHash h = sipHash k bs
    v = remap (fromIntegral f) (fromIntegral h)

    remap :: Integer -> Integer -> Word64
    remap x y = fromIntegral $ (x * y) `shiftR` 64

hashedSetConstruct :: SipKey -> Word64 -> Int -> [ByteString] -> [Word64]
hashedSetConstruct k m n bs = toSet $ hashToRange f k <$> bs
  where
    f = fromIntegral n * m

toSet :: Ord a => [a] -> [a]
toSet = dedup . sort
  where
    dedup = \case
        (x0 : xs@(x1 : _))
            | x0 == x1 -> dedup xs
            | otherwise -> x0 : dedup xs
        xs -> xs

constructGCS ::
    -- | modulus
    Int ->
    -- | sorted list of input values
    [Word64] ->
    ByteString
constructGCS p =
    BiS.toByteString
        . foldMap (golombRiceEncode p)
        . diffs

diffs :: Num a => [a] -> [a]
diffs xs = zipWith (-) xs (0 : xs)

unDiffs :: Num a => [a] -> [a]
unDiffs (x : xs) = scanl (+) x xs
unDiffs [] = []

golombRiceEncode :: Int -> Word64 -> Bitstream Right
golombRiceEncode p v = x <> BiS.singleton False <> y
  where
    q = fromIntegral $ v `shiftR` p
    x = BiS.replicate q True
    y = BiS.fromNBits p v

fromBits :: Num a => [Bool] -> a
fromBits = foldl' onBit 0
  where
    onBit n b = 2 * n + bool 0 1 b

byteBits :: Word8 -> [Bool]
byteBits b = testBit b <$> reverse [0 .. 7]

newtype CompactSize = CompactSize {unCompactSize :: Int} deriving (Eq, Ord, Enum, Num, Real, Integral)

instance Serialize CompactSize where
    get = CompactSize <$> (getWord8 >>= getCS)
      where
        getCS s
            | s < 253 = return $ fromIntegral s
            | s == 253 = fromIntegral <$> getWord16le
            | s == 254 = fromIntegral <$> getWord32le
            | otherwise = fromIntegral <$> getWord64le

    put (CompactSize n)
        | n < 0 = error $ "Invalid CompactSize: " <> show n
        | n < 253 = putWord8 (fromIntegral n)
        | n < bound16 = putWord8 253 >> putWord16le (fromIntegral n)
        | n < bound32 = putWord8 254 >> putWord32le (fromIntegral n)
        | otherwise = putWord8 255 >> putWord64le (fromIntegral n)
      where
        bound16 = 1 `shiftL` 16
        bound32 = 1 `shiftL` 32