Skip to content

Instantly share code, notes, and snippets.

@karmacoma-eth
Created December 12, 2023 02:25
Show Gist options
  • Save karmacoma-eth/de23e3d243e4d7669dfa7a0ec78cecb3 to your computer and use it in GitHub Desktop.
Save karmacoma-eth/de23e3d243e4d7669dfa7a0ec78cecb3 to your computer and use it in GitHub Desktop.
linear-storage-poc.py
from dataclasses import dataclass
from typing import Tuple
bytes32 = int
address = int
@dataclass
class Context:
context_address: address
# the storage pointer points to a "virtual head" in storage
storage_pointer: int
@dataclass(frozen=True)
class StorageUpdate:
'''Represents a storage update for a contract'''
contract_address: address
key: str # for readability of the demo, otherwise would be a bytes32
value: bytes32
def __str__(self):
return f"@{hex(self.contract_address)} {self.key}={self.value}"
class Storage:
'''Represents the storage values for all contracts in the environment'''
def __init__(self, parent: "Storage" = None):
self.updates = list()
self.parent = parent
self.actual_head = 0
self.virtual_head = parent.virtual_head if parent else 0
self.forked = False
# O(1)
# potential optimization: we could aggregate multiple updates of the same slot into a single update
def set(self, context: Context, key: bytes32, value: bytes32):
if self.forked:
raise Exception("Forked storage cannot be modified")
# we have room available in the list, overwrite whatever was there
update = StorageUpdate(context.context_address, key, value)
if self.actual_head < len(self.updates):
self.updates[self.actual_head] = update
else:
self.updates.append(update)
# advance both heads
self.actual_head += 1
self.virtual_head += 1
# O(n), doing a reverse linear scan
def get(self, addr: address, key: bytes32) -> bytes32:
for update in reversed(self.updates[:self.actual_head]):
if update.contract_address == addr and update.key == key:
return update.value
if self.parent is not None:
return self.parent.get(addr, key)
return 0
# O(1) in the happy path
# may incur copies if the storage pointer belongs to a parent
def revert(self, context: Context) -> "Storage":
'''Reverts all storage updates for the given context'''
if self.forked:
raise Exception("Forked storage cannot be modified")
# check that we are well formed
assert self.virtual_head >= self.actual_head
assert len(self.updates) >= self.actual_head
# check that the request makes sense
assert context.storage_pointer >= 0
assert context.storage_pointer <= self.virtual_head
# check if the storage pointer belongs to this storage instance
virtual_base = self.virtual_head - self.actual_head
if context.storage_pointer >= virtual_base:
# number of elements to drop
n = self.virtual_head - context.storage_pointer
# "drop" the elements by just moving the head backwards
# this is O(1), but we leave unused elements in the updates list
self.actual_head -= n
self.virtual_head = context.storage_pointer
return self
# more complex path: the storage pointer belongs to a parent
else:
raise NotImplementedError("revert into parent")
# O(1), does not incur copies
def fork(self) -> Tuple["Storage", "Storage"]:
'''Forks the current storage'''
self.forked = True
return Storage(self), Storage(self)
def __str__(self):
result = "Storage:\n"
result += f" updates:\n"
for update in reversed(self.updates):
result += f" {update}\n"
result += f" head: {self.actual_head}\n"
result += f" vhead: {self.virtual_head}\n"
result += f" forked: {self.forked}\n"
result += f" parent: {str(self.parent)}"
return result
def current_pointer(self) -> int:
return self.virtual_head
def demo():
storage = Storage()
context = Context(0xcafe, storage.current_pointer())
# when write a value
storage.set(context, 'a', 1)
storage.set(context, 'b', 2)
# then we can read it back
assert storage.get(0xcafe, 'a') == 1
assert storage.get(0xcafe, 'b') == 2
# when we overwrite a value
storage.set(context, 'a', 3)
# print(storage)
# Storage:
# updates:
# @0xcafe a=3 <--- this will be returned first
# @0xcafe b=2
# @0xcafe a=1
# head: 2
# vhead: 2
# forked: False
# parent: None
# then we can read it back
assert storage.get(0xcafe, 'a') == 3
# when we open a new context:
subcontext = Context(0xdead, storage.current_pointer())
# then the values don't clash
storage.set(subcontext, 'a', 4)
assert storage.get(0xdead, 'a') == 4
assert storage.get(0xcafe, 'a') == 3
# when we revert the subcontext
storage.revert(subcontext)
# then the values from the subcontext are gone
assert storage.get(0xdead, 'a') == 0
# Storage:
# updates:
# @0xdead a=4 <--- still in the list, but the head was moved back
# @0xcafe a=3
# @0xcafe b=2
# @0xcafe a=1
# head: 3
# vhead: 3
# forked: False
# parent: None
# when we write a value after reverting
storage.set(context, 'a', 4)
# then the value is overwritten in the updates list
assert len(storage.updates) == 4
assert storage.updates[-1] == StorageUpdate(0xcafe, 'a', 4)
# when we fork the storage
storage_fork1, storage_fork2 = storage.fork()
# then we can no longer modify the original storage
try:
storage.set(context, 'c', 5)
assert False
except Exception:
pass
# but we can modify the forks independently
storage_fork1.set(context, 'b', 5)
storage_fork2.set(context, 'b', 6)
# and we get the values we would expect
assert storage_fork1.get(0xcafe, 'b') == 5
assert storage_fork2.get(0xcafe, 'b') == 6
# print(storage_fork1)
# Storage:
# updates:
# @0xcafe b=5
# head: 1 <--- head is 1 in the fork
# vhead: 4 <--- but vhead is 4, because we have 4 logical updates in the environment
# forked: False
# parent: Storage: <--- the parent is the original storage, now immutable
# updates:
# @0xcafe a=4
# @0xcafe a=3
# @0xcafe b=2
# @0xcafe a=1
# head: 3
# vhead: 3
# forked: True
# parent: None
# storage_fork2 is the same, but with b=6
# setup multiple embedded contexts
value_cafe_a_before = storage_fork1.get(0xcafe, 'a')
subcontext1 = Context(0xdead, storage_fork1.current_pointer())
storage_fork1.set(subcontext1, 'a', 7)
subcontext2 = Context(0xdead, storage_fork1.current_pointer()) # this ones stays empty
subcontext3 = Context(0xcafe, storage_fork1.current_pointer())
storage_fork1.set(subcontext3, 'a', 8)
assert storage_fork1.get(0xdead, 'a') == 7
assert storage_fork1.get(0xcafe, 'a') == 8
# when we revert multiple contexts
storage_fork1.revert(subcontext1)
# then all the updates are gone
assert storage_fork1.get(0xdead, 'a') == 0
assert storage_fork1.get(0xcafe, 'a') == value_cafe_a_before
# print(storage_fork1)
# Storage:
# updates:
# @0xcafe a=8
# @0xdead a=7
# @0xcafe b=5
# head: 1
# vhead: 5
# forked: False
# parent: Storage:
# updates:
# @0xcafe a=4
# @0xcafe a=3
# @0xcafe b=2
# @0xcafe a=1
# head: 4
# vhead: 4
# forked: True
# parent: None
demo()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment