fast_rlp.py

import trie
import rlp
import db
import time

trie.rlp_encode = rlp.codec.encode_raw


def int_to_big_endian(value):
    cs = []
    while value > 0:
        cs.append(chr(value % 256))
        value /= 256
    s = ''.join(reversed(cs))
    return s


def encode_optimized(item):
    """RLP encode (a nested sequence of) :class:`Atomic`s."""
    if isinstance(item, bytes):
        if len(item) == 1 and ord(item) < 128:
            return item
        prefix = length_prefix(len(item), 128)
    else:
        item = b''.join(map(encode_optimized, item))
        prefix = length_prefix(len(item), 192)

    return prefix + item


def length_prefix(length, offset):
    """Construct the prefix to lists or strings denoting their length.

    :param length: the length of the item in bytes
    :param offset: ``0x80`` when encoding raw bytes, ``0xc0`` when encoding a
                   list
    """
    if length < 56:
        return chr(offset + length)
    else:
        length_string = int_to_big_endian(length)
        return chr(offset + 56 - 1 + len(length_string)) + length_string


def run():
    st = time.time()
    x = trie.Trie(db.EphemDB())
    for i in range(10000):
        x.update(str(i), str(i**3))
    print 'elapsed', time.time() - st
    return x.root_hash


def main():
    trie.rlp_encode = rlp.encode
    print 'trie.rlp_encode = rlp.encode'
    r = run()
    trie.rlp_encode = rlp.codec.encode_raw
    print 'trie.rlp_encode = rlp.codec.encode_raw'
    r2 = run()
    assert r == r2
    trie.rlp_encode = encode_optimized
    print 'trie.rlp_encode = encode_optimized'
    r3 = run()
    assert r == r3

if __name__ == '__main__':
    main()