linear-storage-poc.py

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()