Skip to content

Commit

Permalink
research: sparse merkle tree optimizations proposed by vitalik
Browse files Browse the repository at this point in the history
  • Loading branch information
kaxxa123 committed Sep 11, 2024
1 parent 1d0edd2 commit 9b1b32f
Show file tree
Hide file tree
Showing 6 changed files with 610 additions and 0 deletions.
162 changes: 162 additions & 0 deletions vitalik_merkle_optimizations/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST

# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
.pybuilder/
target/

# Jupyter Notebook
.ipynb_checkpoints

# IPython
profile_default/
ipython_config.py

# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version

# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock

# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock

# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
.pdm.toml
.pdm-python
.pdm-build/

# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/

# Celery stuff
celerybeat-schedule
celerybeat.pid

# SageMath parsed files
*.sage.py

# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/
.dmypy.json
dmypy.json

# Pyre type checker
.pyre/

# pytype static type analyzer
.pytype/

# Cython debug symbols
cython_debug/

# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
45 changes: 45 additions & 0 deletions vitalik_merkle_optimizations/bintrie_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
import new_bintrie as t1
import new_bintrie_optimized as t2
import new_bintrie_hex as t3
import time
import binascii

keys = [t1.sha3(bytes([i // 256, i % 256])) for i in range(10000)]

d = t1.EphemDB()
r = t1.new_tree(d)
a = time.time()
for k in keys[:1000]:
r = t1.update(d, r, k, k)
print("Naive bintree time to update: %.4f" % (time.time() - a))
print("Root: %s" % binascii.hexlify(r))

d = t2.EphemDB()
r = t2.new_tree(d)
a = time.time()
for k in keys[:1000]:
r = t2.update(d, r, k, k)
print("DB-optimized bintree time to update: %.4f" % (time.time() - a))
print("Root: %s" % binascii.hexlify(r))
print("Writes: %d, reads: %d" % (d.writes, d.reads))
d.reads = 0
for k in keys[:500]:
assert t2.get(d, r, k) == k
for k in keys[-500:]:
assert t2.get(d, r, k) == b'\x00' * 32
print("Reads: %d" % d.reads)

d = t3.EphemDB()
r = t3.new_tree(d)
a = time.time()
for k in keys[:1000]:
r = t3.update(d, r, k, k)
print("DB-optimized bintree time to update: %.4f" % (time.time() - a))
print("Root: %s" % binascii.hexlify(r))
print("Writes: %d, reads: %d" % (d.writes, d.reads))
d.reads = 0
for k in keys[:500]:
assert t3.get(d, r, k) == k
for k in keys[-500:]:
assert t3.get(d, r, k) == b'\x00' * 32
print("Reads: %d" % d.reads)
125 changes: 125 additions & 0 deletions vitalik_merkle_optimizations/new_bintrie.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
from ethereum.utils import sha3, encode_hex

class EphemDB():
def __init__(self, kv=None):
self.kv = kv or {}

def get(self, k):
return self.kv.get(k, None)

def put(self, k, v):
self.kv[k] = v

def delete(self, k):
del self.kv[k]

zerohashes = [b'\x00' * 32]
for i in range(256):
zerohashes.insert(0, sha3(zerohashes[0] + zerohashes[0]))

def new_tree(db):
h = b'\x00' * 32
for i in range(256):
newh = sha3(h + h)
db.put(newh, h + h)
h = newh
return h

def key_to_path(k):
o = 0
for c in k:
o = (o << 8) + c
return o

def descend(db, root, *path):
v = root
for p in path:
if p:
v = db.get(v)[32:]
else:
v = db.get(v)[:32]
return v

def get(db, root, key):
v = root
path = key_to_path(key)
for i in range(256):
if (path >> 255) & 1:
v = db.get(v)[32:]
else:
v = db.get(v)[:32]
path <<= 1
return v

def update(db, root, key, value):
v = root
path = path2 = key_to_path(key)
sidenodes = []
for i in range(256):
if (path >> 255) & 1:
sidenodes.append(db.get(v)[:32])
v = db.get(v)[32:]
else:
sidenodes.append(db.get(v)[32:])
v = db.get(v)[:32]
path <<= 1
v = value
for i in range(256):
if (path2 & 1):
newv = sha3(sidenodes[-1] + v)
db.put(newv, sidenodes[-1] + v)
else:
newv = sha3(v + sidenodes[-1])
db.put(newv, v + sidenodes[-1])
path2 >>= 1
v = newv
sidenodes.pop()
return v

def make_merkle_proof(db, root, key):
v = root
path = key_to_path(key)
sidenodes = []
for i in range(256):
if (path >> 255) & 1:
sidenodes.append(db.get(v)[:32])
v = db.get(v)[32:]
else:
sidenodes.append(db.get(v)[32:])
v = db.get(v)[:32]
path <<= 1
return sidenodes

def verify_proof(proof, root, key, value):
path = key_to_path(key)
v = value
for i in range(256):
if (path & 1):
newv = sha3(proof[-1-i] + v)
else:
newv = sha3(v + proof[-1-i])
path >>= 1
v = newv
return root == v

def compress_proof(proof):
bits = bytearray(32)
oproof = b''
for i, p in enumerate(proof):
if p == zerohashes[i+1]:
bits[i // 8] ^= 1 << i % 8
else:
oproof += p
return bytes(bits) + oproof

def decompress_proof(oproof):
proof = []
bits = bytearray(oproof[:32])
pos = 32
for i in range(256):
if bits[i // 8] & (1 << (i % 8)):
proof.append(zerohashes[i+1])
else:
proof.append(oproof[pos: pos + 32])
pos += 32
return proof
Loading

0 comments on commit 9b1b32f

Please sign in to comment.