Viper_VRTE/_dict.vp

import t, c
import memhub
import string
import json


# Variant type tags (consistent with _variant.py)
define VARIANT_NONE   = 0   # Untyped / null value
define VARIANT_INT    = 1   # int (i64)
define VARIANT_FLOAT  = 2   # float (f32)
define VARIANT_DOUBLE = 3   # double (f64)
define VARIANT_STR    = 4   # str (i8*)
define VARIANT_DICT   = 5   # dict*
define VARIANT_LIST   = 6   # list*
define VARIANT_BOOL   = 7   # bool (i1/i32)
define VARIANT_PTR    = 8   # Custom pointer type

# Variant structure size: vtype(i32, 4 bytes) + padding(4 bytes) + ptr(i8*, 8 bytes) = 16 bytes
define _VARIANT_SIZE = 16


class dict:
    """Heap-allocated dynamic dictionary container based on mbuddy allocator.
    Key type is fixed to str (i8*), values are stored as Variant (type tag + untyped pointer).
    Supports nested dict/list structures and JSON serialization.

    Usage:
        mb: memhub.MemBuddy = memhub.MemBuddy(arena, arena_size)
        d: dict = dict(mb)
        d.set_int("x", 42)
        p: t.CPtr = d["x"]
        v: int = c.Deref((int*)(p))    # int value stored as i32; read it back with an int* view

        # Nested dict
        inner: dict = dict(mb)
        inner.set_int("a", 1)
        d.set_dict("sub", inner)
        sub: dict = c.Deref(d["sub"])
        av: int = c.Deref((int*)(sub["a"]))
    """
    __keys__:       rawptr
    __values__:     rawptr
    __count__:      sizet = 0
    __capacity__:   sizet = 8
    __mbuddy__:     memhub.MemBuddy*
    __iter_index__: sizet = 0

    def __new__(self, mb: memhub.MemBuddy*):
        # Heap allocation for dict object
        return mb.alloc(dict.__sizeof__())

    def __init__(self, mb: memhub.MemBuddy*):
        self.__mbuddy__ = mb
        self.__keys__   = mb.alloc(self.__capacity__ * 8)
        self.__values__ = mb.alloc(self.__capacity__ * _VARIANT_SIZE)

    def __len__(self) -> sizet:
        return self.__count__

    def _find(self, key: str) -> sizet:
        """Find index of target key; returns __count__ if not found"""
        for i in range(self.__count__):
            # Element address: `T*(base + byte_offset)`. rawptr is i8* and `+` scales by 1 byte, byte-identical to the old uint64(base)+n form (no integer round-trip).
            # Same form below (__getitem__/_set_variant_raw/get/get_vtype/__next__/to_json), not annotated again.
            key_ptr: str* = str*(self.__keys__ + i * 8)
            if string.strcmp(key_ptr[0], key) == 0: return i
        return self.__count__

    def __getitem__(self, key: str) -> rawptr:
        """Return untyped pointer to value. Use c.Deref to resolve the actual value.
        Example: p: rawptr = d["x"]; v: int = *(int*(p))"""
        idx: sizet = self._find(key)
        if idx >= self.__count__: return None
        # Variant.ptr is located at offset 8 within Variant structure
        val_ptr: str* = str*(self.__values__ + idx * _VARIANT_SIZE + 8)
        return *val_ptr

    def __setitem__(self, key: str, val: t.CNeedPtr):
        """Store value as VARIANT_PTR. Address is automatically resolved via t.CNeedPtr.
        Note: Lifecycle of the data pointed to by val is managed by caller.
        Use typed methods such as set_int / set_str for persistent storage."""
        self._set_variant_raw(key, VARIANT_PTR, val)

    def _set_variant_raw(self, key: str, vtype: int, ptr: t.CPtr):
        """Internal method: store Variant entry (vtype + ptr). No automatic address resolution."""
        idx: sizet = self._find(key)
        if idx < self.__count__:
            # Update existing key
            vtype_ptr: int* = int*(self.__values__ + idx * _VARIANT_SIZE)
            vtype_ptr[0] = vtype
            val_ptr: str* = str*(self.__values__ + idx * _VARIANT_SIZE + 8)
            val_ptr[0] = ptr
            return
        if self.__count__ >= self.__capacity__:
            new_cap: sizet = self.__capacity__ * 2
            new_keys: rawptr = self.__mbuddy__.alloc(new_cap * 8)
            new_values: rawptr = self.__mbuddy__.alloc(new_cap * _VARIANT_SIZE)
            if new_keys is None or new_values is None: return
            string.memcpy(new_keys,   self.__keys__,   self.__count__ * 8)
            string.memcpy(new_values, self.__values__, self.__count__ * _VARIANT_SIZE)
            self.__keys__     = new_keys
            self.__values__   = new_values
            self.__capacity__ = new_cap
        key_ptr: str* = str*(self.__keys__ + self.__count__ * 8)
        key_ptr[0] = key
        vtype_ptr: int* = int*(self.__values__ + self.__count__ * _VARIANT_SIZE)
        vtype_ptr[0] = vtype
        val_ptr: str* = str*(self.__values__ + self.__count__ * _VARIANT_SIZE + 8)
        val_ptr[0] = ptr
        self.__count__ += 1

    # ============================================================
    # Typed setters — allocate persistent storage inside __mbuddy__
    # ============================================================

    def set_int(self, key: str, val: int):
        """Store integer value (VARIANT_INT). Value is copied to persistent storage within __mbuddy__."""
        storage: rawptr = self.__mbuddy__.alloc(8)
        int_ptr: int* = int*(storage)
        *int_ptr = val
        self._set_variant_raw(key, VARIANT_INT, storage)

    def set_str(self, key: str, val: str):
        """Store string value (VARIANT_STR). Raw string pointer is stored directly."""
        self._set_variant_raw(key, VARIANT_STR, val)

    def set_float(self, key: str, val: float):
        """Store float value (VARIANT_FLOAT). Value is copied to persistent storage within __mbuddy__."""
        storage: rawptr = self.__mbuddy__.alloc(8)
        float_ptr: float* = float*(storage)
        *float_ptr = val
        self._set_variant_raw(key, VARIANT_FLOAT, storage)

    def set_double(self, key: str, val: double):
        """Store double value (VARIANT_DOUBLE). Value is copied to persistent storage within __mbuddy__."""
        storage: rawptr = self.__mbuddy__.alloc(8)
        dbl_ptr: double* = double*(storage)
        *dbl_ptr = val
        self._set_variant_raw(key, VARIANT_DOUBLE, storage)

    def set_dict(self, key: str, val: dict*):
        """Store nested dictionary pointer (VARIANT_DICT)."""
        self._set_variant_raw(key, VARIANT_DICT, val)

    def set_list(self, key: str, val: t.CPtr):
        """Store nested list pointer (VARIANT_LIST)."""
        self._set_variant_raw(key, VARIANT_LIST, val)

    def set_bool(self, key: str, val: int):
        """Store boolean value (VARIANT_BOOL). Value is copied to persistent storage within __mbuddy__."""
        storage: rawptr = self.__mbuddy__.alloc(8)
        bool_ptr: int* = int*(storage)
        *bool_ptr = val
        self._set_variant_raw(key, VARIANT_BOOL, storage)

    # ============================================================
    # Query methods
    # ============================================================

    def get_vtype(self, key: str) -> int:
        """Return Variant type tag associated with specified key."""
        idx: sizet = self._find(key)
        if idx >= self.__count__:
            return VARIANT_NONE
        vtype_ptr: int* = int*(self.__values__ + idx * _VARIANT_SIZE)
        return vtype_ptr[0]

    def get(self, key: str, default: t.CPtr) -> t.CPtr:
        """Return value pointer; return default pointer if key does not exist."""
        idx: sizet = self._find(key)
        if idx >= self.__count__:
            return default
        val_ptr: str* = str*(self.__values__ + idx * _VARIANT_SIZE + 8)
        return *val_ptr

    def contains(self, key: str) -> int:
        """Check whether target key exists."""
        idx: sizet = self._find(key)
        if idx < self.__count__: return 1
        return 0

    def __iter__(self) -> dict*:
        self.__iter_index__ = 0
        return self

    def __next__(self) -> str:
        if self.__iter_index__ >= self.__count__:
            raise StopIteration
        idx: sizet = self.__iter_index__
        self.__iter_index__ = idx + 1
        key_ptr: str* = str*(self.__keys__ + idx * 8)
        return key_ptr[0]

    # ============================================================
    # JSON utilities: serialization / deserialization via includes/json library
    # ============================================================

    def to_json(self, pool: memhub.MemManager*) -> json.JsonValue*:
        """Convert dict instance to JsonValue object (JSON_OBJECT).
        Automatically select corresponding JSON value type according to Variant tag."""
        obj: json.JsonValue* = json.object(pool)
        for i in range(self.__count__):
            key_ptr: str* = str*(self.__keys__ + i * 8)
            vtype_ptr: int* = int*(self.__values__ + i * _VARIANT_SIZE)
            val_ptr: str* = str*(self.__values__ + i * _VARIANT_SIZE + 8)
            vt: int = vtype_ptr[0]
            vp: rawptr = val_ptr[0]
            jv: json.JsonValue* = json.null(pool)
            if jv != None:
                if vt == VARIANT_INT:
                    jv.vtype = 2  # JSON_INT
                    ip: int* = int*(vp)
                    jv.int_val = ip[0]
                elif vt == VARIANT_STR:
                    jv.vtype = 4  # JSON_STRING
                    jv.str_val = vp
                elif vt == VARIANT_BOOL:
                    jv.vtype = 1  # JSON_BOOL
                    bp: int* = int*(vp)
                    jv.bool_val = bp[0]
                elif vt == VARIANT_FLOAT:
                    jv.vtype = 3  # JSON_FLOAT
                    fp: float* = float*(vp)
                    jv.float_val = fp[0]
                elif vt == VARIANT_PTR:
                    # Fallback: treat raw pointer as integer
                    jv.vtype = 2  # JSON_INT
                    ip: int* = int*(vp)
                    jv.int_val = ip[0]
            json.object_set(pool, obj, key_ptr[0], jv)
        return obj

    def from_json(self, pool: memhub.MemManager*, root: json.JsonValue*):
        """Populate dict from JsonValue object.
        Extract values and store them as matching Variant types based on JSON type tag."""
        if root == None:
            return
        if root.vtype != 6:  # JSON_OBJECT
            return
        cur: json.JsonValue* = root.child
        while cur != None:
            if cur.key != None:
                if cur.vtype == 2:  # JSON_INT
                    storage: rawptr = self.__mbuddy__.alloc(8)
                    ip: int* = int*(storage)
                    ip[0] = cur.int_val
                    self._set_variant_raw(cur.key, VARIANT_INT, storage)
                elif cur.vtype == 4:  # JSON_STRING
                    self._set_variant_raw(cur.key, VARIANT_STR, cur.str_val)
                elif cur.vtype == 1:  # JSON_BOOL
                    storage: rawptr = self.__mbuddy__.alloc(8)
                    bp: int* = int*(storage)
                    bp[0] = cur.bool_val
                    self._set_variant_raw(cur.key, VARIANT_BOOL, storage)
                elif cur.vtype == 3:  # JSON_FLOAT
                    storage: rawptr = self.__mbuddy__.alloc(8)
                    fp: float* = float*(storage)
                    fp[0] = cur.float_val
                    self._set_variant_raw(cur.key, VARIANT_FLOAT, storage)
            cur = cur.next

    def dumps(self, pool: memhub.MemManager*) -> str:
        """Serialize dict to JSON string."""
        obj: json.JsonValue* = self.to_json(pool)
        return json.write(pool, obj, False)

    def loads(self, pool: memhub.MemManager*, json_str: str):
        """Parse JSON string and populate dict instance."""
        root: json.JsonValue* = json.parse(pool, json_str)
        self.from_json(pool, root)