conflict-set/conflict_set.py

import ctypes
import enum
import os

from typing import Optional


class _Key(ctypes.Structure):
    _fields_ = [("p", ctypes.POINTER(ctypes.c_ubyte)), ("len", ctypes.c_int)]


class ReadRange(ctypes.Structure):
    _fields_ = [
        ("begin", _Key),
        ("end", _Key),
        ("readVersion", ctypes.c_int64),
    ]


class WriteRange(ctypes.Structure):
    _fields_ = [("begin", _Key), ("end", _Key)]


class Result(enum.Enum):
    COMMIT = 0
    CONFLICT = 1
    TOO_OLD = 2


def _make_key(buf: bytes) -> tuple[_Key, bytearray]:
    """Create a _Key and a backing bytearray that must be kept alive."""
    backing = bytearray(buf)
    array = (ctypes.c_ubyte * len(backing)).from_buffer(backing)
    return _Key(array, len(array)), backing


def write(begin: bytes, end: Optional[bytes] = None) -> WriteRange:
    begin_key, begin_buf = _make_key(begin)
    if end is None:
        end_key = _Key((ctypes.c_ubyte * 0)(), 0)
        end_buf = None
    else:
        end_key, end_buf = _make_key(end)
    result = WriteRange(begin_key, end_key)
    result._begin_buf = begin_buf
    result._end_buf = end_buf
    return result


def read(version: int, begin: bytes, end: Optional[bytes] = None) -> ReadRange:
    begin_key, begin_buf = _make_key(begin)
    if end is None:
        end_key = _Key((ctypes.c_ubyte * 0)(), 0)
        end_buf = None
    else:
        end_key, end_buf = _make_key(end)
    result = ReadRange(begin_key, end_key, version)
    result._begin_buf = begin_buf
    result._end_buf = end_buf
    return result


class ConflictSet:
    def __init__(
        self,
        version: int = 0,
        build_dir: Optional[str] = None,
        implementation: Optional[str] = None,
    ) -> None:
        self._lib = None
        if build_dir is None:
            build_dir = os.path.dirname(__file__) + "/build"
        if implementation is None:
            implementation = "radix_tree"
        for f in (
            build_dir + "/" + implementation + "/libconflict-set.so.0",
            os.path.dirname(__file__)
            + "/build/"
            + implementation
            + "/libconflict-set.0.dylib",
        ):
            try:
                self._lib = ctypes.cdll.LoadLibrary(f)
            except:
                pass

        if self._lib is None:
            import sys

            print(
                "Could not find libconflict-set implementation " + implementation,
                file=sys.stderr,
            )
            sys.exit(1)

        self._lib.ConflictSet_create.argtypes = (ctypes.c_int64,)
        self._lib.ConflictSet_create.restype = ctypes.c_void_p

        self._lib.ConflictSet_check.argtypes = (
            ctypes.c_void_p,
            ctypes.POINTER(ReadRange),
            ctypes.POINTER(ctypes.c_int),
            ctypes.c_int,
        )
        self._lib.ConflictSet_check.restype = None

        self._lib.ConflictSet_addWrites.argtypes = (
            ctypes.c_void_p,
            ctypes.POINTER(WriteRange),
            ctypes.c_int,
            ctypes.c_int64,
        )
        self._lib.ConflictSet_addWrites.restype = None

        self._lib.ConflictSet_setOldestVersion.argtypes = (
            ctypes.c_void_p,
            ctypes.c_int64,
        )
        self._lib.ConflictSet_setOldestVersion.restype = None

        self._lib.ConflictSet_destroy.argtypes = (ctypes.c_void_p,)
        self._lib.ConflictSet_destroy.restype = None

        self._lib.ConflictSet_getBytes.argtypes = (ctypes.c_void_p,)
        self._lib.ConflictSet_getBytes.restype = ctypes.c_int64

        self.p = self._lib.ConflictSet_create(version)

    def addWrites(self, version: int, *writes: WriteRange):
        self._lib.ConflictSet_addWrites(
            self.p, (WriteRange * len(writes))(*writes), len(writes), version
        )

    def check(self, *reads: ReadRange) -> list[Result]:
        r = (ctypes.c_int * len(reads))()
        self._lib.ConflictSet_check(
            self.p, (ReadRange * len(reads))(*reads), r, len(reads)
        )
        return [Result(x) for x in r]

    def setOldestVersion(self, version: int) -> None:
        self._lib.ConflictSet_setOldestVersion(self.p, version)

    def getBytes(self) -> int:
        return self._lib.ConflictSet_getBytes(self.p)

    def __enter__(self):
        return self

    def close(self) -> None:
        if self.p is not None:
            self._lib.ConflictSet_destroy(self.p)
            self.p = None

    def __exit__(self, exception_type, exception_value, exception_traceback):
        self.close()